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

关注公众号，发现CV技术之美

本文分享论文『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的维度扩展到隐藏层大小 ，第二层线性层将维度投影回原来的大小