[Transformer]ViT-ResNAS:Searching for Efficient Multi-Stage Vision Transformers

Abstract

from MIT



ViT证实了用于NLP的Transformer也可以应用于CV任务，并且取得了与CNN相当的性能，而CNN已经在计算机视觉领域被研究、应用了多年。因此自然想到一个问题：如何通过CNN的设计技术来提升ViT的性能。

本文提出的ViT-ResNAS有效的将两种技术结合在一起，用高效的神经架构搜索设计出的多阶段ViT结构。首先本文提出使用残差空间来减少深层次的序列长度，并搭建了多阶段体系结构；在缩减序列长度时，还是用skip connection来提升性能、稳定深层次网络的稳定性。其次，本文使用多架构采样的权重共享NAS，即训练一个超网，然后测试其搜索空间内的所有子网络的性能。

为了高效的训练超网，本文通过采样对具有前后通路的多个子网络进行训练，然后通过进化搜索发现高性能的网络体系结构。
在ImageNet上的实验表明，ViT-ResNAS比原始的DeiT以及其他的基准ViT网络具有更高的性能和吞吐率。

Section I Introduction

自然语言处理的self-attention、Transformer已经广泛用于CV领域，比如图像分类、目标检测、语义分割任务。许多工作使用了CNN与self-attention混合的结构来提升网络性能；另一方面ViT证明了纯Transformer结构也可以用于图像分类任务，DeiT则证明了在ImageNet数据集上训练并且经过正则化后的性能可以超过CNN。

另一方面，神经架构搜素（NAS）则致力于设计高效的神经网络结构，比如借助强化学习、权重共享等。NAS会在预先定义的搜索空间中搜索、训练网络结构的性能，但这十分消耗计算资源。为了加速这一过程，权重共享NAS变得流行起来。权重共享NAS不是在搜索空间中训练单个网络，而是先训练一个超网，其权重在搜索空间中的所有子网络共享。一旦超网训练好，就可以直接使用超网的权值来近似不同网络在搜索空间中的性能。这些方法成功的使得CNN架构优于手工设计的结构。

目前CNN已经成为CV领域的主流框架，并且通过NAS设计高效的CNN结构。而ViT在某些场景下则优于CNN性能，自然而然就引出这样一个问题：是否可以借助CNN的设计方法来优化ViT的性能？

因此本文将这两种方法结合在一起，提出ViT-ResNAS，通过空间压缩和神经架构搜索得到了一种高效的具有残差空间的多阶段ViT结构。


首先，本文提出残差空间缩减降低序列长度，提升深层transformer模块的嵌入尺寸。
Fig 1展示了ViT-ResNAS的结构。通过残差空间缩减来减少序列响度，并且阱网络分成几个不同的阶段，每一个阶段具有同样的序列长度和嵌入层尺寸、包含多个Transformer模块。除了第一阶段，其他阶段均以残差空间缩减模块开始。随着阶段加深，特征图的分辨率减小，通道数增加，此外还是用skip connection来提升性能、保证训练的稳定性。这种包含残差空间缩减的ViT模型被称为ViT-Res.


其次，本文提出了是用多结构采样的权重共享神经架构搜索策略来改进ViT-Res结构，具体如Fig 2所示。首先训练一个超网，搜索空间内定义了所有子网络结构；训练时训练的是超网，然后直接评价其搜索空间内子网的性能。对于每次训练迭代都会给一批样例，通过一次正向-反向传递来训练多个子网络。当超网训练好后，使用进化搜索来发现高性能的ViT-ResNAS结构。



ImageNet数据集上的实验结果证实了ViT–ResNAS的有效性，比如比DeiT，本文的ViT-ResNAS-Tiny网络结构提升了8.6%的精度，虽然MAC操作略有增加；而Small架构在实现了与DeiT-B相近的精度，同时MAC降低了6.3x,吞吐率提升了3.7x。此外，与ViT的其他基线模型相比，ViT-ResNAS也可以更好的权衡精度-吞吐率。






本文的贡献总结如下：
（1）提出了残差空间压缩方法来提升ViT的效率；

（2）采用多结构采样的权重共享NAS来进行架构搜索
;

（3）实验证明搜索得到的ViT-ResNAS可以更好的权衡精度-MAC

Section II Method

本文将首先回顾ViT的结构，随后讨论残差空间压缩和多结构采样的权重共享NAS。

Part 1 Vision Transformer

ViT的主要结构包括：标记化、位置嵌入、多头注意力（MHSA），前馈网络（FFN），层归一化（LN）。MHSA和FFN用于组成Transformer基本模块。

Tokenization


Vit的输入是对固定大小的图像切分成patch后的序列，并且每个patch都被转换成了d维向量（patch embedding/patch token）；这样一张图像就可以看做是patch embedding的序列；为了完成分类，主要是在序列中添加一个可训练的向量，该向量经过Transformer处理后用于预测类别。
Position Embedding

为了保留输入patch的位置信息，还需要在patch embedding追加一个向量用于嵌入位置信息，将二者结合后作为真正送入Transformer模块的输入，详情参见Fig 1左侧。

MHSA
多头注意力模块则是会处理刚才输入的patch 序列，并生成transformer后的向量：key,value,query对于每一个query会计算其与所有N个key的match关系，并且除以维度（sqrt(d)），最后的结果经过softmax获得注意力系数，将注意力系数与value相乘，获得最后的注意力输出：


多头注意力的输出结果级联后经过线性映射获得最终输出。
FFN
FFN模块包含两个线性层和一个GELU做非线性激活。
LN
层归一化会取一个样本不同通道的特征做归一化

Part 2 Residual Spatial Reduction

ViT会在整个网络中保持相同的序列长度，类似于图像中保持图像分辨率不变；而CNN会随着网络层次的加深不断降低特征图的分辨率。前者基于特征图中包含大量空间冗余信息的事实，后者则会为深层抽象特征分配更多同道，在分辨率降低的同时保持相似的计算量。
这种方法已经在高性能CNN中得到广泛应用，也激发了人们将其应用在ViT的兴趣，因此本文提出了残差空间压缩方法，详情参见Fig 3.

考虑到不同patch之间存在空间联系，因此本文将1D的patch序列reshape成2D feature map,随后进行步长卷积和LN来压缩patch序列的长度，同时提升embedding的维度。由于序列长度发生了变化，本文使用了新的位置嵌入来更新相对位置信息；此外为了保持通道尺寸相同，本文还对token使用了LN和线性映射，这些操作都进行在残差分支。

虽然仅通过引入残差分支就可以显著改善精度-MAC之间的需求，但是对深层次网络可能会导致训练不稳定。在DeiT-Distill设定下加入残差分支可以将精度从74.5%提升到79.6%，同时并未显著增加MAC数量；但如果继续训练更深层次网路偶可能导致训练损失为NaN。

为了解决这一问题，本文引入一个额外的跨层连接（skip connection）。对输入序列长度的压缩主要是通过输入2D feature map后进行均值池化实现的压缩，并将所有嵌入后的结果级联，并且通过zero padding进行升维。这一分支作为主分支，有助于提升训练的稳定性。

Fig 3中的残差分支和主分支组成了residual spatial reduction模块，本文在ViT的12个Transformer模块中插入了2个残差空间压缩模块，将整个网络拆分成3个阶段。
此外的设计遵循ResNet的设计规则，当patch embedding的空间分辨率减半或序列长度缩减为1/4时，会将embedding的维度增倍，这样得到的网络结构称之为ViT-Res.
网络其他信息：
输入图像 224 * 224

patch size = 14


embedding的维度：16


经过两次残差压缩后 分辨率降为4

Part 3 Weight-Sharing NAS with Multi- Architectural Sampling

使用神经架构搜索可以高效设计CNN结构，因此本文采用权重共享的NAS来优化ViT-Res结构，比如attention head、transformer block等的数目。Fig 4为一些搜索维度示意图。
算法概述
本文通过增加网络深度和宽度来扩充网络，则深度较浅、宽度较小的子网络们就定义了搜索空间；随后通过训练一个涵盖所有子网络的超网来直接评估各子网的性能。对于每次训练迭代会给定一组输入，本文建议通过一次前向-反向传播中采样多个不同的子网络。随后应用进化策略来发现性能更好或者满足一定资源限制的子网络结构。最后，选择评估结果最佳的子网络作为搜到的网络，并从头开始训练。
搜索空间
本文通过逐渐增加ViT-Res各阶段网络的深度和宽度来构建一个超网。



Fig 4中显示对于每个阶段搜索的是嵌入维度和transformer模块的数量；而每一个transformer block则进一步搜索MHSA中head的数目，FFN中隐藏层的维度d等参数。在每个可搜索的范围会得到不同结构的子网络，对不同的子网络进行采样、训练、评估。


多架构采样的超网训练
为了评估子网络的性能，必须优化他们的网络权重，本文采用的是训练一个超网，作为整个搜索空间的最大结构，权值在子网络之间共享。在超网训练过程中，会选择不同的子网络进行迭代训练。
一般来说，采样越多的子网络，训练后的超网所表示的精度就越准确。
根据之前CNN的NAS研究，会在每次训练迭代时采集一个子网络，经过一定次数的迭代后就可以在一次超网训练找那个采集N个子网络，详情参见Fig 5
.
与CNN中的BN不同，ViT中的层归一化（LN）会选择同一个样本的不同特征维度进行归一化处理。这样使得在一次前向-反向传播中采样多个子网络不会使得他们之间的统计结果发生混叠。对于每次训练，搜采集N个子网络并将每个patch分成N个子patch，每个子网络都用对应的子patch进行训练，可以通过通道掩膜实现（channel mask).
Fig 6展示了如何通过channel mask进行采样。

此外，在采样多个子网络时，本文还会重新校准各个子网络中做LN的统计量，这样可以防止发生超网与子网络训练之间的差异，因为如果不校准，LN要归一化的维度就有可能并不是当前需要归一化的较小的通道尺度。这是因为channel mask值改变了feature map的值，并没有改变维度信息。


为了避免这一问题的发生，本文提出了掩膜层归一化（MLN,Mask Layer Normalization），当输入的是mask后的输入，会重新校准LN的统计量。即MLN不是只取输入的通道数，而是会计算mask的通道数与通道综述的比值，并按照这一比值来重新校准统计数据，详情参见Fig 6(d)。使用MLN可以确保统计数据与单独训练子网络时的数据相同，这样可以在一次前向-反向传播过程中实现多个子网络的采样，这样就可以提升训练超网的效率，进而也提升了搜索网络的性能。







Evolutionary Search




在搜索过程中，我们只采样满足预定义要求的子网络，然后使用超网的权值来评估子网络的性能（如精度）。进化算法是这样的，会维护一个网络集群，每次迭代会选择性能最好的网络作为父代网络，然后通过变异和交叉产生N个新的子网络。
变异指的是：从父网络中随机选择一个子网络，以一定的概率P修改架构参数；交叉指的是：从父网络中随机选择两个子网络，并且随机融合两个体系结构的参数。
通过变异和交叉产生相同数量的新的子网络，并将他们加入到网络集群中；在一段搜索时间T内，选择网络及群众性能最好的网络作为搜索结果，然后从头训练直至收敛。

Part 4 Extra Technique
Token labeling with CutMix and MixUp
Token labeling

会为输入图像的每一个patch打一个标签，最后做预测就是对每一个patch做预测，然后在通过度patch label做Mixup获得输入图像的类别标签。
对于两张图像x1,x2及对应标签y1,y2，会将每张图切分成K个patch，然后将patch展平，级联获得patch序列。

随后随机生成Mask对输入和预测结果进行mask。通过将X重组为2D图像X’，预测结果Y则是通过组合原始的类标签来获得的，其中λ是mask中元素的平均值。


本文选择K = 44,这样等于最后阶段的序列长度。
此外，本文发现使用mixup和tokenlabel可以提升性能，这与前人的研究结果相悖。本文推测 peer work是先进性mixup然后在进行patch级别的MixUp可能会引入噪声；而本文则是会在Mixup和CutMix之间来回切换。



Convolution before Tokenization




本文还在进行tokenization之前添加了3个33的卷积层，并且在第一和倒数第一之间引入了skip connection,引入这部分卷积增加的计算量很少，因为空间分辨率减半，因此patch size设为为7.

Section III Experiments

实验设置
数据集：ImageNet


训练时间：8块GPU 训练300epoch


输入图像 224 * 224


batch size = 1024

Part 1 实验设置

对于NAS部分，分别设计了三种规模的网路，分别是：ViT-ResNAS-Tiny,ViT-ResNAS-Small,ViT-ResNAS-Medium,在各自超网下，Tiny,Small,Medium的drop path rate分别是0.2,0.3,0.4.
超网训练120epoch后，开始在训练过程中每次采样N个子网络，N根据经验设定为16.
对于进化算法部分，Tsearch = 20，父网络数目为75，原始网络种群数目为500，每次生成子网络150个。

Part 2消融实验

使用残差连接和token label的多阶段网络
本文研究了如何借助残差连接和token label提升原始网络的性能。
Table 1展示了在DeiT-Tiny基础网络之上添加卷积层和残差连接之后的结果，可以看到仅引入残差连接就已经提高的精度从78.5%到78.8%，额外的开销可以忽略不计。



当引入token label时精度进一步提升0.5%；结合两者 可以在不增加MAC的情况下提高1.6%的精度，达到80.1%的精度。





本文还探究了在其他实验设定下residual spatial reduction的作用，基础网络是DeiTdistill，加入残差连接确实可以提升精度；但是训练更深层次网络时就会出现训练损失NaN的情况。







Part 3 Weight-Sharing NAS with Multi-Architectural Sampling

为了探究多结构采样的权重共享NAS的效果，在训练超网时会采集N个子网进行训练，采样效率和每个子网络的训练质量由N决定。因此本文没有随意给定N值，而是探究不同N的效果。


Table 2展示了采样不同N个子网络对性能、MAC的影响，N=16时效果最佳。





Part 4 Comparison Results

本文将MAC作为限制条件，搜索到ViT-ResNAS-Tiny,ViT-ResNAS-Small,ViT-ResNAS-Medium三种网络，Table 3展示了搜索到的网络与其他网络在性能、资源商的对比。虽然本文搜索的网络其参数量更大，但注意是因为本文减少了序列长度，增加了通道数导致的，虽然参数更多，但是本文的网络在精度和吞吐量之间有更好的权衡。


与DeiT相比，ViT-ResNAS-Tiny精度提升了8.6%;Small则达到了与DeiT-Base相近的精度，但是内存、吞吐率分别提升了6.3x和3.7x；与其他多阶段网络结构相比，ViT-ResNAS始终能更好的权衡精度-吞吐量。
此外，在小型网络时ViT-ResNAS-Tiny可以更好的权衡精度-吞吐量，但是在MAC=5.0G左右时，Medium这方面的权衡稍有下降，这显示出对于高分辨率特征将注意力机制与卷积这种有效提取局部特征的操作结合在一起对于提升模型泛化性，是十分有必要的。

Section IV Related Works

Vision Transformers


ViT证实了纯Transformer结构可以应用到计算机视觉任务上，但需要在大规模数据集上预训练；DeiT则是通过正则化、增加知识蒸馏减少了对大规模数据集的依赖。其他的工作有通过结合卷积、局部注意力等将局部特征引入Transformer；本文提出的方法是对这些工作的补充，残差空间压缩可以获得更高效的多阶段架构，权重共享NAS可以提高搜索效率。

NAS

NAS也是一大研究热点，而权重共享NAS通过训练一个过参数化的超网，然后通过权重共享得到子网络，这样大大节省了计算成本；但是这些研究工作都聚焦在CNN上；本文则是将NAS应用到多阶段ViT中，这方面目前的额研究还比较稀缺。

Section V Conclusion

本文将CNN中的2种技术，空间压缩和NAS应用到ViT中，提出了更高效的通过NAS搜索得到的ViT-ResNAS结构，其中残差空间压缩有效的在精度-mac中进行权衡，残差连接也可以提升性能和训练稳定性；而基于权重共享的NAS搜索多阶段网络可以更高效的训练超网，从而获得更好的搜索结果。在ImageNet上的实验结果证明了该网络和方法的有效性。

Appendix

NAS of ViT-Res

ViT-Res-Tiny是在DeiT-Tiny基础上引入两个修改，一是RSR模块（Residual Spatial Reduction）,RSU将整个网络结构分成三段；二是tokenization之前加入三层卷积层。Table 4展示了ViT-Res-Tiny的结构。
Search Space
需要将ViT-ResTiny进一步扩充来构建超网，通过增加attention head h以及增加隐藏层的维度，对于卷积层 值搜索了最后一层的输出通道数；每一个阶段包含3个Transformer block,而中间的block可以选择跳过；因此每一个阶段包含3-6个trnasformer模块。此外，通过扩种每一个block的宽度和深度将ViT-Res扩充到Small和Medium。Table 5总结了ViT-ResNAS的搜索空间。
Table 5展示的是ViT-ResNAS的搜索空间。左侧是ViT-Res-Tiny 右侧是Small和Medium 搜索参数是embedding dimension和transformer block的数量；在每一个block中则搜索attention head的数量，FFN中隐藏层的维度。

Training Details

超网的训练
将原始的训练集分成 sub-train，sub-validation集，sub-validation中包含25K张图像，每类25张。其余用于训练集。超网在sub-train上训练，以及在进化搜索阶段，子网络是在sub-validation上进行的评估。

warm up策略：

逐渐热身。即最开始只在搜索空间采样和训练最大的子网络；随着超网训练的进行，再逐渐对宽度和深度较小的子网络进行采样和训练。总训练过25%后，可以对任意宽度和深度的子网络进行采样和训练。
Training Cost

本文训练使用了8块V100，Tiny网络训练用了32小时，其中超网的训练占用了16.7小时，进化搜索占用了5.5小时，子网的搜索用了34.5小时。
Small和Medium中超网训练用了21小时

Small: 进化搜索6小时 子网训练 41.6小时


Medium:进化搜索6小时 子网训练 45小时


Searched ViT-ResNAS Architecture


Table 6,7,8分别展示了搜索到的Tiny,Small,Medium结构







Limitation



本段说明本文工作的一些局限性。
首先，使用NAS得到的网络性能依赖于人工设计的搜索空间。通过均匀增加ViT-Res-Tiny的宽度和深度来建立搜索空间，这种启发式的方法无需微调。与目前CNN的搜索空间相比，本文对搜索空间研究的较少，优化的也较少。如果进一步优化搜索空间，为每一个架构参数设置一个更好的范围，可能会进一步提升性能。
此外，本文使用了与ViT同样的注意力机制，在处理高分辨率特征图，即序列长度很长时对内存需求很大，因此设计更高效的注意力机制可以解决这一问题，同时进一步提升网络性能。