NAS+CNN+Transformer=ViT-Res！MIT团队重磅开源ViT-Res，精度高于DeiT-Ti8.6%

本文分享论文『Searching for Efficient Multi-Stage Vision Transformers』，由 MIT 团队重磅开源 ViT-Res，Tiny 模型精度比 DeiT-Ti高 8.6%。

详细信息如下：

• 论文链接：https://arxiv.org/abs/2109.00642

• 项目链接：https://github.com/yilunliao/vit-search

导言：

Vision Transformer(ViT)证明，用于自然语言处理的Transformer可以应用于计算机视觉任务，并能够达到与卷积神经网络(CNN)相当的性能。这自然地引出了一个问题，如何用CNN的设计技术来改进ViT的性能。为此，作者结合两种CNN中应用的技术并提出了ViT-ResNAS，这是一种由神经结构搜索(NAS)设计的多阶段ViT结构。

首先，作者提出了残差空间缩减（residual spatial reduction） ，利用多阶段体系结构，减少深层block中的token序列长度。当减少序列长度时，作者添加了skip connection，以提高模型性能，并稳定网络的训练。

其次，作者提出了具有多架构采样的权重共享NAS 。作者扩大了网络，并利用其子网（sub-networks）来定义一个搜索空间。然后训练一个覆盖所有子网的超网（super-network），以快速评估其性能。为了有效地训练超网，作者用一次前向-后向传播来对多个子网进行采样和训练。然后，通过进化搜索（evolutionary search）来寻找高性能的网络架构。

在ImageNet上的实验表明，ViT-ResNAS比原始的DeiT和其他较强的ViT 具有更好的精度-MACs和精度-吞吐量trade-off。

      01      

Motivation

Self-attention 和 Transformers起源于自然语言处理(NLP)，已被广泛应用于计算机视觉(CV)任务，包括图像分类、目标检测和语义分割。许多工作利用混合架构，并将Self-attention机制纳入卷积神经网络(CNN)，以建模长距离依赖，提高网络的性能。另一方面，ViT证明，在JFT-300M等大型数据集上训练时，即使没有卷积的纯Transformer结构，也可以在图像分类上取得很不错的性能。此外，DeiT表明，在ImageNet上进行更强的正则化训练时，ViT的表现优于CNN。

同时，另一个研究路线是用神经结构搜索(NAS)设计有效的神经网络。最开始的NAS工作使用强化学习来设计高效的CNN架构。他们在预定义的搜索空间中对许多网络进行采样，并从头开始训练它们，以估计它们的性能，这种方式需要大量的计算资源。为了加速这一过程，权重共享NAS（weight-sharing NAS）出现了。

权重共享NAS不是在搜索空间中训练单个网络，而是训练一个超网，该网络的权重在搜索空间中的所有网络中共享。一旦超网被训练好，我们就可以直接使用它的权值来估计不同网络在搜索空间中的性能。这些方法成功地使CNN的架构优于人工设计的架构。

CNN架构已经在CV中被研究和采用多年，并使用NAS进行了优化，但最近ViT表现出了优于CNN的性能。虽然ViT的性能很不错，但是它采用了和NLP Transformer相同的结构。这自然就引出了如何通过CNN来进一步提高ViT的性能的问题 。因此，作者结合了CNN的两种设计技术，即空间缩减和神经结构搜索 ，提出了ViT-ResNAS，这是一种有效的具有残差空间缩减的多阶段ViT结构，并使用了NAS设计。

首先，作者提出了残差空间缩减（residual spatial reduction） ，以减少序列长度和增加embedding尺寸来生成更深的transformer块的。如上图所示，作者将原始的单阶段的transformer架构转换为多阶段架构 ，每个阶段具有相同的序列长度和embedding大小。

这种架构遵循CNN的架构，更深层的特征图分辨率降低，而通道尺寸增加。此外，作者在减少序列长度时增加了skip connections ，这可以进一步提高性能，稳定训练更深层次的网络。作者将具有残差空间缩减的ViT被命名为ViT-Res。

其次，作者提出了具有多结构采样的权重共享神经结构搜索，以改进ViT-Res的结构（如上图所示）。作者通过增加深度和宽度来扩大ViT-Res网络。它的子网被用来定义为搜索空间。然后，训练一个覆盖所有子网的超网，以直接评估其性能。对于每个训练迭代和给定一个batch的样本，作者通过一次前向-后向传播来对多个子网进行采样和训练，以有效地训练超网。一旦超网被训练，进化搜索就可以用来寻找高性能的ViT-ResNAS网络。

      02      

方法

2.1 Background on Vision Transformer

Vision Transformer主要组件包括tokenization、position embedding、多头自注意(MHSA)、前馈网络(FFN)和layer normalization(LN)。

Tokenization

ViT的输入是一个固定大小的图像，并被分割成多个patch。每个patch用线性层转化为维度为的patch embedding。因此，一个图像可以被表示成一系列的 patch embeddings。为了执行分类任务，将在序列中添加一个的可学习向量（[CLS] token）用于分类。使用transformer blocks处理后，分类token捕获的特征用于预测分类结果。

Position Embedding

为了保持图像patch的相对位置，将预定义或可学习向量的position embedding添加到patch embedding中。然后Transformer Block用于处理position embedding和patch embedding联合的结果。

MHSA

MHSA首先使用三个可学习的矩阵将输入特征转换成key, query和value。然后，对于每个query，计算其与N个key的内积，然后除以，并对结果应用softmax函数得到N个归一化权值。Attention函数的输出是N个value向量的加权和，表示如下：

一个MHSA包含h个平行的Attention head。h个不同的输出concat到一起并通过线性层进行变换投影，得到MHSA的最终输出。

FFN

FFN有两个线性层和一个激活函数GeLU组成。第一层线性层将每个embedding的维度扩展到隐藏层大小，第二层线性层将维度投影回原来的大小。

LN

Layer normalization用来对每个embedding进行归一化：

µ和σ分别是输入特征的均值和方差。

2.2 Residual Spatial Reduction

ViT在整个网络中保持相同的序列长度。相比之下，CNN降低了特征图的分辨率，增加了更深层的通道大小。前者是基于特征图中存在空间冗余的事实，后者是为更深层次的高级特征分配更多的通道，并在分辨率降低时保持相似的计算量。在本文中，作者也希望能够借鉴CNN的这种设计。为此，作者提出了残差空间缩减。

如上图所示，由于patch embedding存在空间关系，作者首先将一维序列重塑为二维特征图，然后用LayerNorm和卷积来减小patch embedding的序列长度，增加通道维度。由于序列长度发生了改变，需要通过添加新的位置embedding来更新相对位置信息。为了保持所有embedding的相同通道维度，作者应用了LayerNorm和一个线性层(图中的“Norm”和“Linear”)得到新的分类token。这些结构构成了残差分支（residual branch）。

虽然只引入残差分支可以显著提高精度-MAC的权衡，但训练更深层次的网络可能是不稳定的。为了解决这一问题，作者引入了一个额外的skip connection。作者使用2D的平均池化来减少特征的长度，然后和零向量进行concat增大通道维度，这样的结构有助于深层Transformer的稳定训练。这些结构构成了主要分支（main branch）。

残差分支和主要分支构成了残差空间缩减（Residual Spatial Reduction）。作者将2个残差空间缩减块插入到包含12个Transformer块的ViT中，并将整个ViT结构均匀地分为3个stage。类似ResNet，空间维度减半时，特征的通道维度翻倍。得到的结构为ViT-Res，如上图所示。

2.3 Weight-Sharing NAS with Multi-Architectural Sampling

CNN中使用的另一种设计技术是神经结构搜索(NAS)。考虑到搜索效率，作者采用权重共享NAS来改进ViT-Res的结构，设计更好的架构参数，比如Attention head和Transformer Block的数量，如上图所示。

Algorithm Overview

作者通过增加深度和宽度来扩大ViT-Res网络。深度和宽度较小的子网定义为搜索空间。然后，训练一个覆盖所有子网的超网来直接评估其性能。对于每一次训练迭代和给定一个batch的样本，作者在一次前向-后向传播中对多个不同的子网进行采样，以有效地训练超网。然后，使用进化搜索来寻找满足一定资源约束的高性能子网架构。最后，搜索得到的最佳子网为搜索网络，并从头开始训练。

Search Space

作者通过均匀增加ViT-Res各阶段的深度和宽度来构建一个大网络，大网络中包含的子网定义了搜索空间。对于每个阶段，需要搜索embedding的维度和Transformer Block的数量。对于不同的Block，在MHSA中搜索不同数量的Attention head 和在FFN中隐藏层通道维度。每个可搜索维度的范围都是预定义的。在超网训练和进化搜索过程中，对不同配置的子网络进行采样以进行训练和评估。

Multi-Architectural Sampling for Super-Network Training

为了评估子网的性能，必须优化它们的权重，以反映其相对质量。因此，作者训练一个超网，其体系结构与定义搜索空间的最大网络相同，其权值在子网之间共享。在超网训练过程中，对于不同的训练迭代过程，作者采样并训练不同的子网。一般来说，子网的采样量越多，训练后的超网所表示的子网的相对质量就越准确。先前对CNN的NAS的研究，每次前向-后向传播中都采样并训练一个子网。在本文中，为了提高训练的效率，作者在每次训练迭代中采样并训练了多个子网，如上图所示（左：对一个子网进行采样；右：多架构抽样（）不同的子网）。

与CNN中的batch normalization不同，ViT中的layer normalization避免了沿batch维度的归一化，这使得在一个前向-后向传播中能够对不同的子网进行采样，而不混合它们的统计数据。具体来说，对于每个训练迭代，对个子网络进行采样，并将一个Batch分成个子Batch。每个子网都用其相应的子Batch进行训练。这可以通过一次前向-后向传播和channel masking来实现。如上图所示，对于不同样本的feature map，生成不同的mask，以模拟不同的子网采样。特征映射的形状保持不变，以方便完整的Batch计算。

此外，作者在采样多个子网时重新校准了LN中的统计数据，以防止超网训练和标准训练之间的差异。如上图（b）和（c）所示，当采样具有较少通道数的网络时，LN在较大的通道维度上就会进行错误的归一化。这是因为通道mask只改变特征映射的值，而不是改变形状。

为了避免这一问题，作者提出了masked layer normalization (MLN)，当LN的输入被mask时，它会重新校准LN的统计量。MLN不只查看输入张量的形状，而是计算mask通道数与通道总数的比值，并使用该比率来重新校准统计数据（如上图d所示）。最终，通过通道mask和MLN，就可以用一次前向-后向传播对多个子网进行采样，这提高了训练ViT-Res超网时的采样效率，从而提高了搜索网络的性能。

Evolutionary Search

在整个搜索过程中，只对满足预定义资源约束(如MACs)的子网进行采样，并使用训练过的超网权值来评估它们的性能。进化搜索维护了搜索的网络群，并通过多次迭代选出performance最好的网络。

首先从随机初始化的网络开始迭代。在每次迭代中，个性能最高的子网作为父网络，通过mutation和crossover生成个新子网。对于mutation，从父网络中随机选择一个子网络，并使用预定义的概率突变来修改每个体系结构参数。

对于crossover，从父网络中选择两个随机的子网络，并随机融合每个体系结构参数。mutation和crossover产生相同数量的新子网络，并将它们添加到网络群中。这个过程重复次，最好的子网络成为最终的搜索网络，然后从头训练到收敛。

2.4 Extra Techniques

除了方面的方法，作者还采用了其他技术来进一步提升性能。

Token Labeling with CutMix and Mixup

作者采用了token labeling和Mixup。token labeling为输入图像中的所有patch提供了label，这使得除了预测具有[CLS] token的类标签，训练ViT还能够预测所有patch tokens的类标签，提升Transformer的精度。token labeling通过patch的CutMix生成输入图像、类标签和patch-wise类标签。

给定图片、和对应的标签、，首先将每个图像分成K个patch，将所有patch进行flatten，并将每个patch与其原始的图像级类标签相关联。这样，我们就得到了图像patch序列，，和，，，以及patch标签序列，，和，，。

然后，我们随机生成一个Mask M来结合图像和相关的patch标签：。然后组成图片的标签为λ（λ），其中λ是二进制Mask M中所有元素的平均值。

作者发现token labeling和Mixup可以提高性能，这与之前工作得到的结论不同。作者猜想这是因为之前的工作先用了Mixup然后再用CutMix，引入了噪声。

Convolution before Tokenization

基于之前的一些工作，作者在本文中也在tokenization之前，采用了一些卷积操作。具体来说，作者添加了三个额外的3×3卷积层，每个层都有C输出通道。第一层的步长设置为2，其他的则为1。在第一层和第三层的输出之间添加一个residual skip connection。

      03      

实验

3.1 Ablation Study

Multi-Stage Network with Residual Connection and Token Labeling

上表展示了本文不同结构的消融实验，证明了本文提出的每一个模块都有非常大的作用。

Weight-Sharing NAS with Multi-Architectural Sampling

上表比较了每次采样子网数量对最终结果的影响，可以看出，当采样数为16时，模型的性能是最好的。

3.2 Comparison with Related Works

上表总结了本文的方法与以往的ViT工作的比较。可以看出，在相似的计算量和参数量下，本文的方法还是具有性能上的优势的。

      04      

总结

在本文中，作者将CNN的两项设计技术（空间缩减和NAS）合并到ViT中，提出了ViT-ResNAS。所提出的残差空间缩减显著提高了精度-mac的权衡，残差连接可以提高性能和稳定训练。多架构采样的权重共享NAS可以更有效地训练超网，从而得到更好的搜索网络。在ImageNet上的实验结果，证明了本文方法的有效性和ViT-ResNAS网络的效率。

作者介绍

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。

知乎/公众号：FightingCV

END,入群👇备注：NAS