[论文精读]VIT:AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

原文链接：https://arxiv.org/abs/2010.11929

摘要

虽然Transformer架构已成为自然语言处理任务的标准模型，但它在计算机视觉领域的应用仍然有限。在视觉领域，注意力机制要么与卷积网络结合使用，要么用来替换卷积网络的某些组件，同时保持其总体结构不变。我们的研究表明，这种对CNN的依赖于并非必要的，直接将纯粹的Transformer应用于图像块序列，也可以在图像分类任务上表现得非常好。当在大量数据上预训练，并转移到多个中等或小型图像识别基准（如ImageNet、CIFAR-100、VTAB等）上时，视觉Transformer（ViT）与最先进的卷积网络相比，获得了优异的结果，同时在训练上需要的计算资源大大减少。

1.引言

基于自注意力的架构，尤其是Transformers，已经成为自然语言处理（NLP）中的首选模型。主流方法是在大型文本语料库上进行预训练，然后在较小的特定任务数据集上进行微调（Devlin等人，2019年）。得益于Transformers的计算效率和可扩展性，现在可以训练前所未有的大型模型，拥有超过1000亿个参数（Brown等人，2020年；Lepikhin等人，2020年）。随着模型和数据集的增长，其性能仍未显示出饱和的迹象。

受到NLP中Transformer扩展成功的启发，我们尝试将标准Transformer直接应用于图像，尽可能少地进行修改。为此，我们将图像分割成多个块，并将这些块的线性嵌入序列作为输入提供给Transformer。图像块被视为与NLP应用中的令牌（词语）相同的处理方式。我们以监督方式训练模型进行图像分类。

当在中等大小的数据集（如ImageNet）上训练时，如果没有强大的正则化，这些模型的准确率会略低于同等大小的ResNets几个百分点。这种看似令人沮丧的结果可能是预料之中的：Transformer缺乏CNN固有的归纳偏见，如平移不变性和局部性，因此在数据量不足时无法很好地泛化。

然而，如果在更大的数据集（1400万至3亿张图片）上训练这些模型，情况就会改变。我们发现大规模训练胜过归纳偏见。我们的视觉Transformer（ViT）在进行足够规模的预训练并转移到数据点较少的任务时，获得了优异的结果。当在公共的ImageNet-21k数据集或内部的JFT-300M数据集上预训练时，ViT在多个图像识别基准上达到或超过了最先进的水平。特别是，最佳模型在ImageNet上达到了88.55%的准确率，在ImageNet-ReaL上达到了90.72%的准确率，在CIFAR-100上达到了94.55%的准确率，在VTAB的19个任务上达到了77.63%的准确率。

2.相关工作

Transformer由Vaswani等人在2017年提出，用于机器翻译，并从此成为许多NLP任务中的最先进方法。基于Transformer的大型模型通常在大型语料库上进行预训练，然后针对具体任务进行微调：BERT（Devlin等人，2019年）使用去噪自监督预训练任务，而GPT系列工作则使用语言建模作为其预训练任务（Radford等人，2018年；2019年；Brown等人，2020年）。

将自注意力机制简单地应用到图像上将要求每个像素与其他像素进行注意力计算。由于这种方法的计算成本与像素数量的平方成正比，它无法扩展到实际的输入大小。因此，在图像处理中应用Transformers时，过去尝试了几种近似方法。Parmar等人（2018年）仅在每个查询像素的局部邻域内应用自注意力，而非全局。这样的局部多头点积自注意力模块可以完全替代卷积（Hu等人，2019年；Ramachandran等人，2019年；Zhao等人，2020年）。在另一条研究路线上，稀疏Transformer（Child等人，2019年）采用可扩展的全局自注意力近似，以适用于图像。另一种扩展注意力的方法是在不同大小的块中应用它（Weissenborn等人，2019年），在极端情况下仅沿单个轴线应用（Ho等人，2019年；Wang等人，2020a）。许多专门在计算机视觉任务上的注意力架构展示了不错的结果，但需要复杂的工程设计来在硬件加速器上有效实现。

与我们的研究最相关的是Cordonnier等人（2020年）的模型，该模型从输入图像中提取2×2大小的块，并在其上应用完全自注意力。这个模型与ViT非常相似，但我们的工作进一步证明了大规模预训练使得普通的transformers能够与（甚至优于）最先进的CNNs竞争。此外，Cordonnier等人（2020年）使用的小块大小为2×2像素，这使得模型仅适用于低分辨率图像，而我们的处理也适用于中等分辨率图像。

也有很多人对结合卷积神经网络（CNNs）和自注意力形式感兴趣，例如通过增强用于图像分类的特征图（Bello等人，2019年）或通过使用自注意力进一步处理CNN的输出，例如用于对象检测（Hu等人，2018年；Carion等人，2020年）、视频处理（Wang等人，2018年；Sun等人，2019年）、图像分类（Wu等人，2020年）、无监督对象发现（Locatello等人，2020年）或统一的文本-视觉任务（Chen等人，2020c；Lu等人，2019年；Li等人，2019年）。

另一个最近相关的模型是图像GPT（iGPT）（Chen等人，2020a），它在降低图像分辨率和颜色空间后将Transformers应用于图像像素。该模型以无监督方式作为生成模型进行训练，生成的表示可以进一步进行微调或线性探测，以评估分类性能，最高在ImageNet上达到72%的准确率。

我们的工作增加了越来越多的论文，这些论文探索了比标准ImageNet数据集更大规模的图像识别。使用额外的数据来源可以在标准基准测试上实现最先进的结果（Mahajan等人，2018年；Touvron等人，2019年；Xie等人，2020年）。此外，Sun等人（2017年）研究了CNN性能随数据集大小的扩展情况，而Kolesnikov等人（2020年）和Djolonga等人（2020年）对CNN从大规模数据集如ImageNet-21k和JFT-300M进行迁移学习进行了实证探索。我们也关注这两个数据集，但是是训练Transformers而不是之前作品中常使用的基于ResNet的模型。

3.方法

在模型设计中，我们尽可能紧密地遵循原始的Transformer（Vaswani等人，2017年）。这种故意简化的设置的一个优点是，可扩展的NLP Transformer架构及其高效的实现几乎可以即拿即用。

3.1视觉Transformer（VIT）

模型的概览如图1所示。标准的Transformer接受一维令牌嵌入序列作为输入。为了处理2D图像，我们将图像$x\in {\mathbb{R}}^{H\times W\times C}$重塑为一个展平的2D块序列 $x_p\in {\mathbb{R}}^{N\times (P^2\cdot C)}$，其中(H, W)是原始图像的分辨率,C是通道数,(P, P) 是每个图像块的分辨率，而 $N=HW/P^2$ 是生成的块数，这也作为Transformer的有效输入序列长度。Transformer在所有层中使用恒定的潜在向量大小D，因此我们将块展平并通过可训练的线性投影映射到D维（方程1）。我们将这个投影的输出称为块嵌入。

类似于BERT的[class]令牌，我们在嵌入图块序列的前面添加一个可学习的嵌入$(z_0^0=x_{class})$，其在Transformer编码器$(z_L^0)$输出时的状态作为图像表示y（方程4）。在预训练和微调期间，分类头都连接在$z_L^0$上。预训练时，分类头由一个带有一层隐藏层的MLP实现，微调时则由一个单一的线性层实现。

为了保留位置信息，我们将位置嵌入添加到图块嵌入中。我们使用标准的可学习的一维位置嵌入，因为我们没有观察到使用更高级的二维位置嵌入能带来显著的性能提升。最终的嵌入向量序列作为编码器的输入。

Transformer编码器（Vaswani等人，2017年）由交替的多头自注意力层（MSA，见附录A）和MLP块（方程2，3）组成。每个块前应用层归一化（LN），每个块后应用残差连接。MLP包含两层，并使用GELU非线性激活函数。
归纳偏好我们注意到，视觉Transformer比CNNs具有更少的图像特定归纳偏好。在CNNs中，局部性、二维邻域结构和平移不变性贯穿整个模型的每一层。而在ViT中，只有MLP层是局部的和平移不变的，而自注意力层则是全局的。二维邻域结构的使用非常有限：只有在模型的开头通过将图像切割成块和微调时为不同分辨率的图像调整位置嵌入时使用（如下所述）。除此之外，初始化时的位置嵌入不携带任何关于图块二维位置的信息，所有图块之间的空间关系都必须从零开始学习。
混合架构作为对原始图像块的替代，输入序列可以由CNN的特征图（LeCun等人，1989年）形成。在这种混合模型中，图块嵌入投影 E（方程1）应用于从CNN特征图中提取的图块。作为一种特殊情况，图块可以具有1x1的空间大小，这意味着输入序列是通过简单地展平特征图的空间维度并投影到Transformer维度获得的。分类输入嵌入和位置嵌入的添加如上所述。

3.2微调和更高的分辨率

通常，我们在大数据集上预训练ViT，然后对（较小的）下游任务进行微调。为此，我们移除预训练的预测头，并添加一个零初始化的 $D\times K$ 前馈层，其中$K$是下游任务的类别数量。在微调时，使用比预训练更高的分辨率通常是有益的（Touvron等人，2019年；Kolesnikov等人，2020年）。当输入更高分辨率的图像时，我们保持图块大小不变，这会导致更长的有效序列长度。视觉Transformer可以处理任意序列长度（在内存限制内），但预训练的位置嵌入可能不再有意义。因此，我们根据预训练位置嵌入在原始图像中的位置进行二维插值。需要注意的是，这种分辨率调整和图块提取是将关于图像二维结构的归纳偏差手动注入视觉Transformer的唯一方式。

4.实验

我们评估了ResNet、视觉Transformer（ViT）和混合模型的表示学习能力。为了了解每种模型的数据需求，我们在不同规模的数据集上进行预训练，并评估了许多基准任务。在考虑预训练模型的计算成本时，ViT的表现非常出色，在大多数识别基准上以较低的预训练成本达到了最先进的水平。最后，我们进行了一个小型的自监督实验，结果表明自监督的ViT在未来具有很大的潜力。

4.1设置

数据集
为了探索模型的可扩展性，我们使用ILSVRC-2012 ImageNet数据集（包含1000类和130万张图片，我们在下文中称之为ImageNet），其超集ImageNet-21k（包含21000类和1400万张图片，参见Deng等，2009），以及JFT（包含18000类和3.03亿张高分辨率图片，参见Sun等，2017）。我们根据Kolesnikov等（2020）的方法，对预训练数据集进行去重，以避免与下游任务的测试集重复。我们将这些数据集上训练的模型迁移到多个基准任务上：ImageNet原始验证标签和清理后的ReaL标签（参见Beyer等，2020），CIFAR-10/100（参见Krizhevsky，2009），Oxford-IIIT Pets（参见Parkhi等，2012），以及Oxford Flowers-102（参见Nilsback & Zisserman，2008）。对于这些数据集的预处理，我们遵循Kolesnikov等（2020）的方法。

我们还在VTAB分类套件的19个任务上进行评估（参见Zhai等，2019b）。VTAB评估在多样化任务上的低数据迁移，每个任务使用1000个训练样本。任务分为三组：自然类任务——如前述的Pets、CIFAR等；专业类任务——如医学和卫星图像；结构类任务——需要几何理解的任务，如定位。
模型变体

我们基于BERT（参见Devlin等，2019）的配置来设定ViT的配置，如表1所总结。“Base”和“Large”模型直接采用BERT的配置，并添加了更大的“Huge”模型。在下文中，我们使用简要符号来表示模型大小和输入补丁大小：例如，ViT-L/16表示“Large”变体，输入补丁大小为16×16。请注意，Transformer的序列长度与补丁大小的平方成反比，因此，补丁大小较小的模型计算量更大。

对于基线CNN模型，我们使用ResNet（参见He等，2016），但将批量归一化层（Batch Normalization，参见Ioffe & Szegedy，2015）替换为组归一化（Group Normalization，参见Wu & He，2018），并使用标准化卷积（standardized convolutions，参见Qiao等，2019）。这些修改提高了迁移性能（参见Kolesnikov等，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比SGD略好）。我们使用线性学习率预热和衰减，详情见附录B.1。对于微调，我们对所有模型使用带动量的SGD，批量大小为512，详见附录B.1.1。对于表2中的ImageNet结果，我们在更高的分辨率下进行微调：ViT-L/16为512，ViT-H/14为518，并使用Polyak & Juditsky（1992）平均法，因子为0.9999（参见Ramachandran等，2019；Wang等，2020b）。

指标我们通过少样本学习或微调准确率报告下游数据集的结果。微调准确率反映了每个模型在相应数据集上进行微调后的性能。少样本学习准确率是通过解决一个正则化的最小二乘回归问题获得的，该问题将训练图像子集的（冻结的）表示映射到{−1, 1}(^K)目标向量。这个公式允许我们以闭式形式恢复精确解。尽管我们主要关注微调性能，但在微调成本过高的情况下，我们有时会使用线性少样本学习准确率进行快速即时评估。

4.2与最先进的结果的比较

我们首先将我们最大的模型——ViT-H/14 和 ViT-L/16——与文献中的最先进的 CNN 进行比较。第一个比较点是 Big Transfer (BiT) (Kolesnikov et al., 2020)，它使用大型 ResNet 执行有监督的迁移学习。第二个是 Noisy Student (Xie et al., 2020)，这是一个使用半监督学习在 ImageNet 和 JFT-300M 上训练的大型 EfficientNet，且标签已被移除。目前，Noisy Student 在 ImageNet 上是最先进的，而 BiT-L 则在其他数据集的表现也同样列出。所有模型均在 TPUv3 硬件上进行训练，我们报告了每个模型预训练所需的 TPUv3 核心天数，即用于训练的 TPU v3 核心（每个芯片 2 个）数量乘以训练时间（天）。

如表2所示。较小的 ViT-L/16 模型在 JFT-300M 上预训练后，在所有任务上都优于 BiT-L（它也在相同的数据集上预训练），同时训练所需的计算资源明显更少。较大的模型 ViT-H/14 进一步提高了性能，特别是在更具挑战性的数据集——ImageNet、CIFAR-100 和 VTAB 套件上。值得注意的是，这个模型的预训练计算量仍然显著少于之前的最先进模型。然而，我们注意到，预训练效率可能不仅受到架构选择的影响，还可能受到其他参数的影响，例如训练计划、优化器、权重衰减等。我们在第4.4节中提供了对不同架构的性能与计算量的对照研究。最终，在公共的 ImageNet-21k 数据集上预训练的 ViT-L/16 模型在大多数数据集上也表现良好，同时预训练所需的资源更少：它可以在大约30天内使用一个标准的8核云 TPUv3 进行训练。

图2将VTAB任务分解为各自的组，并与该基准上的之前最先进方法进行比较：BiT、VIVI——一个在 ImageNet 和 YouTube 上共同训练的 ResNet（Tschannen 等，2020），以及 S4L——在 ImageNet 上进行监督加半监督学习（Zhai 等，2019a）。在自然和结构化任务上，ViT-H/14 的表现优于 BiT-R152x4 和其他方法。在专业化任务上，前两名模型的表现相似。（图2中似乎没有标清楚这几个图分别属于哪个任务？）

4.3预训练数据需求

当在大型的JFT-300M数据集上进行预训练后，视觉Transformer表现良好。比ResNets对视觉的归纳偏差更少，数据集的大小有多重要？我们进行了两组实验。

首先，我们在逐渐增大数据大小的数据集上预训练Vit模型：ImageNet，ImageNet-21K，和JFT-300M。为了提高模型在小数据集上的表现，我们优化了三个基本的正则化参数-权重衰减，dropout和标签平滑。图三显示了在ImageNet上微调之后的结果（其他数据集上的结果在表5中有显示）。当在最小的数据集ImageNet上预训练时，Vit-Large模型相比Vit-Base表现不佳，尽管进行了中等的正则化。当使用ImageNet-21K预训练的时候，他们的表现接近。只有使用JFT-300M的时候，我们才能看到更大模型的全部好处。图3还显示了不同大小的BiT模型所跨越的性能区域。BiT CNNs在ImageNet上的表现优于ViT，但在更大的数据集上，ViT表现更好。

其次，我们在9M、30M和90M的随机子集以及完整的JFT- 300M数据集上训练我们的模型。我们不对较小的子集执行额外的正则化，而是对所有设置使用相同的超参数。这样，我们就可以评估模型的内在性质，而不是正则化的影响。然而，我们确实使用了早停，并报告了在训练期间获得的最佳验证准确性-。为了节省计算量，我们报告了少镜头的线性精度，而不是全微调精度。图4包含了这些结果。在较小的数据集上，Vision Transformers比ResNets过拟合更多。例如，ViT-B/32比ResNet50略快；它在9M子集上表现得更差，但在90M+子集上表现得更好。ResNet152x2和ViT-L/16也是如此。这一结果强化了一种直觉，即卷积归纳偏差对于较小的数据集是有用的，但对于较大的数据集，直接从数据中学习相关模式就足够了，甚至是有益的。

总的来说，ImageNet上的few-shot结果（图4），以及VTAB上的低数据结果（表2）似乎很适合非常低的数据传输。进一步分析ViT的few-shot特性是未来工作的一个令人兴奋的方向。

4.4缩放研究

通过评估JFT-300M的传输性能，我们对不同的模型进行了控制尺度研究。在这种情况下，数据规模并不会成为模型性能的瓶颈，因此我们评估每个模型的性能与预训练成本的关系。模型包括：7个ResNets，R50x1，R50x2，R101x1，R152x1，R152x2，预训练7个epoch，加上 R152x2和R200x3 预训练14个epoch；6个视觉Transformers，ViT-B/32，B/16，L/32，L/16，预训练7轮，加上L/16和H/14预训练了14epoch；还有5个混合模型，R50+ViT-B/32，B/16，L/32，L/16预训练了7轮，加上R50+ViT-L/16预训练了14个epoch（对于混合模型，模型名字后面的数字不是代表patch size，而是ResNet主干中的总的下采样率）。
图5包含了迁移性能和总的预训练计算量的对比。每个模型的详细的结果在附录中的表6中有展示。

可以观察到几个模式。首先，在性能/计算量的权衡上，Vision Transformers (ViT) 优于 ResNets。ViT 使用大约 2 到 4 倍更少的计算量即可达到相同的性能（在5个数据集上的平均值）。其次，在小计算预算下，混合模型略微优于 ViT，但这种差异在较大模型上消失。这一结果有些出人意料，因为人们可能会认为卷积局部特征处理在任何规模下都会对 ViT 有所帮助。第三，Vision Transformers 在尝试的范围内似乎没有达到饱和，激励了未来的扩展研究。

4.5 审视视觉Transformer

为了理解Vision Transformer如何处理图像数据，我们分析了它的内部表征。Vision Transformer的第一层线性地投射了扁平化后的补丁到一个低维空间(Eq 1)。图7（左）显示了学习到的嵌入过滤器的顶部主成分。这些组件类似于在每个小块内进行细结构低维表示的合理基函数。
投影之后，一个学习到的位置嵌入被添加到小块表示中。图7（中间）显示了模型学习将图像内的距离编码到位置嵌入的相似性中，即距离较近的小块往往具有更相似的位置嵌入。此外，行列结构也显现出来；同一行/列中的小块具有相似的嵌入。最后，对于较大的网格，有时会出现正弦结构（附录D）。位置嵌入学习表示二维图像拓扑，这解释了为什么手工制作的二维感知嵌入变体不会带来改进（附录D.4）。

自注意力允许ViT即使在最低的层级上也可以融合整张图片的信息。我们将研究该网络在多大程度上利用了这种能力。具体来说，我们根据注意力权重计算图像空间中信息被整合的平均距离（图7，右）。这种注意距离类比于CNNs中的接受野。我们发现，有些注意头在最底层已经能够关注到图像的大部分，说明模型确实在利用整合全局信息的能力。其他注意头在低层中始终具有较小的注意距离。这种高度局部化的注意力在混合模型中不太明显，这些混合模型在Transformer之前应用了ResNet（图7，右），这表明它可能具有类似于CNN中早期卷积层的功能。此外，注意距离随着网络深度的增加而增加。总体上，我们发现模型会关注到对分类具有语义相关性的图像区域（图6）。

4.6 自监督

Transformers在NLP任务上的表现出了令人印象深刻的表现。然而，他们的成功不仅来自于他们优秀的可伸缩性，而且还来自于大规模的自监督预训练。 我们还对自监督的遮掩补丁预测进行了初步探索，模仿了BERT中使用的遮掩语言建模任务。通过自监督预训练，我们较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率，相对于从零开始训练显著提升了2%，但仍然比监督预训练低4%。附录B.1.2包含更多详细信息。对比预训练的探索留待未来进行。

5.结论

我们探索了将Transformers直接应用于图像识别。与之前在计算机视觉中使用自注意力的工作不同，我们除了初始的图块提取步骤外，并未在架构中引入特定于图像的归纳偏差。相反，我们将图像解释为一系列图块，并通过标准的Transformer编码器处理，就像在NLP中使用的一样。这种简单但可扩展的策略在与大数据集上的预训练结合时表现出惊人的效果。因此，视觉Transformer在许多图像分类数据集上达到了或超过了最先进的水平，同时预训练成本相对较低。
虽然这些初步结果令人鼓舞，但仍有许多挑战。一个挑战是将ViT应用于其他计算机视觉任务，如检测和分割。我们的结果以及Carion等人（2020年）的研究表明了这种方法的前景。另一个挑战是继续探索自监督预训练方法。我们的初步实验显示，自监督预训练有一定的改进，但自监督与大规模监督预训练之间仍有很大差距。最后，进一步扩展ViT的规模可能会带来性能的提升。