论文笔记《AN IMAGE IS WORTH 16X16 WORDS：TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》

Karl2024-04-152024-04-20

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

原文链接：论文笔记《AN IMAGE IS WORTH 16X16 WORDS：TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》 | Karl的博客

CSDN链接：论文笔记《AN IMAGE IS WORTH 16X16 WORDS：TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》-CSDN博客

论文链接：[2010.11929] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (arxiv.org)

代码链接：google-research/vision_transformer (github.com)

ABSTRACT

虽然Transformer架构已经成为自然语言处理任务的事实上的标准，但它在计算机视觉上的应用仍然有限。在视觉方面，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组件，同时保持其整体结构不变。我们证明这种对CNNs的依赖是不必要的，直接应用于图像patch序列的纯Transformer可以很好地完成图像分类任务。当对大量数据进行预训练并转移到多个中型或小型图像识别基准（ImageNet，CIFAR-100，VTAB等）时，Vision Transformer（ViT）与最先进的卷积网络相比获得了出色的结果，同时需要更少的计算资源来训练。

Fine-tuning code and pre-trained models are available at https://github.com/google-research/vision_transformer.

1 INTRODUCTION

基于自注意力的架构，特别是Transformers（Vaswani et al., 2017），已经成为自然语言处理（NLP）的首选模型。主要的方法是在大型文本语料库上进行预训练，然后在较小的特定任务数据集上进行微调（Devlin et al., 2019）。由于Transformers的计算效率和可扩展性，它可以训练前所未有的规模，超过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应用程序中，图像patch的处理方式与令牌（单词）相同。我们以监督的方式对模型进行图像分类训练。

当在中等规模的数据集（如ImageNet）上进行训练时，没有进行强正则化，这些模型产生的精度比同等规模的ResNets低几个百分点。这个看似令人沮丧的结果是可以预料的：Transformers缺乏CNNs固有的一些归纳偏差，例如平移不变性和局部性，因此在数据量不足的情况下训练时不能很好地泛化。

然而，如果模型在更大的数据集（14M-300M图像）上训练，情况就会发生变化。我们发现大规模训练胜过归纳偏见。我们的Vision Transformer（ViT）在足够的规模上进行预训练并转移到具有更少数据点的任务时获得了出色的结果。当在公共ImageNet-21k数据集或室内JFT-300M数据集上进行预训练时，ViT在多个图像识别基准上接近或超过了最先进的水平。特别是，最佳模型在ImageNet上达到88.55%的准确率，在ImageNet- real上达到90.72%的准确率，在CIFAR-100上达到94.55%的准确率，在VTAB套件的19个任务上达到77.63%的准确率。

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

对图像进行注意力的朴素应用会要求每个像素都注意其他像素。由于像素数量的代价是二次的，因此不能按实际的输入大小进行缩放。因此，为了在图像处理的背景下应用Transformers，过去已经尝试了几种近似方法。Parmar等人（2018）仅在局部邻域对每个查询像素应用自注意力，而不是全局。这种局部多头点积自注意块可以完全取代卷积（Hu et al., 2019; Ramachandran et al., 2019; Zhao et a., 2020）。在不同的工作中，稀疏Transformers（Child et al., 2019）采用可扩展的全局自注意力近似，以便适用于图像。另一种扩展注意力的方法是将其应用于不同大小的块（Weissenborn et al., 2019），在极端情况下，仅沿着单个轴（Ho et al., 2019; Wang et al., 2020a）。许多这些专门的注意力架构在计算机视觉任务上显示出有希望的结果，但需要复杂的工程才能在硬件加速器上有效地实现。

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

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

另一个最近的相关模型是图像GPT（iGPT）（Chen et al., 2020a），它在降低图像分辨率和色彩空间后，将Transformers应用于图像像素。该模型以无监督的方式作为生成模型进行训练，然后可以对结果表示进行微调或线性探测以获得分类性能，在ImageNet上实现72%的最大准确率。

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

3 METHOD

在模型设计中，我们尽可能地遵循原始Transformer（Vaswani et al., 2017）。这种故意简单设置的一个优点是，可扩展的NLP Transformer架构——以及它们的高效实现——几乎可以开箱即用。

3.1 VISION TRANSFORMER (VIT)

该模型的概述如图1所示：

标准Transformer接收标记嵌入的1D序列作为输入。为了处理2D图像，我们将图像$\mathbf{x} \in \mathbb{R}^{H \times W \times C}$重塑为一系列平坦的2D patch $\mathbf{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$，其中$(H, W)$为原始图像的分辨率，$C$为通道数，$(P, P)$为每个图像patch的分辨率，$N = \frac{HW}{P^2}$为得到的patch数，它也作为Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的潜在向量大小$D$，因此我们用可训练的线性投影：

$\mathbf{z}_0 = \left[\mathbf{x}_\text{class};\ \mathbf{x}_p^1\mathbf{E};\ \mathbf{x}_p^2\mathbf{E};\ \cdots;\ \mathbf{x}_p^N\mathbf{E}\right] + \mathbf{E}_{pos}, \quad \mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D},\ \mathbf{E}_{pos} \in \mathbb{R}^{(N + 1) \times D}$

将patch展平并映射到$D$维。我们将该投影的输出称为patch embedding。

与BERT的[class]令牌类似，我们在嵌入的patch序列（$\mathbf{z}_0^0 = \mathbf{x}_\text{class}$）上添加一个可学习的嵌入，其在Transformer编码器（$\mathbf{z}_L^0$）输出处的状态作为图像表示$\mathbf{y}$：

$\mathbf{y} = \text{LN}(\mathbf{z}_L^0)$

在预训练和微调过程中，$\mathbf{z}_L^0$都附加了一个分类头。分类头在预训练时由一个隐藏层的MLP实现，在微调时由一个线性层实现。

将位置嵌入添加到patch嵌入中以保留位置信息。我们使用标准的可学习的1D位置嵌入，因为我们没有观察到使用更先进的2D感知位置嵌入的显著性能提升（附录D.4）。得到的嵌入向量序列作为编码器的输入。

Transformer编码器（Vaswani et al., 2017）由多头自注意（MSA，见附录A）和MLP块交替层组成：

$\begin{aligned} \mathbf{z}^\prime_\ell &= \text{MSA}(\text{LN}(\mathbf{z}_{\ell - 1})) + \mathbf{z}_{\ell - 1} & & \ell = 1 \dots L \\ \mathbf{z}_\ell &= \text{MSA}(\text{LN}(\mathbf{z}^\prime_\ell)) + \mathbf{z}^\prime_\ell & & \ell = 1 \dots L \end{aligned}$

在每个区块之前应用层Layernorm（LN），在每个区块之后应用残差连接（Wang et al., 2019; Baevski & Auli, 2019）。

MLP包含两个具有GELU非线性的层：

$\begin{aligned} \mathbf{z}_0 &= \left[\mathbf{x}_\text{class};\ \mathbf{x}_p^1\mathbf{E};\ \mathbf{x}_p^2\mathbf{E};\cdots\ ;\ \mathbf{x}_p^N\mathbf{E}\right] + \mathbf{E}_{pos}, & & \mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}, \mathbf{E}_{pos} \in \mathbb{R}^{(N + 1) \times D} \\ \mathbf{z}^\prime_\ell &= \text{MSA}(\text{LN}(\mathbf{z}_{\ell - 1})) + \mathbf{z}_{\ell - 1}, & & \ell = 1 \dots L \\ \mathbf{z}_\ell &= \text{MLP}(\text{LN}(\mathbf{z}^\prime_\ell)) + \mathbf{z}^\prime_\ell, & & \ell = 1 \dots L \\ \mathbf{y} &= \text{LN}(\mathbf{z}^0_L) \end{aligned}$

Inductive bias. 我们注意到Vision Transformer具有比CNNs少得多的图像特定的感应偏置。在CNNs中，局部性、二维邻域结构和平移等方差被嵌入到整个模型的每一层中。在ViT中，只有MLP层是局部和平移不变的，而自注意力层是全局的。二维邻域结构的使用非常少：在模型开始时通过将图像切割成小块，以及在微调时用于调整不同分辨率图像的位置嵌入（如下所述）。除此之外，初始化时的位置嵌入不携带关于patch的二维位置信息，所有patch之间的空间关系都需要从头学习。

Hybrid Architecture. 作为原始图像patch的替代方案，输入序列可以由CNN的特征映射形成（LeCun et al., 1989）。在该混合模型中，将patch embedding投影$\mathbf{E}$应用于从CNN feature map中提取的patch。作为一种特殊情况，patch的空间大小可以是$1 \times 1$，这意味着输入序列是通过简单地将feature map的空间维度平坦化并投影到Transformer维度得到的。如上所述添加分类输入嵌入和位置嵌入。

3.2 FINE-TUNING AND HIGHER RESOLUTION

通常，我们在大型数据集上预训练ViT，并对（较小的）下游任务进行微调。为此，我们移除预训练的预测头，并附加一个零初始化的$D \times K$前传层，其中$K$是下游类的数量。与预训练相比，在更高分辨率下进行微调通常是有益的（Touvron et al., 2019; Kolesnikov et al., 2020）。当输入更高分辨率的图像时，我们保持patch大小不变，从而获得更大的有效序列长度。Vision Transformer可以处理任意序列长度（直到内存限制），但是，预训练的位置嵌入可能不再有意义。因此，我们根据预训练的位置嵌入在原始图像中的位置对其进行二维插值。请注意，此分辨率调整和patch提取是将图像的二维结构的感应偏置手动注入Vision Transformer的唯一点。

4 EXPERIMENTS

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

4.1 SETUP

Datasets. 为了探索模型的可扩展性，我们使用了ILSVRC-2012 ImageNet数据集，其中包含1k个类和130万张图像（我们将其称为ImageNet），其超集ImageNet-21k包含21k个类和14M张图像（Deng et al., 2009），以及JFT（Sun et al., 2017）包含18k个类和303M张高分辨率图像。我们在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个训练示例。任务分为三组：自然任务，如上述任务，宠物，CIFAR等；专业任务——医学和卫星图像；以及结构任务——需要几何理解的任务，比如定位。

Model Variants. 我们基于BERT使用的ViT配置（Devlin et al., 2019），如表1所示：

“Base”和“Large”模型直接取自BERT，我们加入了较大的“Huge”模型。在接下来的内容中，我们使用简短的符号来表示模型大小和输入patch大小：例如，ViT-L/16表示具有$16 \times 16$输入patch大小的“Large”变体。请注意，Transformer的序列长度与patch大小的平方成反比，因此patch大小较小的模型在计算上更昂贵。

对于基线CNNs，我们使用ResNet（He et al., 2016），但将批归一化层（Ioffe & Szegedy, 2015）替换为组归一化层（Wu & He, 2018），并使用标准化卷积（Qiao et al., 2019）。这些修改改善了迁移（Kolesnikov et al., 2020），我们将修改后的模型命名为“ResNet (BiT)”。对于混合架构，我们将中间特征映射以一个“像素”的patch大小馈送到ViT中。为了实验不同的序列长度，我们：

要么获取普通ResNet50的阶段4的输出；
要么删除阶段4，在阶段3中放置相同数量的层（保持层总数），并获取这个扩展阶段3的输出。

选项2的结果是4倍长的序列长度和更昂贵的ViT模型。

Training & Fine-tuning. 我们使用Adam（Kingma & Ba, 2015）训练所有模型，包括ResNets，其中$\beta_1 = 0.9, \beta_2 = 0.999$，批大小为4096，并应用0.1的高权重衰减，我们发现这对所有模型的转移都很有用（附录D.1显示，与常规做法相比，Adam在我们的设置中比SGD对ResNets的效果略好）。我们使用线性学习率预热和衰减，详见附录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）。

Metrics. 我们通过小样本精度或微调精度报告下游数据集的结果。微调精度捕获每个模型在各自的数据集上进行微调后的性能。通过求解正则化最小二乘回归问题，将训练图像子集的（冻结）表示映射到$\{-1, 1\}^K$个目标向量，获得小样本精度。这个公式使我们可以得到封闭形式的精确解。虽然我们主要关注微调性能，但我们有时会使用线性的小样本精度来进行快速的动态评估，因为微调的成本太高。

4.2 COMPARISON TO STATE OF THE ART

我们首先将我们最大的模型——ViT-H/14和ViT-L/16，与文献中最先进的CNNs进行比较。第一个比较点是Big Transfer（BiT）（Kolesnikov et al., 2020），它使用大型ResNets执行监督迁移学习。第二个是Noisy Student（Xie et al., 2020），它是一个大型的高效网络，使用半监督学习在ImageNet和JFT300M上进行训练，去掉了标签。目前，Noisy Student在ImageNet上是最先进的，BiT-L在这里报道的其他数据集上是最先进的。所有模型都在TPUv3硬件上进行训练，我们报告预训练每个模型所需的TPUv3-core-days，即用于训练的TPUv3核数（每个芯片2个）乘以训练时间（天）。

表2显示了结果：

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

图2将VTAB任务分解为各自的组：

并在此基准上与之前的SOTA方法进行比较：BiT，VIVI在ImageNet和Youtube上共同训练的ResNet（Tschannen et al., 2020），以及S4L在ImageNet上监督和半监督学习（Zhai et al., 2019a）。ViT-H/14在自然任务和结构任务上优于BiT-R152x4和其他方法。在专业任务上，前两个型号的性能相似。

4.3 PRE-TRAINING DATA REQUIREMENTS

Vision Transformer在大型JFT-300M数据集上进行预训练时表现良好。与ResNets相比，对于视觉的归纳偏差更少，那么数据集的大小有多重要呢？我们进行了两个系列的实验。

首先，我们在越来越大的数据集上预训练ViT模型：ImageNet、ImageNet-21k和JFT300M。为了提高在较小数据集上的性能，我们优化了三个基本的正则化参数——权重衰减、dropout和标签平滑。图3显示了调优到ImageNet后的结果（表5显示了其他数据集上的结果）：

注意，ImageNet预训练的模型也进行了微调，但还是在ImageNet上。这是因为在微调过程中分辨率的提高提高了性能。

当在最小的数据集ImageNet上进行预训练时，尽管（适度）正则化，ViT-Large模型的表现仍不如ViT-Base模型。使用ImageNet-21k预训练，它们的性能是相似的。只有使用JFT-300M，我们才能看到更大型号的全部好处。图3还显示了不同大小的BiT模型所跨越的性能区域：

BiT CNNs在ImageNet上的表现优于ViT，但在更大的数据集上，ViT超越了它。

其次，我们在9M、30M和90M的随机子集以及完整的JFT300M数据集上训练我们的模型。我们没有对较小的子集执行额外的正则化，并对所有设置使用相同的超参数。这样，我们评估的是模型的内在属性，而不是正则化的影响。然而，我们确实使用了早期停止，并报告了在训练期间实现的最佳验证准确性。为了节省计算，我们报告了小样本线性精度而不是全微调精度。图4包含了结果：

在较小的数据集上，Vision Transformers比ResNets过拟合得多，计算成本相当。例如，ViT-B/32略快于ResNet50；它在9M的子集上表现得更差，但在90M以上的子集上表现得更好。ResNet152x2和ViT-L/16也是如此。这个结果强化了卷积归纳偏差对较小数据集有用的直觉，但对于较大的数据集，直接从数据中学习相关模式是足够的，甚至是有益的。

总的来说，ImageNet上的小样本结果（图4）：

以及VTAB上的小样本结果（表2）：

似乎表明小样本迁移学习很有希望。进一步分析ViT的小样本特性是未来工作的一个令人兴奋的方向。

4.4 SCALING STUDY

我们通过评估JFT-300M的迁移学习性能，对不同模型进行了控制缩放研究。在这种情况下，数据大小不会成为模型性能的瓶颈，我们根据每个模型的预训练成本来评估性能。模型集包括：

7个ResNets：R50x1、R50x2、R101x1、R152x1、R152x2，预训练7个epoch，加上R152x2和R200x3预训练14个epoch；
6个Vision Transformers，ViT-B/32、B/16、L/32、L/16，预训练7个epoch，加上L/16和H/14预训练14个epoch；
和5个混合架构，R50+ViT-B/32、B/16、L/32、L/16 预训练7个epoch，加上R50+ViT-L/16预训练14个epoch（对于混合架构，数量为模型名称的末尾不代表patch大小，而是代表ResNet主干中的总下采样率）。

图5包含了迁移学习性能与总预训练计算的对比（参见附录D.5有关计算成本的详细信息）：

每个模型的详细结果见附录表6：

可以观察到一些模式。首先，Vision Transformers在性能/计算权衡方面比ResNets有优势。ViT节省了大约2-4倍的计算来获得相同的性能（平均超过5个数据集）。其次，混合架构在较小的计算预算下略优于ViT，但对于较大的模型，这种差异就消失了。这个结果有些令人惊讶，因为人们可能期望卷积局部特征处理能够帮助任何规模的ViT。第三，Vision Transformers在尝试的范围内似乎没有饱和，这激励了未来的扩展努力。

4.5 INSPECTING VISION TRANSFORMER

为了开始理解Vision Transformer如何处理图像数据，我们分析其内部表示。Vision Transformer的第一层将展平后的patch线性投影到较低维空间：

图7（左）显示了学习到的嵌入过滤器的顶部主成分：

这些成分类似于每个patch内精细结构的低维表示的可信基函数。

投影后，将学习到的位置嵌入添加到patch表示中。图7（中）显示，该模型在位置嵌入的相似性中学习对图像内的距离进行编码，即距离越近的patch往往有更多相似的位置嵌入。此外，出现了行-列结构：同一行/列中的patch具有相似的嵌入。最后，对于较大的网格，有时会出现正弦结构（附录D）。位置嵌入学习表示二维图像拓扑解释了为什么手工制作的二维感知嵌入变体没有产生改进（附录D.4）。

自注意力允许ViT在整个图像中集成信息，即使是在最低层。我们研究了网络在多大程度上利用了这种能力。具体来说，我们计算图像空间中信息被整合的平均距离，这是基于注意力权重的（图7，右）。这种“注意距离”类似于CNNs的感受野大小。我们发现，一些人的大脑已经注意到了最低层的大部分图像，这表明该模型确实使用了整体整合信息的能力。其他注意头在低层的注意距离一直很小。在Transformer之前应用ResNet的混合模型中，这种高度局部化的注意力不那么明显（图7，右），这表明它可能与CNN中的早期卷积层具有类似的功能。注意距离随网络深度的增加而增加。从全局来看，我们发现该模型关注与分类在语义上相关的图像区域（图6）：

4.6 SELF-SUPERVISION

Transformers在NLP任务上表现出色。然而，它们的成功不仅源于出色的可扩展性，还源于大规模的自监督预训练（Devlin et al., 2019; Radford et al., 2018）。我们还模拟BERT中使用的掩码语言建模任务，对自我监督的掩码patch预测进行了初步探索。通过自我监督预训练，我们较小的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

我们探索了Transformers在图像识别中的直接应用。与之前在计算机视觉中使用自注意的工作不同，除了初始patch提取步骤外，我们没有将特定于图像的归纳偏差引入体系结构。相反，我们将图像解释为一系列patch，并通过NLP中使用的标准Transformer编码器对其进行处理。当与大型数据集的预训练相结合时，这种简单但可扩展的策略效果出奇地好。因此，Vision Transformer在许多图像分类数据集上达到或超过了最先进的水平，同时预训练相对便宜。

虽然这些初步结果令人鼓舞，但仍存在许多挑战。一是将ViT应用于其他计算机视觉任务，如检测和分割。我们的结果，加上Carion等人（2020）的结果，表明了这种方法的前景。另一个挑战是继续探索自我监督的预训练方法。我们的初步实验显示了自监督预训练的改进，但自监督预训练与大规模监督预训练之间仍有很大差距。最后，进一步扩展ViT可能会提高性能。

ACKNOWLEDGEMENTS

这项工作在柏林、苏黎世和阿姆斯特丹执行。我们感谢谷歌的许多同事的帮助，特别是Andreas Steiner，他在基础设施和代码的开源发布方面提供了至关重要的帮助；Joan Puigcerver和Maxim Neumann提供大规模训练基础设施的帮助；Dmitry Lepikhin，Aravindh Mahendran，Daniel Keysers，Mario Lucic，Noam Shazeer，Ashish Vaswani和Colin Raffel进行了有益的讨论。

APPENDIX

A MULTIHEAD SELF-ATTENTION

标准$\mathbf{qkv}$自注意（SA, Vaswani et al.(2017)）是神经架构的流行构建块。对于输入序列$\mathbf{z} \in \mathbb{R}^{N \times D}$，中的每个元素，我们计算序列中所有值$\mathbf{v}$的加权和。注意力权重$A_{ij}$基于序列中两个元素的成对相似度及其各自的查询$\mathbf{q}^i$和键$\mathbf{k}^j$表示。

$\begin{aligned} \left[\mathbf{q}, \mathbf{k}, \mathbf{v}\right] &= \mathbf{z}\mathbf{U}_{qkv} & & \mathbf{U}_{qkv} \in \mathbb{R}^{D \times 3D_h}, \\ A &= \text{softmax}(\frac{\mathbf{qk}^\mathrm{T}}{\sqrt{D_h}}) & & A \in \mathbb{R}^{N \times N}, \\ \text{SA}(\mathbf{z}) &= A\mathbf{v}. \end{aligned}$

多头自注意（MSA）是自注意的扩展，其中我们并行运行$k$个自注意操作，称为“头”，并投影它们的连接输出。为了在改变$k$时保持计算量和参数数量不变，通常将$D_h$设为$\frac{D}{k}$。

$\text{MSA}(\mathbf{z}) = \left[\text{SA}_1(z); \text{SA}_2(z);\cdots\ ; \text{SA}_k(z)\right]\mathbf{U}_{msa} \quad \mathbf{U}_{msa} \in \mathbb{R}^{k \cdot D_k \times D}$

B EXPERIMENT DETAILS

B.1 TRAINING

表3总结了不同模型的训练设置：

我们发现在ImageNet上从头开始训练模型时，强正则化是关键。使用时，Dropout应用于除qkv投影之外的每个密集层之后，以及直接在添加位置到patch嵌入之后应用。混合模型使用与ViT模型完全相同的设置进行训练。最后，所有的训练都是在分辨率224上完成的。

B.1.1 FINE-TUNING

我们使用动量为0.9的SGD对所有ViT模型进行微调。我们在学习率上运行一个小的网格搜索，参见表4中的学习率范围。为此，我们使用训练集中的小分支（宠物和鲜花为10%，CIFAR为2%，ImageNet为1%）作为开发集，并在剩余数据上进行训练。对于最终结果，我们对整个训练集进行训练，并对各自的测试数据进行评估。对于微调ResNets和混合模型，我们使用完全相同的设置，唯一的例外是ImageNet，我们在学习率扫描中添加另一个值0.06。此外，对于ResNets，我们还运行Kolesnikov等人（2020）的设置，并在本次运行和扫描中选择最佳结果。最后，如果没有特别提到，所有的微调实验都在384分辨率下运行（在不同于训练的分辨率下运行微调是常见的做法（Kolesnikov et al., 2020））。

当将ViT模型转移到另一个数据集时，我们删除整个头部（两个线性层）并将其替换为单个，零初始化的线性层，输出目标数据集所需的类数。我们发现这比简单地重新初始化最后一层要健壮一些。

对于VTAB，我们遵循Kolesnikov等人（2020）的协议，并对所有任务使用相同的超参数设置。我们使用学习率为0.01，训练2500步（表4）：

我们通过在两个学习率和两个时间表上运行一个小扫描来选择这个设置，并在200个示例验证集中选择VTAB得分最高的设置。我们遵循Kolesnikov等人（2020）使用的预处理，除了我们不使用特定于任务的输入分辨率。相反，我们发现Vision Transformer受益于所有任务的高分辨率（$384 \times 384$）。

B.1.2 SELF-SUPERVISION

我们采用掩模patch预测目标进行初步的自我监督实验。为了做到这一点，我们通过用可学习的[mask]嵌入（80%），随机的其他patch嵌入（10%）或保持原样（10%）替换它们的嵌入来破坏50%的patch嵌入。这种设置与Devlin等人（2019）使用的语言非常相似。最后，我们使用各自的patch表示来预测每个损坏patch的3位平均颜色（即总共512种颜色）。

我们在JFT上训练了1M步（约14个epoch）的自监督模型，批大小为4096。我们使用Adam，其基本学习率为$2 \cdot 10^{−4}$，预热为10k步，余弦学习率衰减。作为预训练的预测目标，我们尝试了以下设置：

仅预测平均的3位颜色（即512种颜色的1种预测）；
预测$16 \times 16$ patch的$4 \times 4$缩小版本，并行使用3位颜色（即512种颜色的16种预测）；
使用L2在完整patch上进行回归（即在3个RGB通道上进行256种回归）。

令人惊讶的是，我们发现所有这些都运行得很好，尽管L2稍微差一些。我们只报告选项1的最终结果，因为它显示了最佳的少数镜头性能。我们还用Devlin等人（2019）使用的15%腐败率进行了实验，但在我们的少数几个指标上，结果也略差。

最后，我们想指出的是，我们的掩码patch预测实例化不需要如此大量的预训练，也不需要像JFT这样的大型数据集，以便在ImageNet分类上获得类似的性能提升。也就是说，我们观察到在10万步预训练后下游性能的收益递减，并且在ImageNet上预训练时看到类似的收益。

C ADDITIONAL RESULTS

我们报告了与论文中给出的数字相对应的详细结果。表5对应于本文的图3：

显示了不同ViT模型在ImageNet、ImageNet-21k和JFT-300M数据集上预训练的迁移训练性能。表6对应于本文的图5：

显示了不同大小的ViT、ResNet和混合模型的迁移学习性能，以及预训练的估计计算成本。

D ADDITIONAL ANALYSES

D.1 SGD VS. ADAM FOR RESNETS

ResNets通常使用SGD进行训练，我们使用Adam作为优化器是非常非常规的。在这里，我们展示了激发这一选择的实验。也就是说，我们比较了SGD和Adam在JFT上预训练的两个ResNets-50x1和152x2的微调性能。对于SGD，我们使用Kolesnikov等人（2020）推荐的超参数。结果如表7所示：

Adam预训练在大多数数据集和平均水平上优于SGD预训练。这证明了选择Adam作为在JFT上预训练ResNets的优化器是合理的。请注意，绝对数字低于Kolesnikov等人（2020）报告的数字，因为我们只预训练了7个epoch，而不是30个epoch。

D.2 TRANSFORMER SHAPE

我们对Transformer架构的不同维度进行了扩展，以找出最适合扩展到非常大的模型的维度。图8显示了不同配置下ImageNet上的5次性能测试：

所有配置都基于8层的ViT模型，$D = 1024, D_{MLP} = 2048$，patch大小为32，所有线相交。我们可以看到，缩放深度带来了最大的改进，这一点在64层之前非常明显。然而，在16层之后，收益递减已经很明显了。有趣的是，扩展网络的宽度似乎只会导致最小的变化。在不引入参数的情况下，减小patch大小从而增加有效序列长度显示出惊人的鲁棒性改进。这些发现表明，计算可能比参数数量更好地预测性能，并且缩放应该强调深度而不是宽度（如果有的话）。总体而言，我们发现按比例缩放所有维度会产生稳健的改进。

D.3 HEAD TYPE AND `CLASS` TOKEN

为了尽可能接近原始的Transformer模型，我们使用了一个额外的[class]令牌，它被用作图像表示。然后，该令牌的输出通过一个小型多层感知器（MLP）转换为类预测，其中$\tanh$为单个隐藏层中的非线性。

这种设计继承自文本的Transformer模型，我们在整个主要论文中都使用它。最初的尝试是只使用图像patch嵌入，全局平均池化（GAP），然后是线性分类器——就像ResNet的最终特征图一样——表现非常糟糕。然而，我们发现这既不是由于额外的令牌，也不是由于GAP操作。相反，性能上的差异完全可以通过对不同学习率的需求来解释，参见图9：

D.4 POSITIONAL EMBEDDING

我们使用位置嵌入对不同的空间信息编码方式进行了消融。我们尝试了以下情况：

不提供位置信息：将输入视为a bag of patchs；
一维位置嵌入：将输入视为栅格顺序的patch序列（本文中所有其他实验的默认值）；
二维位置嵌入：将输入视为二维的patch网格。在这种情况下，学习了两组嵌入，每组用于一个轴，$X$嵌入和$Y$嵌入，每个轴的大小为$\frac{D}{2}$；然后，基于输入路径上的坐标，我们将$X$嵌入和$Y$嵌入连接起来，得到该patch的最终位置嵌入。
相对位置嵌入：考虑patch之间的相对距离来编码空间信息，而不是它们的绝对位置。为此，我们使用一维相对注意，其中我们定义了所有可能的patch对的相对距离。因此，对于每个给定的对（一个作为查询，另一个作为注意机制中的键/值），我们有一个偏移量$p_q − p_k$，其中每个偏移量都与嵌入相关联。然后，在应用softmax之前，我们使用相对注意的logit作为偏差项，并将其添加到主要注意（基于内容的注意）的logit中。

除了不同的空间信息编码方式，我们还尝试了不同的方式将这些信息整合到我们的模型中。对于一维和二维位置嵌入，我们尝试了三种不同的情况：

在它们模型的主干之后和将输入馈给Transformer编码器之前向输入添加位置嵌入（本文中所有其他实验的默认值）；
学习并在每层开始的输入中添加位置嵌入；
在每层开始的输入中添加一个学习到的位置嵌入（层与层之间共享）。

表8总结了在ViT-B/16模型上的消融研究结果：

我们可以看到，虽然没有位置嵌入的模型和有位置嵌入的模型在性能上有很大的差距，但是不同的位置信息编码方式之间几乎没有差别。我们推测，由于我们的Transformer编码器在patch级输入上操作，而不是像素级输入，因此如何编码空间信息的差异不太重要。更准确地说，在patch级输入中，空间维度比原始像素级输入小得多，例如，$14 \times 14$而不是$224 \times 224$，并且对于这些不同的位置编码策略来说，学习在这种分辨率下表示空间关系同样容易。即便如此，网络学习到的位置嵌入相似度的具体模式取决于训练超参数（图10）：

D.5 EMPIRICAL COMPUTATIONAL COSTS

我们还对硬件上架构的实际速度感兴趣，由于通道宽度和缓存大小等细节，理论FLOPs并不总是很好地预测。为此，我们在TPUv3加速器上对感兴趣的主要模型执行推理速度计时；推理速度和反传速度之间的差异是一个恒定的模型无关因素。

图12（左）显示了在不同的输入大小下，一个核每秒可以处理多少图像：

每个单点都是指在广泛的批大小范围内测量的峰值性能。可以看到，ViT与图像大小的理论双二次缩放仅在最大分辨率下的最大模型中才刚刚开始发生。

另一个感兴趣的量是每个模型可以容纳在一个核心上的最大批处理大小，更大的批处理更适合扩展到大型数据集。图12（右）显示了同一组模型的数量。这表明大型ViT模型在内存效率方面比ResNet模型有明显的优势。

D.6 AXIAL ATTENTION

轴向注意力（Huang et al., 2020; Ho et al., 2019）是一种简单而有效的技术，可以在组织为多维张量的大输入上运行自注意力。轴向注意力的一般思想是执行多个注意力操作，每个操作沿着输入张量的单个轴进行，而不是将一维注意力应用于展平后的输入。在轴向注意力中，每个注意力沿特定轴混合信息，同时保持沿其他轴的信息独立。沿着这条路线，Wang等人（2020b）提出了AxialResNet模型，其中ResNet50中所有内核大小为$3 \times 3$的卷积都被轴向自注意力（即行和列注意力）取代，并通过相对位置编码增强。我们已经实现了AxialResNet作为一个基线模型。

我们的实现基于https://github.com/csrhddlam/axial-deeplab中的开源PyTorch实现。在我们的实验中，我们在准确性方面复制了（Wang et al., 2020b）中报告的分数，然而，我们的实现与开源实现类似，在TPUs上非常慢。因此，我们无法将其用于广泛的大规模实验。这些可以通过精心优化的实现来解锁。

此外，我们修改了ViT以处理二维形状的输入，而不是一维序列的patch，并合并了轴向Transformer块，其中不是自注意力后面跟着MLP，而是行自注意力加上MLP，然后是列自注意力加上MLP。

图13展示了在JFT数据集上进行预训练时，轴向ResNet、轴向ViT-B/32和轴向ViT-B/16在ImageNet 5shot线性上与预训练计算的性能，包括FLOPs数和推理时间（每秒示例数）。正如我们所看到的，就性能而言，轴向ViT-B/32和轴向ViT-B/16都比它们的ViT-B对应物做得更好，但这是以更多的计算为代价的。这是因为在轴向ViT模型中，每个具有全局自注意力的Transformer块被两个轴向Transformer块取代，一个具有行自注意力，一个具有列自注意力，尽管自注意力操作的序列长度在轴向情况下较小，但每个轴向ViT块都有一个额外的MLP。对于AxialResNet，尽管它在精度/计算权衡方面看起来是合理的（图13，左），但是这种朴素的实现在TPUs上非常慢（图13，右）：