AutoFormer: Searching Transformers for Visual Recognition

摘要

最近，基于Transformers的模型在图像分类和检测等视觉任务中显示出了巨大的潜力。 然而，变压器网络的设计是具有挑战性的。 已经观察到，深度、嵌入尺寸和头部的数量在很大程度上影响视觉变形器的性能。 以前的模型基于手工手工配置这些维度。 在这项工作中，我们提出了一个新的一次性架构搜索框架，即AutoFormer，专门用于视觉Transformers搜索。 在超网训练期间，自动前缠绕不同块的重量在同一层。 受益于该战略，训练有素的超级网络允许数千个子网得到非常好的训练。 具体来说，这些继承自超级网络权重的子网的性能与那些从头开始重新训练的子网相当。 此外，搜索模型，我们参考的AutoFormers，超过了最近的先进水平，如ViT和DeiT。 特别是AutoFormer-tiny/small/base在ImageNet上实现了74.7%/81.7%/82.4%的top-1精度，分别为5.7M/22.9M/53.7M参数。 最后，我们通过提供下游基准和蒸馏实验的性能来验证自动成形机的可移植性。  

1、介绍

视觉转换器由于其高模型能力和在获取长期依赖方面的优越潜力，近年来在计算机视觉领域引起了广泛的关注。 在Transformers的基础上，现代最先进的模型，如ViT和DeiT，能够从图像中学习强大的视觉表示，与以前的卷积神经网络模型相比，取得了非常有竞争力的性能。  

然而，Transformers神经结构的设计不是简单的。 例如，如何选择最佳的网络深度、嵌入尺寸和/或头数? 这些因素都是提高模型容量的关键，但找到一个好的组合是困难的。 如图2所示，增加Transformers的深度、头数和MLP比值(多层感知器中隐藏维数与嵌入维数的比值)，首先可以获得较高的准确率，但在达到峰值后会出现过拟合。 增大嵌入维数可以提高模型的性能，但对于较大的模型，精度增益趋于平稳。 这些现象说明了设计最优Transformers架构的挑战。  

                  

以前设计视觉变形器的工作都是基于手工制作，这在很大程度上依赖于人类的专业知识，通常需要大量的试错。 有一些关于使用神经体系结构搜索(NAS)自动化变压器设计的工作。 然而，它们都集中在自然语言任务，如机器翻译，这是相当不同于计算机视觉任务。 因此，很难推广现有的自动搜索算法来找到有效的视觉变换架构。  

在这项工作中，我们提出了一个新的架构搜索算法，名为AutoFormer，致力于寻找纯视觉Transformer模型。 我们的方法主要解决Transformer搜索中的两个挑战。 1)如何将Transformer网络深度、埋入尺寸、机头数等关键因素进行良好组合? 2)如何高效找出适合不同资源约束和应用场景的各种Transformer模型?  

为了应对这些挑战，我们构建了一个涵盖变压器主要变化维度的大搜索空间，包括嵌入维度、头数、查询/键/值维度、MLP比率和网络深度。 这个空间包含了大量具有不同结构和模型复杂性的转换器。 特别是，它允许使用不同结构的积木构造变压器，从而打破了在变压器设计中所有积木共享相同结构的惯例。  

为了解决效率问题，受BigNAS和纤细网络的启发，我们提出了一种超网络训练策略，称为权重纠缠，专门用于变压器架构。 其核心思想是使不同的Transformer块在每一层中共享其公共部件的权重。 一个块中权重的更新会影响其他所有的块作为一个整体，这样不同块的权重在训练过程中最大程度地纠缠在一起。 该策略不同于大多数一次性NAS方法，同一层不同区块的权重是独立的，如图5所示。  

当我们使用所提出的权值纠缠进行Transformer超网训练时，我们观察到一个令人惊讶的现象:它允许超网中的大量子网得到很好的训练，因此这些继承自超网权值的子网的性能与重新训练的子网相当。 这种优势使得我们的方法可以获得数以千计的体系结构，这些体系结构可以满足不同的资源约束，同时保持与从头训练相同的精度水平。 我们在3.4节中详细讨论了重量纠缠的潜在原因。  

 

我们在训练有素的超级网络上执行一个具有模型尺寸约束的进化搜索，以找到有前途的Transformer。 在ImageNet上的实验表明，我们的方法比手工制作的最先进的Transformer模型具有更好的性能。 例如，如图1所示，在22.9M参数下，Autoformer的top-1精度为81.7%，比DeiT-S[50]和it - s /16[13]分别提高了1.8%和2.9%。 此外，当转移到下游视觉分类数据集时，我们的AutoFormer也在更少的参数下表现良好，获得更好的或可与最佳卷积模型(如EfficientNet)相比的结果。  

综上所述，我们在本文中做出了三个主要贡献。 1)据我们所知，这项工作是第一次努力设计一个自动搜索算法寻找视觉变形模型。 2)我们提出了一个简单而有效的训练Transformer超网络的框架。 无需额外的微调或再训练，经过训练的超级网络能够通过直接继承它的权重来生产数千个高质量的变压器。 这种优点使得我们的方法可以搜索不同的模型来适应不同的资源约束。 3)我们搜索的模型AutoFormers在视觉变形器ImageNet上取得了最先进的结果，并在下游任务上表现出了很好的可移植性。  

2、背景

在介绍我们的方法之前，我们首先简要回顾了视觉转换器和一次性NAS的背景。  

2.1、视觉Transformer

Transformer最初是为自然语言任务设计的[52,28,12]。 最近的工作，如ViT和DeiT，显示了它在视觉识别方面的巨大潜力。 在下面，我们给出了视觉转换器的基本管道，它作为AutoFormer的基础架构。  

给定一幅2D图像，我们首先将其均匀地分割成一系列2D补丁，就像自然语言处理中的标记一样。 然后，我们通过线性投影或几个CNN层将补丁平化并转换为d维向量，称为补丁嵌入。 一个可学习的[类]嵌入被注入到序列的头部，以表示整个图像。 位置嵌入被添加到补丁嵌入中以保留位置信息。 然后将组合嵌入馈送到下面描述的变压器编码器。 最后，采用线性层进行最终分类。  

变压器编码器由多头自注意(MSA)块和多层感知器(MLP)块交替组成。 LayerNorm (LN)在每个块之前应用，剩余连接在每个块之后应用。 MSA和MLP的详细资料如下。 多线程Self-Attention (MSA)。 在一个标准的自关注模块中，输入序列将首先线性转换为查询，键和值，其中N是标记的数量，D是嵌入维数，是Q-K-V维数。 然后计算序列中每个元素的所有值的加权和。 权重或注意力基于序列中两个元素之间的成对相似性:  

                  

其中是比例因子。 最后，应用一个完全连接的层。 Multihead self-attention将查询、键和值分割到不同的head中，并并行执行self-attention，并将其连接的输出投影出来。 多层感知器(MLP)。 MLP块由两个完全连接的层组成，具有激活功能，通常为GELU。 在本工作中，我们关注于寻找每一层MLP比率的最佳选择。  

2.2、One-Shot NAS

一次性NAS通常采用权重共享策略，避免对每个子网从零开始进行训练。 将体系结构搜索空间A编码在超网络中，记为，其中W是超级网的权重。  

W在所有候选体系结构中共享，即N中的子网。 在一次性NAS中搜索最优体系结构通常被表述为两阶段优化问题。 第一步是优化权重W ：

          

其中，表示训练数据集上的损失函数。 为了减少内存使用，一次性方法通常从N中采样子网进行优化。 第二阶段是根据中学习到的权重对子网的性能进行排序来搜索架构：

        

其中，采样子网α继承了WA的权重w, 表示体系结构α在验证数据集上的精度为top-1。 由于不可能枚举所有架构进行评估，以往的工作借助于随机搜索[34,4]、进化算法[43,16]或强化学习[40,48]来找到最有前途的架构。  

3、AutoFormer

在本节中，我们首先证明了按照经典的权值共享策略[16]，对每层不同的块使用不同的权值，直接应用一次性NAS进行transformer搜索是不现实的，因为收敛速度慢，性能不理想。在此基础上，提出了视觉变换的权值纠缠策略。最后给出了搜索空间和搜索流水线。

                      

3.1、One-Shot NAS with Weight Entanglement

先前的一次性NAS方法通常在超网络训练期间在架构间共享权值，同时在同一层解耦不同操作者的权值。该策略在卷积神经网络空间上搜索架构时表现良好然而，在transformer搜索空间中，这种经典的搜索策略遇到了困难。1)收敛速度慢。如图4(左)所示，超网络的训练损耗收敛缓慢。原因可能是变压器块的独立训练导致权值在有限的时间内更新。2)低性能。继承经典权值共享策略下一次性超网权值的子网的性能远远低于从头开始训练时的真实性能(如图4右图所示)，这限制了超网的排名能力。此外，在搜索之后，由于权值并没有完全优化，还需要对搜索到的结构进行额外的再训练。受BigNAS和slimmable networks的启发，我们提出了用于视觉转换器架构搜索的权值纠缠训练策略。其核心思想是使不同的变压器块在每一层中共享其公共部件的权重。具体地说，对于具有l层堆栈的子网，我们将其结构和权值表示为：

                           

其中表示第i层的采样块，为块权值。在体系结构搜索过程中，每个层都有多个块可供选择。因此，和实际上是从属于搜索空间的n个候选块集合中选择出来的，其表达式为：

                         

其中是搜索空间中的一个候选块，是它的权值。 权值纠缠策略强制同一层中的不同候选块共享尽可能多的权值。 这要求，对于任意两个块和在同一层，我们有：

                  

              

这样的层内权值共享使得和的权值更新相互纠缠。 任何块的训练都会影响其他块相交部分的权值，如图5所示。 这与一次性NAS中的经典权重共享策略不同，后者在同一层中的构建块是隔离的。 换句话说，在经典的权值共享中，对于任意两个块和，我们。  

需要注意的是，我们提出的权值纠缠策略是针对同构构建块的，比如不同头数的自注意模块，以及不同隐藏维度的多层感知器。 根本原因是同构块在结构上是兼容的，因此权重可以彼此共享。 在实现过程中，对于每一层，我们只需要存储n个齐次候选块中最大块的权值。 剩下的较小的构件可以直接从最大的构件中提取权重。 单次NAS采用了重量纠缠，能够高效、有效地搜索变压器架构，如图4所示。 与经典的权值共享方法相比，我们的权值纠缠策略有三个优势。 1)更快的收敛。 重量纠缠允许每个块比以前的独立训练策略更新更多的时间。 2)内存成本低。 我们现在只需要存储每一层的最大构建块的参数，而不是存储空间中的所有候选元素。 3)更好的子网性能。 我们发现，采用权重纠缠度训练的子网可以达到与从头训练相同的性能。  

3.2、搜索空间

我们设计了一个大的transformer搜索空间，在变压器构件块中包含五个变量因子:嵌入维数、Q-K-V维数、头数、MLP比、网络深度，详见表1和图3。 所有这些因素对模型能力都很重要。 例如，在注意力层中，不同的头部被用来捕捉不同的依赖关系。 然而，最近的研究表明，许多头是多余的。 因此，我们使注意头的数量具有弹性，使每个注意模块可以决定其必要的注意头数量。 另一方面，由于不同的层具有不同的特征表示能力，在构建新模型时，层中不同的隐藏维度可能比固定的隐藏维度更好。 此外，AutoFormer在搜索空间中添加了新的Q-K-V维，并固定了QK-V维与每个块中头部数量的比率。 这种设置使得注意力计算中的比例因子对头部的数量不变性，稳定了梯度，并解耦了不同头部的含义。 此外，我们还结合MLP比和嵌入维数来确定每个块的隐藏维数，这比固定值的搜索空间更大。  

             

按照一次性NAS方法，我们将搜索空间编码为一个超级网络。 也就是说，空间中的每个模型都是超级网络的一部分/子集。 所有子网共享其公共部分的权重。 supernet是空间中最大的模型，其结构如图3所示。 特别地，超级网堆叠空间中定义的最大嵌入维数、Q-KV维数和MLP比的最大变压器块数。 在训练过程中，对所有可能的子网进行统一采样，并更新相应的权值。 根据模型参数的约束，我们将大规模搜索空间划分为三个部分，并将其编码为三个独立的超网，如表1所示。 这样的划分使得搜索算法可以集中于寻找特定参数范围内的模型，用户可以根据其可用资源和应用需求对模型进行专门化。 总的来说，我们的超级网包含了超过的候选架构，覆盖了广泛的模型尺寸。  

3.3、搜索流水线

我们的搜索管道包括两个顺序的阶段。  

阶段1:超网训练与重量纠缠。 在每次训练迭代中，我们均匀采样一个子网，并更新其相应的权值在超级网的重量，同时冻结其余。 具体算法见补充资料附录A。  

第二阶段:资源约束下的进化搜索。 在得到训练后的超网络后，我们对其进行进化搜索以获得最优子网。 根据进化算法的管理者来评估和选择子网。 我们的目标是在最小化模型尺寸的同时最大化分类精度。 在进化搜索的开始阶段，我们选取N个随机结构作为种子。 选取前k个架构作为亲本，通过交叉和变异生成下一代。 对于一个交叉，两个随机选择的候选人被挑选和交叉产生一个新的在每一代。 对于突变，候选序列以的概率先突变其深度。 然后，它以的概率对每个块进行变异，以产生一个新的体系结构。  

3.4、讨论

为什么重量纠缠起作用? 我们猜想有两个根本原因。 1)训练中的正则化。 与卷积神经网络不同，transformer根本没有卷积操作。 它的两个基本组件MSA和MLP只使用完全连接的层。 权重纠缠可以看作是transformer的一种正则化训练策略，在某种程度上类似于dropout的影响。 对小子网进行抽样时，对应的单元不能依赖其他隐藏单元进行分类，从而降低了单元的依赖性。 2)深度薄子网优化。 最近的研究表明，深变难以训练，这与我们在图2中的观察一致。 这是因为在反向传播过程中，深度薄网络中的梯度可能会爆炸或消失。 增加网络的宽度或“过度参数化”将有助于优化。 我们的权重纠缠训练策略以类似的方式帮助优化瘦子网。 宽子网所反向的梯度有助于更新薄子网的权值。 此外，弹性深度与随机深度和深度监督的效应相似，后者也监督浅层。  

4、实验

在本节中，我们首先介绍实现细节和演化搜索设置。 然后，我们分析了提出的权值纠缠策略，并提供了大量从超网络中采样的训练良好的子网来验证其有效性。 最后，我们在几个基准上对AutoFormer的性能进行了评估，并与人工或自动设计的最新模型进行了比较。  

             

              

4.1、实现细节

Supernet训练：我们使用与DeiT类似的配方训练超级网。 具体如表3所示。 数据增强技术，包括RandAugment， Cutmix， Mixup[69]和随机擦除，采用与DeiT相同的超参数，除了重复增强[20]。 图像被分割成16x16大小的补丁。 所有模型都使用PyTorch 1.7实现，并在Nvidia Tesla V100 gpu上进行训练。 进化搜索。 进化搜索的实现遵循与SPOS相同的协议。 为了进行公平的比较，我们保留了用于测试的ImageNet验证集，并将10,000个训练示例(每个类100张图片)的子样本作为验证数据集。 我们将种群大小设置为50，代数设置为20。 每一代我们选择前10个架构作为父母，通过突变和交叉生成子网络。 突变概率Pd和Pm分别设为0.2和0.4。  

进化搜索：进化搜索的实现遵循与SPOS相同的协议。 为了进行公平的比较，我们保留了用于测试的ImageNet验证集，并将10,000个训练示例(每个类100张图片)的子样本作为验证数据集。 我们将种群大小设置为50，代数设置为20。 每一代我们选择前10个架构作为父母，通过突变和交叉生成子网络。 突变概率和分别设为0.2和0.4。  

4.2、消融研究和分析

                 

              

The Efficacy of Weight Entanglement.

我们将AutoFormer与随机搜索和SPOS(经典的权重共享)基线进行比较，以证明权重纠缠的有效性。 对于随机搜索，我们从搜索空间中随机抽取结构，以满足模型尺寸的约束。 对于SPOS，我们将其适应于表1中定义的Transformer搜索空间，其余配置保持与原方法一致。 换句话说，在SPOS超级网络的每一层中，具有不同体系结构参数的Transformer块不共享权值。 例如，在MLP块中，有多个隐藏维度的选择。 每个MLP选择都有自己的权重，彼此独立。 在对SPOS超级网络进行训练后，我们应用同样的进化过程来寻找最有前途的候选架构，并使用与我们的AutoFormer相同的设置对其进行再训练。  

表2给出了ImageNet上的比较。 我们可以观察到:1)经过再训练后，随机搜索和SPOS比我们的方法差2.3%和1.3%，这说明了我们提出的方法的优越性。 2)在不进行再训练的情况下，即直接从超网络中继承权值，权值纠缠训练策略的效果明显优于经典的权值共享。 纠缠权值可以产生训练良好的子网，与重新训练的子网非常接近。 我们推测SPOS在Transformer空间的性能相对较差主要是由于训练不足。 我们还观察到，如果我们在SPOS中训练更多的超级网，性能可以缓慢地提高。 然而，它的训练成本大大高于我们提出的权值纠缠策略。 图6绘制了在搜索过程中从训练过的超网络中采样的结构的数量上的精度。 每一代都描述了前50名候选人。 很明显，在超级网络上的进化搜索比随机搜索基线更有效。  

      

4.3、ImageNet上的结果

我们在ImageNet上搜索AutoFormer，找到多个不同参数尺寸的transformer型号。 所有这些模型都直接从超网络继承权重，而不需要额外的再训练和其他后处理。 性能报告见表5和图1。 很明显，我们的AutoFormer模型家庭实现更高的准确性比最近手工制作的最先进的transformer模型，如ViT和DeiT。 特别是，使用约23M参数，我们的小型模型，即AutoFormer-S，实现了81.7%的顶级精度，分别比DeiT-S和it - s /16好1.8%和2.9%。  

与vanilla CNN模型相比，AutoFormer也是有竞争力的。 如图1所示，我们的AutoFormers比手工设计的ResNet、ResNeXt和DenseNet表现得更好，这表明纯transformer模型在视觉表示方面的潜力。 然而，基于transformer的视觉模型，包括AutoFormer，现在仍然不如基于倒立残块的模型，如MobileNetV3和EfficientNet。 原因是反向残差是针对边缘设备进行优化的，所以模型尺寸和FLOPs都比视觉变形器小得多。  

4.4、迁移学习结果

分类：

我们将Autoformer转换到一个常用的识别数据集:1)一般分类:CIFAR-10和CIFAR-100; 2)细粒度分类:Stanford Car， FLowers和OxfordIII Pets。 我们遵循与DeiT相同的训练设置，它采用ImageNet预先训练的检查点，并对新的数据集进行微调。 表6显示了top-1精度的结果:1)与最先进的ConvNets相比，AutoFormer在参数较少的情况下接近最佳结果，差距可以忽略; 2)与基于变形器的模型相比，AutoFormer在所有数据集上都能获得更好的结果，而且参数更少(约4x)。 一般来说，我们的AutoFormer能够在比现有型号(包括ResNet、RegNet、EfficientNet、ViT和DeiT)少一个数量级的参数下达到相同或更好的精度。  

   

蒸馏：

AutoFormer也与知识精馏(KD)正交，因为我们专注于寻找一个有效的架构，而KD则专注于更好地训练给定的架构。 通过提取RegNetY-32GF教师的硬标签，将KD与AutoFormers结合，可以进一步提高成绩，分别从74.7%/81.7%/82.4%提高到75.7%/82.4%/82.9%。  

5、相关工作

Vision Transformer

Transformer最初是为语言建模而提出的，最近被应用于计算机视觉。 它在识别、检测和分割等多种任务中显示出了很好的潜力。 在视觉中使用transformer的一种简单方法是将卷积层与自我注意模块相结合。 最近的研究也在这方面取得了进展。 最近，Dosovitskiy等人引入了Vision Transformer (ViT)，一种用于视觉识别的纯Transformer架构，不使用任何卷积。 ViT在非重叠图像补丁的线性投影上堆叠Transformer块。 当使用广泛的、不公开的图像数据集(JFT-300M, 3亿张图像)进行训练时，它呈现了很有希望的结果。 最近的DeiT模型验证了大规模数据是不需要的。 仅使用Imagenet也可以产生具有竞争力的无卷积变压器。 然而，现有的Transformer模型都建立在手工设计的基础上，这是昂贵的工程和容易出错。 在这项工作中，我们提出了第一个努力自动化设计的视觉Transformer与神经结构搜索。  

神经架构搜索

NAS用于自动化网络设计的兴趣越来越大。 早期的方法使用强化学习或进化算法搜索网络。 最近的研究大多采用一次性权重共享策略来摊销搜索成本。 其关键思想是训练一个过度参数化的超网络模型，然后在各个子网之间共享权值。 其中SPOS比较简单 在每次迭代中，它只采样一个随机路径，并使用一个批处理数据训练路径。 一旦训练过程完成，就可以通过继承共享权值来对子网进行排序。 然而，大多数权重共享方法在确定最佳架构后需要额外的再训练步骤。  

最近的工作，OFA， BigNAS和slimmable网络通过训练一个一次性的超级网络来缓解这个问题。 尽管AutoFormer与这些方法在训练一个统一超级网络的概念上有相似之处，但这些方法的设计目的是搜索卷积网络，而不是视觉Transformer，后者具有非常不同的架构特征。 具体来说，AutoFormer考虑了多头自注意和MLP的设计，这是Transformer模型的独特之处，并给出了搜索尺寸的专用设计，如第3.2节所述。 此外，BigNAS采用了多种精心制作的技术，如三明治训练、原位蒸馏、正则化等。 OFA提出了一种逐步收缩的方法，通过逐步提取整个网络来获得更小的子网。 相比之下，AutoFormer是简单和有效的，实现一次性所有的训练没有这些技术，并可以很容易地扩展到搜索其他视觉Transformer变体。  

对于变压器，应用NAS改进其架构的研究很少。 这些方法主要针对自然语言处理任务。 其中，和我们最相似的是HAT。 除了任务之间的区别，HAT在搜索后需要额外的再训练或微调步骤，而AutoFormer不需要，这是关键的区别。 另一个区别是搜索空间。 HAT搜索一个编码器-解码器转换器结构，而我们的是一个纯编码器。 有两个并行的作品，BossNAS和CvT，探索不同于我们的搜索空间。 BossNAS搜索cnn的混合Transformer，而我们的搜索纯Transformer。 CvT提出了一个新的体系结构家族，并寻找它们的步幅和内核大小。 由于搜索空间的不同，我们在本工作中不进行比较。  

6、结论

在这项工作中，我们提出了一种新的一次性架构搜索方法，AutoFormer，专门用于Transformer搜索。 AutoFormer配备了训练策略，重量纠缠。 在这种策略下，搜索空间中的子网几乎完全训练。 大量实验表明，所提出的算法可以改善超网络的训练，找到有前途的体系结构。 我们搜索的AutoFormers在ImageNet视觉Transformer中实现了最先进的结果。 此外，自动成型机可以很好地完成下游的几个分类任务，并可以通过蒸馏进一步改进。 在未来的工作中，我们有兴趣进一步丰富搜索空间，包括卷积作为新的候选算子。 将权值纠缠应用于卷积网络搜索或对权值纠缠进行理论分析是另一个潜在的研究方向。