An image is worth 16x16 words: Transformers for image recognition at scale
郑重声明：原文参见标题，如有侵权，请联系作者，将会撤销发布！
Published as a conference paper at ICLR 2021
ABSTRACT
虽然Transformer结构已成为自然语言处理任务的事实标准，但其在计算机视觉中的应用仍然有限。在视觉中，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组件，同时保持其整体结构。我们表明，这种对CNN的依赖是不必要的，并且直接应用于图像块序列的纯Transformer可以在图像分类任务中表现得很好。当在大量数据上进行预训练并转移到多个中型或小型图像识别基准（ImageNet、CIFAR-100、VTAB等）时，与最先进的卷积网络相比，视觉Transformer (ViT)获得了优异的结果，同时需要更少的计算资源进行训练。1
1 微调代码和预训练模型可在https://github.com/google-research/vision_transformer中获取
1 INTRODUCTION
基于自注意的架构，特别是Transformers（Vaswani et al., 2017），已经成为自然语言处理（NLP）中的首选模型。主要的方法是在大型文本语料库上进行预训练，然后在较小的特定任务数据集上进行微调（Devlin et al., 2019）。由于Transformer的计算效率和可扩展性，训练具有超过100B参数的空前规模的模型成为可能（Brown et al., 2020; Lepikhin et al., 2020）。随着模型和数据集的增长，仍然没有性能饱和的迹象。
然而，在计算机视觉中，卷积架构仍然占主导地位（LeCun et al., 1989; Krizhevsky et al., 2012; He et al., 2016）。受NLP成功的启发，多个工作尝试将类CNN架构与自关注力相结合（Wang et al., 2018; Carion et al., 2020），其中一些工作完全取代了卷积（Ramachandran et al., 2019; Wang et al., 2020a）。后一种模型虽然理论上有效，但由于使用了专门的注意力模式，尚未在现代硬件加速器上有效扩展。因此，在大规模图像识别中，经典的类ResNet架构仍然是最先进的（Mahajan et al., 2018; Xie et al., 2020; Kolesnikov et al., 2020）。
受NLP中Transformer缩放成功的启发，我们尝试将标准Transformer直接应用于图像，并尽可能减少修改。为了做到这一点，我们将图像分割成块，并提供这些块的线性嵌入序列作为Transformer的输入。图像块的处理方式与NLP应用中的token（单词）相同。我们以监督的方式对模型进行图像分类训练。
当在中等大小的数据集（如ImageNet）上训练时，如果没有强大的正则化，这些模型的精度会比同等大小的ResNet低几个百分点。这种看似令人沮丧的结果可能是意料之中的：Transformer缺乏CNN固有的一些归纳偏差，如平移等变性和局部性，因此在数据量不足的情况下训练时不能很好地泛化。
然而，如果在较大的数据集（14M-300M图像）上训练模型，则图片会发生变化。我们发现大规模训练胜过归纳偏差。当以足够的规模进行预训练并迁移到数据点较少的任务中时，我们的视觉Transformer（ViT）会获得优异的结果。当在公共ImageNet-21k数据集或内部JFT-300M数据集上进行预训练时，ViT在多个图像识别基准上接近或优于现有技术。特别是，最佳模型在ImageNet上达到88.55%的准确率，在ImageNet-ReaL上达到90.72%，在CIFAR-100上达到94.55%，在由19个任务组成的VTAB套件上达到77.63%。
2 RELATED WORK
3 METHOD
在模型设计中，我们尽可能地遵循原始Transformer (Vaswani et al., 2017)。这种有意简化的设置的一个优点是，可扩展的NLP Transformer架构及其高效实现几乎可以开箱即用。
3.1 VISION TRANSFORMER (VIT)
该模型的概述如图1所示。标准Transformer接收token嵌入的1D序列作为输入。为了处理2D图像，我们将图像
重塑为平坦的2D块(patch)序列
，其中(H, W)是原始图像的分辨率，C是通道数，(P, P)是每个图像块的分辨率，N=HW/P2是得到的块数，这也是Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的隐矢量大小D，因此我们将块压平，并使用可训练的线性投影映射到D维(公式1)。我们将此投影的输出称为块嵌入。
类似于BERT的[class] token，我们为嵌入块序列预先准备了可学习嵌入(
)，其在Transformer编码器的输出处的状态(
)用作图像表征y (公式4)。在预训练和微调期间，分类头都连接到。分类头在预训练时由具有一个隐藏层的MLP实现，在微调时由单个线性层实现。
位置嵌入被添加到块嵌入以保留位置信息。我们使用标准的可学习1D位置嵌入，因为我们没有观察到使用更先进的2D感知位置嵌入带来的显著性能增益(附录D.4)。所得到的嵌入矢量序列用作编码器的输入。
Transformer编码器(Vaswani et al., 2017)由多头自注意力(MSA，见附录A)和MLP块(公式2, 3)的交替层组成。在每个块之前应用层归一化(LN)，在每个块之后应用残差连接(Wang et al., 2019; Baevski & Auli, 2019)。
MLP包含两个具有GELU非线性的层。
感应偏差。我们注意到，视觉Transformer比CNN具有小得多的图像特异性感应偏差。在CNN中，局部性、二维邻域结构和平移同变性被融入整个模型的每一层。在ViT中，只有MLP层是局部且平移同变性的，而自注意层是全局的。二维邻域结构的使用非常谨慎：在模型开始时，通过将图像切割成块，并在微调时调整不同分辨率图像的位置嵌入(如下所述)。除此之外，在初始化时的位置嵌入不携带关于块的2D位置的信息，并且块之间的所有空间关系必须从头开始学习。
混合架构。作为原始图像块的替代方案，输入序列可以由CNN (LeCun et al., 1989)的特征图形成。在该混合模型中，将块嵌入投影E (公式1)应用于从CNN特征图提取的块。作为一种特殊情况，块可以具有1x1的空间大小，这意味着通过简单地将特征图的空间维度平坦化并投影到Transformer维度来获得输入序列。如上所述，添加了分类输入嵌入和位置嵌入。
3.2 FINE-TUNING AND HIGHER RESOLUTION
通常，我们在大型数据集上预训练ViT，并微调到(较小的)下游任务。为此，我们移除预训练的预测头，并添加零初始化的D×K前馈层，其中K是下游类的数量。以比预训练更高的分辨率进行微调通常是有益的(Touvron et al., 2019; Kolesnikov et al., 2020)。当馈送更高分辨率的图像时，我们保持块大小相同，这导致更大的有效序列长度。视觉Transformer可以处理任意序列长度(最高可达内存约束)，然而，预先训练的位置嵌入可能不再有意义。因此，我们根据预训练的位置嵌入在原始图像中的位置对其进行2D插值。注意，这种分辨率调整和块提取是关于图像的2D结构的感应偏差被手动注入到视觉Transformer中的唯一点。
4 EXPERIMENTS
我们评估了ResNet、视觉Transformer (ViT)和混合体的表征学习能力。为了了解每个模型的数据需求，我们在不同大小的数据集上进行了预训练，并评估了许多基准任务。当考虑到预训练模型的计算成本时，ViT表现非常好，以较低的预训练成本在大多数识别基准上达到了最先进的水平。最后，我们使用自监督进行了一个小实验，并表明自监督ViT在未来充满希望。
4.1 SETUP
数据集。为了探索模型的可扩展性，我们使用了ILSVRC-2012 ImageNet数据集，该数据集具有1k个类和1.3M个图像(我们在下文中将其称为ImageNet)，其超集ImageNet-21k具有21k个类和14M个图像(Deng et al., 2009)，以及具有18k个类，303M个高分辨率图像的JFT (Sun et al., 2017)。我们根据Kolesnikov等人(2020)，将预训练数据集与下游任务的测试集进行去重。我们将在这些数据集上训练的模型转移到几个基准任务中：原始验证标签和清理后的ReaL标签上的ImageNet (Beyer et al., 2020)、CIFAR-10/100 (Krizhevsky, 2009)、Oxford IIIT Pets (Parkhi et al., 2012)和Oxford Flowers-102 (Nilsback & Zisserman, 2008)。对于这些数据集，按照Kolesnikov等人(2020)进行预处理。
我们还对19-任务VTAB分类套件进行了评估(Zhai et al., 2019b)。VTAB评估不同任务的低数据传输，每个任务使用1000个训练示例。任务分为三组：Natural——任务(如上述)，宠物、CIFAR等。Specialized——任务(如医疗和卫星图像)以及Structured——需要几何理解的任务(如定位)。
模型变体。如表1所示，我们基于用于BERT的ViT配置(Devlin et al., 2019)。“Base”和“Large”模型直接取自BERT，我们增加了更大的“Huge”模型。在下文中，我们使用简短的符号来表示模型大小和输入块大小：例如，ViT-L/16表示输入块大小为16×16的“Large”变体。请注意，Transformer的序列长度与块大小的平方成反比，因此具有较小块大小的模型在计算上更昂贵。
对于基线CNN，我们使用ResNet (He et al., 2016)，但将批归一化层(Ioffe & Szegedy, 2015)替换为组归一化(Wu & He, 2018），并使用标准化卷积(Qiao et al., 2019)。这些修改改善了传输(Kolesnikov et al., 2020)，我们将修改后的模型称为“ResNet (BiT)”。对于混合体，我们将中间特征图输入到具有一个“像素”的块大小的ViT中。为了实验不同的序列长度，我们要么（i）取常规ResNet50的第4阶段的输出，要么（ii）去除第4阶段，在第3阶段放置相同数量的层（保持层的总数），并取该扩展的第3阶段的输出。选项(ii)导致4倍的序列长度和更昂贵的ViT模型。
训练和微调。我们使用Adam (Kingma & Ba, 2015)训练包括ResNet在内的所有模型，其中β1=0.9，β2=0.999，批量大小为4096，并应用0.1的高权重衰减，我们发现这对所有模型的迁移都很有用(附录D.1显示，与常见做法相比，Adam在我们的设置中比用于ResNet的SGD表现得略好)。我们使用线性学习率预热和衰减，详见附录B.1。对于微调，我们对所有模型使用具有动量的SGD，批量大小为512，请参阅附录B.1.1。对于表2中的ImageNet结果，我们以更高的分辨率进行了微调：ViT-L/16为512，ViT-H/14为518，还使用了Polyak & Juditsky (1992)的平均值，因子为0.9999 (Ramachandran et al., 2019; Wang et al., 2020b)。
度量。我们通过小样本或微调精度在下游数据集上报告结果。微调精度捕捉每个模型在各自数据集上进行微调后的性能。通过解决将训练图像子集的(冻结的)表征映射到{−1, 1}K目标矢量的正则化最小二乘回归问题，可以获得小样本精度。这种公式使我们能够以闭合项恢复精确解。虽然我们主要关注微调性能，但有时我们会使用线性小样本精度进行快速实时评估，因为微调成本太高。
4.2 COMPARISON TO STATE OF THE ART
4.3 PRE-TRAINING DATA REQUIREMENTS
视觉Transformer在大型JFT-300M数据集上进行预训练时表现良好。与ResNet相比，视觉的归纳偏差更少，数据集的大小有多重要？我们进行了两个系列的实验。
首先，我们在越来越大的数据集上预训练ViT模型：ImageNet、ImageNet-21k和JFT300M。为了提高较小数据集的性能，我们优化了三个基本的正则化参数——权重衰减、随机失活和标签平滑。图3显示了微调到ImageNet后的结果(其他数据集的结果如表5所示)2。当在最小的数据集ImageNet上预训练时，尽管(适度)正则化，但与ViT-Base模型相比，ViT-Large模型表现不佳。通过ImageNet-21k的预训练，他们的表现是相似的。只有JFT-300M，我们才能看到更大模型的全部好处。图3还显示了性能由不同大小的BiT模型跨越的区域。BiT CNN在ImageNet上的表现优于ViT，但在更大的数据集中，ViT的表现更好。
其次，我们在9M、30M和90M的随机子集以及完整的JFT300M数据集上训练我们的模型。我们不对较小的子集执行额外的正则化，并对所有设置使用相同的超参数。通过这种方式，我们评估了固有的模型属性，而不是正则化的影响。然而，我们确实使用了早期停止，并报告了在训练过程中实现的最佳验证准确性。为了节省计算，我们报告了很少的线性精度，而不是完全的微调精度。图4包含了结果。在较小的数据集上，视觉Transformer比具有可比计算成本的ResNet过拟合更多。例如，ViT-B/32比ResNet50稍快；它在9M子集上表现得更差，但在90M+子集上表现更好。ResNet152x2和ViT-L/16也是如此。这一结果强化了一种直觉，即卷积归纳偏差对于较小的数据集是有用的，但对于较大的数据集，直接从数据中学习相关模式是足够的，甚至是有益的。
总体而言，ImageNet上的小样本结果(图4)以及VTAB上的低数据结果(表2)似乎有望实现非常低的数据传输。进一步分析ViT的小样本特性是未来工作的一个令人兴奋的方向。
2 请注意，ImageNet预训练的模型也经过了微调，但还是在ImageNet上。这是因为在微调过程中分辨率的提高也提高了性能。
4.4 SCALING STUDY
4.5 INSPECTING VISION TRANSFORMER
为了开始了解视觉Transformer是如何处理图像数据的，我们分析了它的内部表征。视觉Transformer的第一层将压扁的块线性投影到较低维空间中(公式1)。图7(左)显示了学到的嵌入滤波器的顶部主要组件。这些成分类似于每个块内精细结构的低维表征的合理基函数。
在投影之后，学到的位置嵌入被添加到块表征中。图7(中)显示，该模型学习在位置嵌入的相似性中对图像内的距离进行编码，即更接近的块往往具有更相似的位置嵌入。此外，出现行-列结构；同一行/列中的块具有相似的嵌入。最后，正弦结构有时对于较大的网格是明显的(附录D)。位置嵌入学会了表示2D图像拓扑，这解释了为什么手工制作的2D感知嵌入变体不能产生改进(附录D.4)。
自注意力允许ViT集成整个图像中的信息，即使是在最底层。我们调查网络在多大程度上利用了这种能力。具体来说，我们根据注意力权重计算图像空间中信息整合的平均距离(图7，右)。这种“注意距离”类似于CNN中的感受野大小。我们发现，最底层的一些头已经处理了大部分图像，这表明模型确实使用了全局集成信息的能力。其他注意力头在低层中的注意力距离一直很小。这种高度局部的注意力在Transformer之前应用ResNet的混合模型中不太明显(图7，右)，这表明它可能与CNN中的早期卷积层具有类似的功能。此外，注意力距离随着网络深度的增加而增加。在全局范围内，我们发现该模型关注与分类在语义上相关的图像区域(图6)。
4.6 SELF-SUPERVISION
Transformer在NLP任务中表现出令人印象深刻的性能。然而，它们的成功很大程度上不仅源于它们出色的可扩展性，还源于大规模的自监督预训练(Devlin et al., 2019; Radford et al., 2018)。我们还模仿BERT中使用的掩码语言建模任务，对用于自监督的掩码块预测进行了初步探索。通过自监督预训练，我们较小的ViT-B/16模型在ImageNet上实现了79.9%的准确率，与从头开始训练相比显著提高了2%，但仍落后于监督预训练4%。附录B.1.2包含更多详细信息。我们将对比预训练的探索留给未来的工作(Chen et al., 2020b; He et al., 2020; Bachman et al., 2019; Henaff et al., 2020)。
5 CONCLUSION
我们已经探索了Transformer在图像识别中的直接应用。与先前在计算机视觉中使用自注意力的工作不同，除了最初的块提取步骤外，我们没有在架构中引入图像特定的归纳偏差。相反，我们将图像解释为一系列块，并通过NLP中使用的标准Transformer编码器进行处理。这种简单但可扩展的策略与大型数据集上的预训练相结合，效果出奇地好。因此，视觉Transformer在许多图像分类数据集上达到或超过了现有技术，同时预训练相对便宜。
虽然这些初步成果令人鼓舞，但仍存在许多挑战。一种是将ViT应用于其他计算机视觉任务，如检测和分割。我们的结果，加上Carion等人(2020)的结果表明了这种方法的前景。另一个挑战是继续探索自监督的预训练方法。我们的初步实验表明，自监督预训练有所改进，但与大规模监督预训练相比仍有很大差距。最后，ViT的进一步扩展可能会提高性能。
APPENDIX
A MULTIHEAD SELF-ATTENTION
B EXPERIMENT DETAILS
B.1 TRAINING
表3总结了我们针对不同模型的训练设置。我们发现，当在ImageNet上从头开始训练模型时，强正则化是关键。当使用时，Dropout应用于除qkv投影外的每个密集层之后，并直接应用于添加位置嵌入到块嵌入之后。混合模型的训练与ViT模型完全相同。最后，所有的训练都是在224分辨率上完成的。
B.1.1 FINE-TUNING
我们使用动量为0.9的SGD对所有ViT模型进行微调。我们对学习率进行了小型网格搜索，见表4中的学习率范围。为此，我们使用来自训练集的小型子分割(10%的Pets and Flowers，2%的CIFAR，1%的ImageNet)作为开发集，并对剩余数据进行训练。对于最终结果，我们对整个训练集进行训练，并对各自的测试数据进行评估。对于微调ResNet和混合模型，我们使用完全相同的设置，唯一的例外是ImageNet，我们在学习率扫描中添加了另一个值0.06。此外，对于ResNet，我们还运行了Kolesnikov等人(2020)的设置，并在这次运行和我们的横扫中选择最佳结果。最后，如果没有另行提及，所有微调实验都以384分辨率运行(以与训练不同的分辨率运行微调是常见的做法(Kolesnikov et al., 2020))。
当将ViT模型转移到另一个数据集时，我们移除整个头(两个线性层)，并用一个零初始化的线性层替换它，该线性层输出目标数据集所需的类数。我们发现这比简单地重新初始化最后一层要稳健一些。
对于VTAB，我们遵循Kolesnikov等人(2020)的协议，并对所有任务使用相同的超参数设置。我们使用0.01的学习率，并训练2500步(表4)。我们通过对两个学习率和两个时间表进行小范围扫描，并在200个示例验证集中选择VTAB得分最高的设置，来选择该设置。我们遵循Kolesnikov等人(2020)中使用的预处理，只是我们不使用特定于任务的输入分辨率。相反，我们发现视觉Transformer从所有任务的高分辨率(384×384)中获益最多。
B.1.2 SELF-SUPERVISION
C ADDITIONAL RESULTS
D ADDITIONAL ANALYSES
D.1 SGD VS. ADAM FOR RESNETS
D.2 TRANSFORMER SHAPE
我们对Transformer架构的不同维度进行了消融，以找出哪些最适合扩展到非常大的模型。图8显示了不同配置下ImageNet上的5-样本性能。所有配置都基于具有8层的ViT模型，D=1024，DMLP=2048，块大小为32，所有线的交点。我们可以看到，缩放深度会带来最大的改进，这些改进在64层之前都是清晰可见的。然而，在16层之后已经可以看到收益递减。有趣的是，缩放网络的宽度似乎会导致最小的变化。在不引入参数的情况下，减小块大小并因此增加有效序列长度显示出令人惊讶的稳健改进。这些发现表明，计算可能比参数的数量更能预测性能，而且缩放应该强调深度而不是宽度(如果有的话)。总的来说，我们发现按比例缩放所有维度会带来稳健的改进。
D.3 HEAD TYPE AND CLASS TOKEN
为了尽可能接近原始Transformer模型(类似于BERT)，我们使用了一个额外的[class] token，该token被用作图像表征(位于最前)。然后，通过小型多层感知器(MLP)将该token的输出转换为类预测，其中tanh作为单个隐藏层中的非线性。
这个设计是从文本的Transformer模型中继承的，我们在整个主论文中都使用它。最初尝试只使用图像块嵌入，全局平均池化(GAP)，然后使用线性分类器——就像ResNet的最终特征图一样——效果非常差。然而，我们发现这既不是由于额外的token，也不是由于GAP操作。相反，性能上的差异完全可以用不同学习率的要求来解释，见图9。
D.4 POSITIONAL EMBEDDING
D.5 EMPIRICAL COMPUTATIONAL COSTS
D.6 AXIAL ATTENTION
D.7 ATTENTION DISTANCE
D.8 ATTENTION MAPS
D.9 OBJECTNET RESULTS
D.10 VTAB BREAKDOWN
