Press "Enter" to skip to content

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模型的实地应用。

 

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注