搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（十五）

专栏目录 https://zhuanlan.zhihu.com/p/348593638

本文目录

33 HR-NAS：使用轻量级Transformer的高效搜索高分辨率神经架构
(来自 香港大学、字节跳动和中国人民大学)
33.1 HR-NAS原理分析

34 AutoFormer：搜索用于视觉识别的Transformer
(来自 微软)
34.1 AutoFormer原理分析

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

本文介绍的2篇文章都是把视觉 Transformer 和 NAS 结合的工作。

33 HR-NAS：使用轻量级Transformer的高效搜索高分辨率神经架构

论文名称：HR-NAS: Searching Efficient High-Resolution Neural Architectures
with Lightweight Transformers

论文地址：

https://arxiv.org/pdf/2106.06560.pdf

33.1 HR-NAS原理分析：

本文是香港大学、字节跳动的研究员们提出的一种针对密集预测任务，搜索 Transformer 架构的 NAS 方法，入选了 CVPR2021 (Oral)。

密集预测任务粗鲁的 NAS 方法

其实以往针对密集预测任务 (语义分割，姿态估计) 的 NAS 方法的搜索空间的设计比较粗鲁，就是直接拓展分类任务的 NAS 方法的搜索空间。所以搜索空间主要是分类任务的搜索空间加上特征融合的head的搜索。而网络的 Backbone 则看成是由一个个 cell 构成的单分支架构。这种 NAS 的方法看上去就很粗鲁，因为分类任务的搜索空间并不一定对密集预测任务是最好的，并且将 Backbone 看成是一个个 cell 构成的单分支架构大大限制了网络设计的灵活性。这就导致了这种粗鲁的 NAS 方法得到的结果是次优的，其性能不如手工设计的网络。

密集预测任务的特点

一般而言密集预测任务很看重两点：其一是全局信息的完整性；其二是高分辨率的特征表示，它们对于细节信息的精确预测很关键，比如语义信息的边界以及关键点的位置。但是分类任务对于这两点考虑欠佳，尤其是第2点高分辨率的特征表示，最直接的解决办法是在网络的末端手动去结合不同尺度的特征信息。HRNet是一种在网络的 Backbone 部分就处理多尺度信息 (multi-scale feature) 的方法。但是 multi-scale feature 却无法保证第1点全局信息的完整性。因为密集的预测任务通常具有高输入分辨率，但是网络通常只覆盖固定的感受野。因此，一些网络借助 Attention 机制，像 SENet ，non-local network等等。

现在由于 Transformer 的兴盛，作者借助 Transformer 来保证第1点全局信息的完整性。对于第2点，作者借鉴 HRNet 的思路，在网络的 Backbone 部分就处理多尺度信息 (multi-scale feature)。基于以上两点作者提出了 HR-NAS 方法，它的特点是：

• 搜索空间是多分支的。

• 搜索空间中同时包含 Convolutions 和 Transformers 的结构。

• 搜索策略考虑了资源的预算和限制。

多分支的搜索空间

作者为密集预测任务设计了一个多分支的搜索空间，它会考虑多种不同尺度特征的信息，既能保证全局信息的完整性，又可以在整个网络保持特征的高分辨率。

HR-NAS 多分支的搜索空间如下图1所示。我们看到超网络主要包含2中模块：并行模块(Parallel module) 和融合模块 (Fusion module)，且它们都是可以搜索的。网络一开始先过2层卷积，将输入特征的维度降低为  ，再从这个高分辨率分支开始，通过融合模块逐渐增加高到低分辨率的分支，融合模块的作用是在不同的分支之间交换信息。再通过并行模块连接这个多分支架构。到了网络的末端，所有的分支被 Resize 并 Concat 到一起，不需要任何额外的分类头/回归头。具体的搜索空间参数如下：

• 分支数：  。
• 各个分支的卷积数：  。
• 各个卷积的通道数：  。

融合模块的作用是在不同的分支之间交换信息，除此之外每个融合模块都会另外产生出一个低分辨率分支，直到分辨率达到1/32。每个融合模块的输入分支数是  ，输出分支数是  。融合模块输出时每个分支都会融合相邻两个分支的信息，例如，1/8 输出分支集成了1/4、1/8 和 1/16 输入分支的信息。

一个可搜索的 Block 的架构

一个可搜索的 Block 的架构如下图2所示。包含2条路，一条是混合卷积支路，另一条是 lightweight Transformer 支路 (目的是提供更多的全局信息)。卷积的通道数和 Transformer 的 token 数都是可以搜索的参数。我们首先看混合卷积支路：输入特征  通道数是  ，经过卷积操作  之后通道数变为  ，而后变成3个分支，每个分支拥有  个通道。这3个分支将分别通过 3×3，5×5，7×7 的卷积  ，输出会再 concat 在一起，经过卷积  得到卷积支路的输出。对于 lightweight Transformer 支路，通过具有  个 token 的 lightweight Transformer  ，得到输出，并与卷积支路的输出加在一起作为这个 Block 最终的输出。整个 Block 的这个流程可以表达为：

轻量级 Transformer 架构 lightweight Transformer

这一小节主要介绍搜索空间中的 lightweight Transformer 是什么。首先要明白作者为什么不复用 ViT 的 Transformer 架构而是重新设计了 lightweight Transformer，因为模型的特征的分辨率很高，那么势必会导致 patch 的数量非常多。如果按照 ViT 的 Transformer 架构设计，计算量随着 patch 数的平方增长，导致计算量过大，对诸多资源受限的设备都非常不友好。所以作者重新设计了 lightweight Transformer，如下图3所示，它由 Projector，Encoder 和 Decoder 组成。

Projector 部分首先把输入特征  和位置编码  concat在一起，再通过投影函数  。  具体就是 Point-wise convolution，将通道数从  降低到  。再通过双线性插值操作 (bilinear interpolation) 将空间维度从  降低到  。位置编码就是简单的归一化位置信息操作：

最终得到的特征的维度是  。

经过了投影层之后，原来的特征  变成了特征  ，它包含  个 tokens，且每个 token 都是长度为  的向量。特征  再输入到 Transformer 的 Encoder 里面。这里面通过一系列的 Multi-Head Self-Attention 建模全局信息，如下式4所示：

Feed-Forward Network (FFN) 表示为  ，使用 ReLU 作为激活函数，表达式是：

式中，  。

Decoder 的结构和 Encoder 完全一致。和 DETR 的设计思路一致，Decoder 使用 Encoder 的输出作为 key 和 value，并使用一个可学习嵌入矩阵  作为 Query。

Decoder 的输出经过反投影变换  (具体就是 Point-wise convolution 和双线性插值操作) 将特征变换回合适的维度。

以上就是 lightweight Transformer 的流程，它的时间复杂度有多大呢？

MHSA 和 FFN 的计算复杂度分别是  和  。因为  是投影变换之后的空间维度，所以总的计算复杂度和  成近似的线性关系。

资源感知的细粒度搜索策略

这个工作的搜索实质上是对 channel 和 token 的剪枝。由于每个 Block 的 channel 数都是独立的，所以各个 Block 的 channel 数不互相影响。似地，如果Transformer的1个 token 被丢弃，我们只需要从投影  和反投影  的 1×1 卷积中移除相应的权重，并从 query 中移除相应的 embedding 即可。把一个 channel 或者一个 token 定义为一个 search unit，则一个 Block 里面就有  个 search unit。

遵循 Darts 的方法，作者引入了一个重要性因子  ，它可以与搜索块的每个搜索单元的网络权重一起学习。然后逐步丢弃那些重要性较低的 channel 或 token，同时保持整体性能。受已有的 channel 剪枝工作的启发，作者在  上增加了一个  范数的罚项。有效地将高计算成本的 search unit 对应的  推为零。具体而言，一个 search unit 的  范数的罚项是通过下式来计算 (实质就是不同操作的计算量)：

尽管计算量 FLOPs 来衡量网络的延迟不是最优解，但是它是最常用和最简单的选择，总的损失函数是：

式中，  是对应任务的损失，后面那一项就是罚项。

在训练期间，在每隔几个 epochs 之后，作者逐步移除重要性因子  低于预定阈值 (0.001) 的 search unit，并重新校准批次归一化 (BN) 层的运行统计值。当搜索结束时，剩余的结构不仅代表了最佳的精度-效率的权衡，而且对于特定的任务具有最佳的 低级/高级(多分支) 和 局部/全局(包含卷积和 Transformer) 特征组合。此外，由于网络训练和架构搜索是以统一的端到端方式进行的，因此可以直接使用生成的网络，而无需微调。

Experiment：

HR-NAS 这个工作的实验部分特别丰富，作者在4种不同的的任务上面验证方法的有效性：

1. 图像分类：数据集：ImageNet
2. 姿态估计：数据集：COCO keypoint
3. 语义分割：数据集：Cityscapes，ADE20K
4. 3D目标检测：数据集：KITTI

所有的5个数据集使用相同的超网络。HR-NAS的超网络如上图4所示，网络一开始是2个stride=2, c=24的3×3卷积，后面接5个并行模块(Parallel module)，且每2个并行模块之间穿插一个融合模块 (Fusion module)，以得到多尺度的特征，4个不同分支的 channel 数分别是 18, 36, 72, 144。Transformer的超参数  。

实验结果：

图像分类：数据集：ImageNet

语义分割：数据集：Cityscapes

语义分割：数据集：ADE20k

姿态估计：数据集：COCO keypoint

3D目标检测：数据集：KITTI

HR-NAS 方法在所有三个密集预测任务上获得了最先进的性能，在分类任务上获得了竞争结果。与其他任务相比，分类通常较少受益于多尺度和全局上下文特征，因为它通过全局池化聚集位置不变的特征。

搜索结果可视化

如下图10-13是4种任务的搜索出来的架构的可视化结果。其中蓝色，红色，黄色，绿色和灰色部分分别代表3×3, 5×5, 7×7卷积的channel数，Transformer 的 query 数，以及剪掉的channel/queries。我们可以观察到，HR-NAS 可以为不同的任务找到不同的架构，表明它可以自动适应各种任务：

1. 在图像分类任务和3D检测任务中，在高分辨率分支 (即第1和第2分支) 处，HR-NAS找到的模型去除了大部分搜索单元；一些搜索块甚至被完全移除，如图10中完全灰色的圆圈所示。原因是在这两个任务中，全局语义信息比局部信息更重要。
2. 用于分割任务的模型仍然保留了来自前两个分支的计算，因为保持用于语义分割任务的高分辨率图像是重要的。
3. 人体姿态估计模型主要利用第二和第三分支，这意味着它可以更多地依赖中分辨率的特征，而不是高分辨率语义。
4. Transformer 更多地用于分割和人体关键点估计任务，表明这些密集的预测任务更多地受益于全局上下文。

小结

HR-NAS 遵循 HRNet 的方法设计了一种包含卷积和 Transformer 的多分支架构，使之可以兼顾全局信息的完整性和高分辨率的特征表示。并且使用 NAS 方法对每个分支的每个 Block 的 channel 数或者 token 数使用 Darts 的方法进行搜索，这个搜索本质上是一个剪枝的过程。HR-NAS 在多种任务中达到了 SOTA 的性能，缺点是搜索空间只包含 channel (对应 Transformer 的 token) 维度的搜索，对于网络的连接方式依然采用手工设计。

34 AutoFormer：搜索用于视觉识别的Transformer

论文名称：AutoFormer: Searching Transformers for Visual Recognition

论文地址：

https://arxiv.org/pdf/2107.00651.pdf

34.1 AutoFormer原理分析：

四个维度的搜索空间

本文是标准的使用 One-Shot NAS 策略搜索视觉 Transformer 模型的工作，搜索空间和韩松老师组在 ACL 2020年的工作 HAT (作者也引用了这个文章) 一模一样，都是从以下 4 个维度搜索：

• Transformer模型的嵌入维度：Embedding dimension
• Transformer模型的层数：Depth
• Transformer模型 MLP 的 Expansion ratio
• Transformer模型 Head 的数量

并且作者探索了这4个维度的不断增加对模型性能的影响，如下图14-15所示。

结果发现增加变压器的深度、head 数和 MLP 的 Expansion ratio 有助于最初获得更高的精度，但在达到峰值后会过拟合。扩大嵌入维度可以提高模型能力，但对于更大的模型，精度会停滞不前。这说明有必要对这几个参数的最优值来个精确的搜索。

One-Shot NAS 方法的基本思路是先预训练一个超网络 SuperNet，这个超网络包含了若干个不同结构的子网络，其中各个子网络的权重也得到了良好的训练，其性能与 train from scratch 的子网络差不多。在这个工作里面作者使用了权重纠缠 (weight entanglement) 的方式来训练这个超网络，它可以让超网中的大量子网络得到很好的训练，这样，这些从超网络继承权重的子网络的性能可以与从头开始重新训练的子网络相媲美。在超网络训练好以后，再按照 One-Shot NAS 方法的步骤使用 Evolutionary Search 搜索最适宜的子网络架构。

搜索空间如下图16所示，一共有4个维度，每个维度的具体配置如图17所示。比如说对于 Tiny 模型的架构搜索，Embedding dimension的选择是 (192, 240)，以24为间隔。模型的层数的选择是 (12, 14)，以1为间隔，等等。最终得到的模型的参数量的范围在4-9M之间。

遵循 One-Shot NAS 的方法，我们将搜索空间编码到超网络中。也就是说，空间中的每个模型都是超网络的一部分/子集。所有子网共享其公共部分的权重。超网是空间中最大的模型，其架构如图16所示。特别地，超网堆叠具有空间中定义的最大的 Embedding dimension、QKV 维数 和 MLP 的 Expansion ratio。在训练过程中，更新了所有可能的子网络相应的权重。

超网络训练：权重纠缠 weight entanglement

One-Shot NAS 方法一般在训练 Supernet 时使用权重共享 (weight sharing)，它的特点是：不同的子网络共享参数，但是在同一层中的不同的 operator 的权重是解耦的。这种做法在 Transformer 模型里面是怎么用的呢？

假如说某一层的 Embedding dimension 有3个候选值，分别是  ，给这一层分别构造 3个 Block，它们的输出之和代表这一层的输出，就相当于是超网络的这一层具有3个分支，各个分支的 Embedding dimension 分别是  。按照这样的方式来训练超网络。

权重共享 weight sharing 对卷积模型很有效，但是对 Transformer 模型而言会遇到一些问题：其一就是超网络收敛缓慢，如下图19 (左) 所示。这是因为不同的 Block 得到训练的次数太少了。其二就是从单次超网继承权重的子网的性能，在经典的权重共享策略下训练，远远低于它们 train from scratch 的真实性能，如下图19 (右) 所示。在搜索之后，仍然有必要对搜索到的架构执行额外的再训练，因为权重没有被完全优化。

那权重纠缠这种做法在 Transformer 模型里面是怎么用的呢？

假如说某一层的 Embedding dimension 有3个候选值，分别是  ，给这一层分别构造1个 Block，它的输出代表这一层的输出，这3个候选 Block 相当于分别取它的前  维。

比如对于具有  层的超网络  来讲，我们将其结构和权重表达为：

式中，  是第  层选择的块，  是该块的权值。在架构搜索过程中，每层都有多个块选择。因此，  和  实际上是从属于搜索空间的一组  个候选块中选择的，写成公式为：

式中，  是第  层第  个候选块，  是该块的权值。权重纠缠方法强制同一层中的不同候选块共享尽可能多的权重。这要求，对于同一层中的任意两个块  和  ，我们有：

这种层内权重共享使得  和  的权重更新相互纠缠。任何块的训练都将影响其他块的相交部分的权重，如上图16所示。这与One-Shot NAS中的经典权重共享策略不同，后者对于同一层中的任意两个块  和  ，我们有：

权重纠缠方法的好处：其一就是超网络收敛很快，如上图17 (左) 所示。这是因为它允许每个块比以前的独立训练策略更新更多次。其二是内存成本低，我们现在只需要为每一层存储最大的 Embedding dimension，而不是分别存储3个分支。其三是子网络性能更好。

注意，权重纠缠策略专用于同质块 (homogeneous building blocks)，例如具有不同头部数量的 Self-attention 模块，以及具有不同隐藏维度的多层感知器。根本原因是同质块在结构上是兼容的，因此权重可以相互共享。

搜索策略：资源约束下的进化算法搜索

搜索策略也和 HAT 保持一致，是基于 Evolution Search 的方法。在获得训练好的超网络之后，对其进行 Evolution Search 以获得最子网络。子网根据进化算法的管理器进行评估和选择。目标是在最小化模型尺寸的同时最大化分类精度。在进化搜索的开始，我们挑选  个随机架构作为种子。通过交叉和变异选择  个最好的架构作为父代来产成下一代。两个随机选择的候选人被挑选和交配，以产生新的子代，一共迭代20代。

Experiments：

超网络训练使用的实验配置和 DeiT 保持一致，进化算法架构搜索使用的实验配置和 SPOS 保持一致。

ImageNet 实验结果

如下图20所示是ImageNet 实验结果。AutoFormer的性能在三种不同尺寸的模型上都超过了手工设计的 ViT 和 DeiT 的性能。但依然无法超过 MobileNetV3 和 EfficientNet 这类卷积最强模型。

对比实验：Weight Entanglement 和 Evolution Search 的作用

如下图21所示，作者探索了Weight Entanglement 和 Evolution Search 的作用。具体来讲有4种不同的组合，我们发现使用传统 Weight Sharing 策略得到的 ImageNet 子网络继承权重的性能和独立训练的性能有较大的差距，而 Weight Entanglement 则不存在这个问题。权重纠缠训练策略可以产生明显优于经典权重共享的结果。纠缠权重产生的子网络与从头开始重新训练的子网络的性能非常接近。如果进一步微调或重新训练ImageNet上搜索到的子网络，性能提升非常小，甚至可以忽略不计。

小结

AutoFormer 遵循 HAT (用于 NLP 任务) 的方法从嵌入维度：Embedding dimension，层数：Depth，MLP 的 Expansion ratio，Head 的数量这4个维度对视觉 Transformer 架构进行了搜索，搜索空间与 HAT 一致。搜索策略依然使用 SPOS 的 Evolution Search。总的来讲本文其实相当于把 HAT 的方法成功地用在了 ViT 上面，在多个数据集上取得了超越 ViT/DeiT 的性能。所提出的 Weight Entanglement 的方法其实在 HAT 中已经使用 (HAT没有强调)。

如果觉得有用，就请分享到朋友圈吧！