ViT开山之作解读:An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

1 minute read

Published:

论文地址:An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

这篇论文由Google AI团队提出,是Vision Transformer(ViT)的开山之作,将Transformer模型成功应用于图像分类任务,证明了在大规模数据集上预训练的ViT可以超越传统的卷积神经网络(CNN)。

摘要

虽然Transformer架构已成为自然语言处理任务的实际标准,但其在计算机视觉领域的应用仍然有限。在视觉领域,注意力机制要么与卷积网络结合使用,要么用来替换卷积网络的某些组件,同时保持其整体结构不变。我们展示了这种对CNN的依赖是不必要的,直接应用于图像块序列的纯Transformer可以在图像分类任务上表现非常出色。当在大量数据上进行预训练并迁移到多个中等规模或小型图像识别基准测试(如ImageNet、CIFAR-100、VTAB等)时,视觉Transformer(ViT)与最先进的卷积网络相比,取得了卓越的结果,同时在训练过程中需要的计算资源要少得多。

1 introduction

基于自注意力的架构,特别是Transformer(Vaswani等人,2017年),已经成为自然语言处理(NLP)的首选模型。主流的方法是在大型文本语料库上进行预训练,然后在较小的任务特定数据集上进行微调(Devlin等人,2019年)。由于Transformer的计算效率和可扩展性,现在可以训练具有前所未有的规模的模型,参数超过100B(Brown等人,2020年;Lepikhin等人,2020年)。随着模型和数据集的增长,性能的增长还没有出现饱和的迹象。

然而,在计算机视觉领域,卷积架构仍然占据主导地位(LeCun等人,1989年;Krizhevsky等人,2012年;He等人,2016年)。受到NLP成功的启发,许多研究尝试将类似CNN的架构与自注意力结合起来(Wang等人,2018年;Carion等人,2020年),有些甚至完全取代了卷积(Ramachandran等人,2019年;Wang等人,2020年)。后者的模型虽然在理论上是高效的,但由于使用了专门的注意力模式,还没有在现代硬件加速器上有效地扩展。因此,在大规模图像识别中,经典的类似ResNet的架构仍然是最先进的(Mahajan等人,2018年;Xie等人,2020年;Kolesnikov等人,2020年)。

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

当在没有强正则化的情况下在中等规模的数据集(如ImageNet)上训练时,这些模型的准确率比同等规模的ResNet低几个百分点。这种看似令人沮丧的结果是可以预料的:Transformer缺乏CNN固有的一些归纳偏置,例如平移等变性和局部性,因此在数据量不足的情况下泛化能力不强。

然而,如果模型在更大的数据集(1400万-3亿图像)上进行训练,情况就会发生变化。我们发现,大规模训练胜过归纳偏置。我们的Vision Transformer(ViT)在足够规模的预训练后转移到数据点较少的任务上时,取得了优异的结果。当在公开的ImageNet-21k数据集或内部的JFT-300M数据集上进行预训练时,ViT在多个图像识别基准测试中接近或超越了最先进的水平。特别是,最佳模型在ImageNet上的准确率达到了88.55%,在ImageNet-ReaL上达到90.72%,在CIFAR-100上达到94.55%,在包含19个任务的VTAB套件上达到77.63%。


与卷积神经网络(CNN)相比,Transformer模型缺少一些内在的归纳偏置(inductive biases),这些归纳偏置是指模型结构中固有的、对数据的某些假设,这些假设能够帮助模型在有限数据的情况下也能较好地泛化。具体来说,本文提到的两个归纳偏置是:

  1. 平移等变性(Translation Equivariance):这是指模型对于输入数据的平移变化具有不变性。在图像处理中,这意味着无论图像中的对象出现在哪个位置,模型都应该能够识别出来。CNN通过使用滑动窗口的卷积操作来实现这一点,卷积核在图像上滑动时,能够捕捉到局部特征,并且这些特征的响应是相对于卷积核位置的。这种结构使得CNN在处理图像时具有平移等变性。

  2. 局部性(Locality):这是指模型倾向于关注输入数据中的局部特征。在图像中,这意味着相邻的像素之间往往有更强的相关性。CNN通过局部连接的卷积核来捕捉这种局部性,每个卷积核只关注它覆盖的局部区域的特征。

由于Transformer在结构上是基于全局注意力机制的,它在处理图像时会考虑图像的所有部分,而不像CNN那样有局部感受野。这意味着Transformer在处理图像时不会自动地捕捉到局部特征和位置信息,因此在数据量不足的情况下,它可能无法像CNN那样有效地泛化。

为了克服这些局限性,Transformer模型可能需要额外的技术,如位置编码(positional encoding)来提供序列中的位置信息,或者通过在大量数据上进行预训练来学习数据的分布特征,从而提高模型的泛化能力。在预训练阶段,Transformer可以通过在大规模数据集上学习,逐渐获得对数据的更深层次理解,这有助于在面对较小的数据集时,模型也能够较好地泛化。


2 相关工作

Transformer最初由Vaswani等人(2017年)提出,用于机器翻译,自此之后,它已成为许多自然语言处理(NLP)任务中的最新技术方法。大型基于Transformer的模型通常在大型语料库上进行预训练,然后针对手头的任务进行微调:BERT(Devlin等人,2019年)使用去噪自监督预训练任务,而GPT系列则使用语言建模作为其预训练任务(Radford等人,2018年;2019年;Brown等人,2020年)。

直接将自注意力应用到图像上,会要求每个像素都关注其他每个像素。由于像素数量的二次方成本,这在实际输入尺寸上是无法扩展的。因此,为了在图像处理的背景下应用Transformer,过去尝试了几种近似方法。Parmar等人(2018年)对每个查询像素仅在局部邻域内应用自注意力,而不是全局应用。这样的局部多头点积自注意力块可以完全替代卷积(Hu等人,2019年;Ramachandran等人,2019年;Zhao等人,2020年)。在另一系列研究中,稀疏Transformer(Child等人,2019年)采用可扩展的全局自注意力近似方法,以便应用于图像。另一种扩展注意力的方法是将其应用于不同大小的块(Weissenborn等人,2019年),在极端情况下,仅沿单个轴应用(Ho等人,2019年;Wang等人,2020年)。这些专门的注意力架构在计算机视觉任务上显示出了有希望的结果,但需要复杂的工程才能在硬件加速器上高效实现。

与我们的工作最相关的是Cordonnier等人(2020年)的模型,它从输入图像中提取2×2大小的块,并在顶部应用完全自注意力。这个模型与ViT非常相似,但我们的工作进一步证明了大规模预训练使普通的Transformer能够与(甚至优于)最先进的CNN竞争。此外,Cordonnier等人(2020年)使用的是2×2像素的小尺寸块,这使得模型只适用于小分辨率图像,而我们同时处理中等分辨率图像。

在将卷积神经网络(CNN)与自注意力形式结合方面也有许多兴趣,例如通过增强图像分类的特征图(Bello等人,2019年),或者通过自注意力进一步处理CNN的输出,例如用于目标检测(Hu等人,2018年;Carion等人,2020年)、视频处理(Wang等人,2018年;Sun等人,2019年)、图像分类(Wu等人,2020年)、无监督对象发现(Locatello等人,2020年)或统一的文本-视觉任务(Chen等人,2020年;Lu等人,2019年;Li等人,2019年)。

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

我们的工作增加了探索比标准ImageNet数据集更大规模图像识别的论文数量。使用额外的数据源可以在标准基准测试中实现最新技术结果(Mahajan等人,2018年;Touvron等人,2019年;Xie等人,2020年)。此外,Sun等人(2017年)研究了CNN性能如何随着数据集规模的增长而扩展,Kolesnikov等人(2020年);Djolonga等人(2020年)对从大规模数据集(如ImageNet-21k和JFT-300M)进行CNN迁移学习的情况进行实证探索。我们也关注这两个后的数据集,但训练的是Transformer而不是先前工作中使用的基于ResNet的模型。

3 方法

在模型设计上,我们尽可能紧密地跟随原始的Transformer(Vaswani等人,2017年)。这种有意设计的简单设置的一个优点是,可扩展的NLP Transformer架构及其高效实现几乎可以直接使用。

3.1 Vision Transformer

模型的概览如图1所示。标准的Transformer接收一个1D的token嵌入序列作为输入。为了处理2D图像,我们将图像$x ∈ R^{H×W×C}$重塑成一个扁平化的2D块序列$x_p ∈ R^{N×(P^2·C)}$,其中(H,W)是原始图像的分辨率,C是通道数,(P,P)是每个图像块的分辨率,$N = HW/P^2$是结果中的块数量,这也作为Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的潜在向量大小D,因此我们通过可训练的线性投影(公式1)将块扁平化并映射到D维度。我们将这个投影的输出称为块嵌入。

类似于BERT的[class] token,我们在嵌入块序列前加上一个可学习的嵌入($z_0^0 = x_class$),其在Transformer编码器($z_L^0$)的输出状态作为图像表示y(公式4)。在预训练和微调期间,都会附加一个分类头到$z_L^0$。分类头在预训练时由一个带有隐藏层的MLP实现,在微调时由一个单层线性层实现。

位置嵌入被添加到块嵌入中以保留位置信息。我们使用标准的可学习1D位置嵌入,因为我们没有观察到使用更高级的2D感知位置嵌入会带来显著的性能提升(附录D.4)。嵌入向量的结果序列作为编码器的输入。

Transformer编码器(Vaswani等人,2017年)由交替的多头自注意力(MSA,见附录A)和MLP块层组成(公式2,3)。在每个块之前应用Layernorm(LN),并在每个块之后应用残差连接(Wang等人,2019年;Baevski & Auli,2019年)。MLP包含两个带有GELU非线性的层。

归纳偏置。我们注意到,视觉Transformer比CNN具有更少的图像特定归纳偏置。在CNN中,局部性、二维邻域结构和平移等变性被嵌入到整个模型的每一层中。在ViT中,只有MLP层是局部的和平移等变的,而自注意力层是全局的。二维邻域结构的使用非常有限:在模型的开始通过将图像切割成块,以及在微调时调整不同分辨率图像的位置嵌入(如下所述)。除此之外,初始化时的位置嵌入不包含有关块的2D位置的信息,所有块之间的空间关系都必须从头开始学习。

混合架构。作为原始图像块的替代,输入序列可以从CNN的特征图(LeCun等人,1989年)中形成。在这个混合模型中,块嵌入投影E(公式1)应用于从CNN特征图提取的块。作为一个特殊情况,块可以具有1x1的空间尺寸,这意味着输入序列是通过简单地扁平化特征图的空间维度并投影到Transformer维度来获得的。如上所述添加分类输入嵌入和位置嵌入。

3.2 微调以及更高的分辨率

通常,我们会在大型数据集上预训练ViT,然后对其进行微调以适应(规模较小的)下游任务。为此,我们移除预训练的预测头,并附加一个零初始化的 D×K 完全连接层,其中 K 是下游任务的类别数。在更高分辨率下进行微调通常更有益(Touvron 等人,2019;Kolesnikov 等人,2020)。当输入更高分辨率的图像时,我们保持patch大小不变,这会导致更大的有效序列长度。ViT可以处理任意序列长度(直到内存限制),然而,预训练的位置嵌入可能不再有意义。因此,我们根据它们在原始图像中的位置,对预训练的位置嵌入进行二维插值。注意,这种分辨率调整和patch提取是唯一手动将关于图像二维结构的归纳偏差注入到ViT中的点。

4 实验

我们评估了ResNet、ViT以及混合模型的表示学习能力。为了理解每个模型对数据的需求,我们在不同大小的数据集上进行预训练,并在许多基准任务上进行评估。在考虑预训练模型的计算成本时,ViT的表现非常有利,在大多数识别基准测试中以较低的预训练成本达到最先进的水平。最后,我们进行了一个小的自监督实验,表明自监督ViT对未来充满希望。

4.1 实验设置

数据集。为了探索模型的可扩展性,我们使用了具有1000个类别和130万张图片的ILSVRC-2012 ImageNet数据集(以下简称为ImageNet),它的超集ImageNet-21k包含21k个类别和1400万张图片(Deng等人,2009年),以及包含18k个类别和3.03亿高分辨率图片的JFT(Sun等人,2017年)。我们遵循Kolesnikov等人(2020年)的方法,针对下游任务的测试集对预训练数据集进行了去重。我们将在这些数据集上训练的模型转移到几个基准任务上:使用原始验证标签和经过清理的ReaL标签的ImageNet(Beyer等人,2020年),CIFAR-10/100(Krizhevsky,2009年),Oxford-IIIT Pets(Parkhi等人,2012年),以及Oxford Flowers-102(Nilsback & Zisserman,2008年)。对于这些数据集,预处理遵循Kolesnikov等人(2020年)的方法。

我们还对19项任务的VTAB分类套件进行了评估(Zhai等人,2019b)。VTAB评估了低数据迁移到多样化任务的能力,每个任务使用1000个训练样本。任务被分为三组:自然——像上述的Pets、CIFAR等任务,专业——医疗和卫星图像,以及结构化——需要几何理解的任务,如定位。

模型变体。我们基于BERT(Devlin等人,2019年)中使用的配置来设定ViT的配置,如表1所总结。“基础版”和“Large”模型直接采用了BERT的配置,我们还增加了更大的“Huge”型号。在后续内容中,我们使用简短的符号来表示模型大小和输入patch大小:例如,ViT-L/16表示采用16×16输入patch大小的“Large”变体。请注意,ViT的序列长度与patch大小的平方成反比,因此使用较小patch大小的模型在计算上更昂贵。

对于基线CNN模型,我们使用的是ResNet(何恺明等人,2016年),但是将BatchNorm层(Ioffe & Szegedy,2015年)替换为GroupNorm(Wu & He,2018年),并且使用了标准化卷积(Qiao等人,2019年)。这些修改改善了模型的迁移能力(Kolesnikov等人,2020年),我们将修改后的模型称为“ResNet (BiT)”。对于混合模型,我们将中间特征图以一个“像素”的patch大小输入到ViT中。为了尝试不同的序列长度,我们要么(i)取常规ResNet50的第四阶段输出,要么(ii)移除第四阶段,将相同数量的层放置在第三阶段(保持总层数不变),然后取这个扩展的第三阶段的输出。选项(ii)导致序列长度变为原来的4倍,并且需要一个更昂贵的ViT模型。


Batch Normalization(批量归一化,简称BN)和Group Normalization(分组归一化,简称GN)都是深度学习中用于提高训练速度、稳定性和性能的技术。它们通过归一化层的输入来减少内部协变量偏移(internal covariate shift),即确保网络的每一层输入数据的分布保持相对稳定。下面分别介绍这两种技术:

  1. Batch Normalization(BN): 批量归一化是2015年由Ioffe和Szegedy提出的一种技术,用于改善神经网络的训练过程。它通过规范化层的输入来加速训练,减少训练时间,并有助于避免梯度消失或爆炸的问题。具体来说,BN对每个小批量(mini-batch)数据进行归一化处理,计算该批量数据的均值和方差,然后使用这些统计量来归一化当前层的输入。归一化公式如下:

其中,x是层的输入,mu是该批量数据的均值,sigma^2是方差,epsilon 是一个很小的常数,用来防止除以零。归一化后,BN层还会学习两个参数,缩放因子(gamma)和偏移量(beta),以便网络可以恢复数据的原始表示能力。

  1. Group Normalization(GN): 分组归一化是由Wu和He在2018年提出的一种技术,旨在解决BN对批量大小敏感的问题。在BN中,如果批量大小过小,那么计算得到的均值和方差可能不够准确,从而影响模型的性能。GN不依赖于批量大小,它将输入通道分成多个组,然后在每个组内进行归一化。具体来说,对于输入特征图的每个组,GN计算该组所有通道的均值和方差,然后使用这些统计量来归一化该组内每个通道的数据。归一化公式如下:

其中,x是组内的输入,mu是该组的均值,sigma^2是方差。与BN类似,GN也会学习缩放因子和偏移量来恢复数据的表示能力。

总的来说,BN和GN都是有效的归一化技术,但它们在不同的场景下有不同的优势。BN在大规模数据集和较大的批量大小下表现良好,而GN由于不依赖于批量大小,因此在小批量或无批量(如测试时)情况下更为稳定。


训练和微调。我们使用Adam优化器(Kingma & Ba,2015)训练所有模型,包括ResNets,设置β1 = 0.9,β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等人,2019;Wang等人,2020b)。

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

4.2 与目前最优的方法的比较

我们首先比较了我们最大的模型——ViT-H/14和ViT-L/16——与文献中报道的最先进CNN模型。第一个比较对象是Big Transfer (BiT)(Kolesnikov等人,2020年),它使用大型ResNets进行监督迁移学习。第二个比较对象是Noisy Student(Xie等人,2020年),这是一个大型的EfficientNet,使用半监督学习在删除了标签的ImageNet和JFT-300M上进行训练。目前,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节提供了不同架构在性能与计算方面的对照研究。最后,ViT-L/16模型在公共ImageNet-21k数据集上预训练,在大多数数据集上也表现良好,同时预训练所需的资源更少:它可以使用标准的8核云TPUv3在大约30天内训练完成。

图2将VTAB任务分解为各自的组别,并与此基准上的先前SOTA方法进行了比较:BiT、VIVI——在ImageNet和Youtube上共同训练的ResNet(Tschannen等人,2020年)以及S4L——ImageNet上的监督加半监督学习(Zhai等人,2019a)。ViT-H/14在自然(Natural)和结构化(Structured)任务上的表现超过了BiT-R152x4和其他方法。在专业(Specialized)任务上,前两个模型的性能相似。

4.3 预训练数据需求

视觉Transformer(ViT)在大规模数据集JFT-300M上进行预训练时表现良好。与ResNets相比,ViT对视觉的归纳偏好较少,那么数据集的规模对ViT的影响有多大呢?我们进行了两组实验。

首先,我们在规模逐渐增大的数据集上预训练ViT模型:ImageNet、ImageNet-21k和JFT-300M。为了提升在较小数据集上的性能,我们优化了三个基本的正则化参数——权重衰减、Dropout和label smoothing。图3(迁移到ImageNet的情况,虽然在小数据集上预训练时较大的ViT模型的表现不如BiT ResNets(阴影区域),但它们在大型数据集上预训练时表现突出。同样,随着数据集规模的增长,更大的ViT变体也会超过较小的变体。)显示了在ImageNet上微调后的结果(其他数据集的结果见表5(不同架构在预训练计算资源与性能对比:视觉Transformer、ResNets和混合模型。在相同的计算预算下,视觉Transformer通常比ResNets表现更好。对于较小模型尺寸,混合模型相较于纯Transformer有所提升,但在更大的模型中,这种差距消失了。 ))。当在最小的数据集ImageNet上预训练时,ViT-Large模型的表现不如ViT-Base模型,尽管进行了适度的正则化。在使用ImageNet-21k预训练后,它们的性能变得相似。只有在JFT-300M上预训练时,我们才看到了更大模型的全面优势。图3还显示了不同大小的BiT模型的性能范围。BiT CNN在ImageNet上的表现超过了ViT,但是随着数据集规模的增大,ViT逐渐赶超。同时,我们也可以看到,在这个场景下,较小的patch size的精确较高。

其次,我们在9M、30M和90M的随机子集以及完整的JFT-300M数据集上训练我们的模型。我们没有对较小的子集进行额外的正则化,并且对所有设置使用相同的超参数。这样,我们评估的是模型的内在属性,而不是正则化的效果。然而,我们使用了早停法,并报告了训练过程中达到的最佳验证准确率。为了节省计算资源,我们报告的是少样本线性准确率,而不是完整的微调准确率。图4(在ImageNet上的线性少样本评估与预训练大小的关系中,ResNets在较小的预训练数据集上表现更好,但很快达到平台期,而ViT在较大的预训练数据上表现更佳。ViT-b是将ViT-B模型的所有隐藏维度减半的变体。 )包含了结果。与计算成本相当的ResNets相比,视觉Transformer在较小数据集上更容易过拟合。例如,ViT-B/32比ResNet50略快;在9M子集上的表现要差得多,但在90M+子集上的表现更好。ResNet152x2和ViT-L/16的情况也是如此。这一结果加强了这样的直觉:卷积归纳偏好对于较小的数据集很有用,但对于更大的数据集,直接从数据中学习相关模式就足够了,甚至更有利。

总的来说,ImageNet上的少样本结果(图4)以及VTAB上的低数据结果(表2)对于极低数据迁移看起来很有前景。对ViT的少样本属性进行进一步分析是未来工作的一个令人兴奋的方向。

4.4 扩展性研究

我们通过评估JFT-300M的迁移性能,对不同模型进行了受控的扩展性研究。在这种设置中,数据规模不会限制模型的性能,我们评估每个模型的性能与预训练成本。模型组包括:7个ResNets,R50x1、R50x2、R101x1、R152x1、R152x2,预训练7个周期,加上R152x2和R200x3预训练14个周期;6个视觉Transformer,ViT-B/32、B/16、L/32、L/16,预训练7个周期,加上L/16和H/14预训练14个周期;以及5个混合模型,R50+ViT-B/32、B/16、L/32、L/16预训练7个周期,加上R50+ViT-L/16预训练14个周期(对于混合模型,模型名称末尾的数字不代表patch大小,而是代表ResNet骨干网络中的总下采样比率)。图5包含了迁移性能与总预训练计算的对比(有关计算成本的详细信息,请参阅附录D.5)。每个模型的详细结果在附录的表6中提供。可以观察到几个模式。首先,视觉Transformer在性能/计算权衡方面优于ResNets。ViT在达到相同性能时大约使用2-4倍更少的计算资源(5个数据集的平均值)。其次,混合模型在较小的计算预算下略优于ViT,但在更大的模型中,这种差异消失了。这一结果有些令人惊讶,因为人们可能会期望卷积局部特征处理在任何规模上都能帮助ViT。第三,视觉Transformer似乎在尝试的范围内没有饱和,这激发了未来的扩展努力。

4.5 探究ViT

为了开始理解视觉Transformer如何处理图像数据,我们分析了其内部表示。视觉Transformer的第一层将展平的图像块线性映射到较低维的空间(公式1)。图7(左图)显示了学习到的嵌入滤波器的顶部主成分。这些成分类似于每个图像块内精细结构的低维表示的合理基函数。

在映射之后,学习到的位置嵌入被添加到图像块表示中。图7(中间)显示,模型学会了在位置嵌入的相似性中编码图像内的距离,即更接近的图像块倾向于具有更相似的位置嵌入。此外,还出现了行-列结构;同一行/列中的图像块具有相似的嵌入。最后,对于较大的网格,有时可以明显看到正弦波结构(附录D)。位置嵌入学会表示二维图像拓扑解释了为什么手工制作的二维感知嵌入变体没有带来改进(附录D.4)。

自注意力允许ViT即使在最低层也能整合整个图像的信息。我们调查了网络在多大程度上利用这种能力。具体来说,我们根据注意力权重(图7,右图)计算了图像空间中信息整合的平均距离。这种“注意力距离”类似于CNN中的接受场大小。我们发现,一些注意力头在最低层就已经关注了大部分图像,这表明模型确实使用了全局整合信息的能力。其他注意力头在低层的注意力距离一直保持很小。这种高度局部化的注意力在混合模型中不太明显,这些模型在Transformer之前应用了ResNet(图7,右图),这表明它可能与CNN中的早期卷积层具有类似的功能。此外,随着网络深度的增加,注意力距离也会增加。总体而言,我们发现模型关注的是对于分类语义相关的图像区域(图6)。

4.6 自监督

Transformers在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; Hénaff et al., 2020)的工作留作未来研究。

5 总结

我们已经探索了将Transformer直接应用于图像识别。与以前在计算机视觉中使用自注意力的工作不同,我们除了最初的图像块提取步骤外,没有在架构中引入特定于图像的归纳偏好。相反,我们将图像解释为一系列图像块,并使用标准Transformer编码器对其进行处理,就像在自然语言处理(NLP)中一样。这种简单但可扩展的策略在与大规模数据集的预训练结合时表现出奇地好。因此,视觉Transformer(ViT)在许多图像分类数据集上匹配或超过了最新技术水平,同时预训练成本相对较低。

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

6 附录

A 多头注意力

标准的qkv自注意力(SA,Vaswani et al. 2017)是神经网络架构的一个流行构建块。对于输入序列z中的每个元素z ∈ R^{N×D},我们计算与序列中所有值v的加权和。注意力权重Aij基于序列中两个元素及其各自的查询q^i和键k^j表示之间的成对相似性。

多头自注意力(MSA)是自注意力(SA)的一个扩展,我们在其中并行运行k个自注意力操作,称为“头”,并连接它们的输出进行投影。为了在改变k时保持计算量和参数数量不变,通常将D_h(公式5)设置为D/k。

B 实验细节

B.1 训练

表3总结了我们不同模型的训练设置。我们发现,当在ImageNet上从头开始训练模型时,强烈的正则化非常关键。如果使用Dropout,除了qkv投影和在positional- to patch嵌入后直接应用外,每个密集层之后都会应用。混合模型的训练与它们对应的ViT模型完全相同。最后,所有训练都在224分辨率下完成。

B.1.1 微调

我们使用动量为0.9的SGD对所有ViT模型进行微调。我们对学习率进行了小范围的网格搜索,详见表4中的学习率范围。为此,我们使用训练集的小部分子集(宠物和花卉10%,CIFAR 2%,ImageNet 1%)作为开发集,并在剩余数据上进行训练。对于最终结果,我们在整个训练集上进行训练,并在各自的测试数据上进行评估。对于ResNets和混合模型的微调,我们使用完全相同的设置,唯一的例外是在ImageNet上我们在学习率扫描中增加了另一个值0.06。此外,对于ResNets,我们还运行了Kolesnikov等人(2020)的设置,并选择了这次运行和我们的扫描中的最佳结果。最后,除非另有说明,所有微调实验都在384分辨率下进行(与训练不同分辨率的微调是常见的做法)。

当将ViT模型迁移到另一个数据集时,我们移除整个头部(两个线性层),并用一个单一的、零初始化的线性层替换它,该层输出目标数据集所需的类别数量。我们发现这种方法比简单地重新初始化最后一层更为稳健。

对于VTAB,我们遵循Kolesnikov等人(2020)的协议,并为所有任务使用相同的超参数设置。我们使用0.01的学习率,训练2500步(表4)。我们通过在两个学习率和两个计划之间进行小范围的扫描,并选择在200个样本验证集上VTAB得分最高的设置来选择这个设置。我们遵循Kolesnikov等人(2020)中使用的预处理方法,但我们不使用特定任务的输入分辨率。相反,我们发现对于所有任务,视觉Transformer最受益于高分辨率(384×384)。

B.1.2 自监督

我们在初步的自监督实验中采用了掩蔽图像块预测目标。为此,我们通过以下方式损坏了50%的图像块嵌入:80%的情况下用可学习的[mask]嵌入替换它们,10%的情况下用随机的其他图像块嵌入替换,还有10%的情况下保持不变。这个设置与Devlin等人在2019年用于语言的设置非常相似。最后,我们使用它们各自的图像块表示预测每个损坏块的3位平均颜色(即总共512种颜色)。

我们在JFT上以批量大小4096训练了我们的自监督模型100万步(大约14个周期)。我们使用Adam优化器,基础学习率为2×10^-4,预热10k步,并使用余弦学习率衰减。作为预训练的预测目标,我们尝试了以下设置:1)仅预测平均3位颜色(即1次预测512种颜色),2)同时预测16×16图像块的4×4缩小版和3位颜色(即16次预测512种颜色),3)使用L2对完整图像块进行回归(即对3个RGB通道进行256次回归)。令人惊讶的是,我们发现所有方法都相当有效,尽管L2稍微差一些。我们只报告选项1)的最终结果,因为它在少样本性能上表现最好。我们还尝试了Devlin等人在2019年使用的15%损坏率,但在我们的少样本指标上结果也稍微差一些。

最后,我们想指出,我们的掩蔽图像块预测实现不需要如此大量的预训练或像JFT这样的大型数据集就能在ImageNet分类上获得类似的性能提升。也就是说,我们在10万步预训练后观察到下游性能的递减回报,并且在ImageNet上预训练时看到了类似的增益。

C 附加结果

我们报告了与论文中所示图表相对应的详细结果。表5对应于论文中的图3,显示了在越来越大的数据集上预训练的不同ViT模型的迁移性能:ImageNet、ImageNet-21k和JFT-300M。表6对应于论文中的图5,显示了不同大小的ViT、ResNet和混合模型的迁移性能,以及它们预训练的估计计算成本。

D 附加分析

D.1 ResNets使用SGD与Adam

ResNets通常使用SGD进行训练,但在这里,作者选择使用Adam优化器,这在某种程度上是不寻常的。为了证明这个选择的合理性,作者进行了实验,比较了在JFT数据集上使用SGD和Adam进行预训练的两种ResNet模型(50x1和152x2)的微调性能,见表7。实验结果表明,使用Adam进行预训练在大多数数据集上都比使用SGD表现得更好,平均性能也更佳。这证实了在JFT数据集上使用Adam作为ResNet的预训练优化器是有道理的。值得注意的是,与Kolesnikov等人(2020)报告的数字相比,这些绝对数值较低,因为我们的模型仅预训练了7个epoch,而不是30个epoch。

D.2 transoformer 形状(结构)

我们对Transformer架构的不同维度进行了缩放消融实验,以找出最适合大规模模型缩放的维度。图8显示了不同配置在ImageNet上的5-shot性能。所有配置都基于一个ViT模型,具有8层,D=1024,D_MLP=2048和一个patch大小为32,是所有线的交点。我们可以看到,扩展深度会带来最大的改进,直到64层仍然明显。然而,在16层之后,收益递减已经明显可见。有趣的是,缩放网络的宽度似乎会产生最小的变化。减小patch大小,从而增加有效序列长度,展现出惊人的稳健改善,而不会引入额外的参数。这些发现表明,计算能力可能是性能的更好预测因素,而不是参数数量,如果需要缩放,则应强调深度而不是宽度。总的来说,我们发现按比例缩放所有维度会带来稳健的改进。

D.3 头的类型和CLASS token

为了尽可能接近原始的Transformer模型,我们使用了额外的[class]标记,该标记被用作图像表示。然后,这个标记的输出通过一个小型的多层感知器(MLP)转化为类别预测,其中单个隐藏层使用tanh作为非线性激活函数。这种设计是从文本的Transformer模型继承而来的,我们在主要论文中使用了它。最初尝试仅使用图像块嵌入,然后对它们进行全局平均池化(GAP),后面跟着一个线性分类器——就像ResNet的最终特征图一样——效果非常差。然而,我们发现这既不是因为额外的标记,也不是因为GAP操作。相反,性能差异完全可以通过不同的学习率要求来解释,详见图9。

D.4 positional embedding

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

  • 提供没有位置信息:将输入视为一堆图像块。
  • 1维位置嵌入:将输入视为一系列按光栅顺序排列的图像块(本文中所有其他实验的默认设置)。
  • 2维位置嵌入:将输入视为二维网格中的图像块。在这种情况下,学习两组嵌入,每组对应一个轴,X轴嵌入和Y轴嵌入,每个嵌入的大小为D/2。然后,根据输入图像块在路径上的坐标,我们连接X和Y嵌入,以获得该图像块的最终位置嵌入。
  • 相对位置嵌入:考虑图像块之间的相对距离来编码空间信息,而不是它们的绝对位置。为此,我们使用1维相对注意力,其中我们定义了所有可能的图像块对的相对距离。因此,对于每一对给定的图像块(一个作为查询,另一个作为注意力机制中的键/值),我们有一个偏移量p_q - p_k,每个偏移量都与一个嵌入相关联。然后,我们简单地运行额外的注意力,其中我们使用原始的查询(查询的内容),但使用相对位置嵌入作为键。然后我们使用相对注意力的对数几率作为偏置项,并在应用softmax之前将其添加到主要内容注意力(基于内容的注意力)的logits上。

除了不同编码空间信息的方式,我们还尝试了将这些信息整合到我们的模型中的不同方式。对于1维和2维位置嵌入,我们尝试了三种不同的情况:(1)在模型的主干之后立即将位置嵌入添加到输入中,并在将输入馈送到Transformer编码器之前(本文中所有其他实验的默认设置);(2)在每个层的开始学习并添加位置嵌入到输入中;(3)在每个层的开始添加学习到的位置嵌入到输入中(在层之间共享)。

表8总结了这项在ViT-B/16模型上进行的消融研究的结果。正如我们所看到的,虽然没有位置嵌入的模型与具有位置嵌入的模型之间的性能有很大的差距,但不同编码位置信息的方式之间几乎没有差异。我们推测,由于我们的Transformer编码器操作的是块级输入,而不是像素级,所以编码空间信息的不同方式之间的差异不那么重要。更准确地说,在块级输入中,空间维度比原始像素级输入要小得多,例如,14×14相对于224×224,在这个分辨率下学习表示空间关系对于这些不同的位置编码策略来说同样容易。尽管如此,网络学习到的位置嵌入相似性的具体模式取决于训练超参数(图10)。

D.5 经验计算成本

我们对于架构在硬件上的实际运行速度也感兴趣,这并不总是能够通过理论计算的浮点操作数(FLOPs)准确预测,因为还需要考虑诸如lane widths和缓存大小等细节。为此,我们对主要感兴趣的模型在TPUv3加速器上进行了推理速度的计时;推理速度和反向传播速度之间的差异是一个恒定的模型无关因素。


“Lane widths”(车道宽度)在计算机硬件和架构中通常指的是PCIe(Peripheral Component Interconnect Express)总线中的通道数量。PCIe是一种高速串行计算机扩展总线标准,用于连接计算机主板和各种扩展卡、设备以及组件。PCIe链路可以由多个通道(lanes)组成,每个lane可以看作是数据传输的一条路径。车道宽度指的是这些lane的数量,例如PCIe x1、x4、x8、x16,其中数字表示lane的数量。车道数量越多,数据传输的带宽就越大,但也需要更多的物理空间和可能更高的功耗。

在PCIe总线中,每个lane由一组差分信号对组成,包括一个发送对和一个接收对,用于高速串行数据传输。差分信号传输可以提高信号的完整性和抗干扰能力,从而使得PCIe链路能够在较长距离上保持可靠的数据传输。随着技术的发展,PCIe的不同版本支持不同的数据传输速率,从早期的每个lane 2.5GT/s(PCIe 1.0)到最新的32GT/s(PCIe 6.0),并且使用了不同的编码方式来提高传输效率。


图12(左图)显示了每个核心每秒钟可以处理的图像数量,涵盖了各种输入尺寸。每个点都表示在广泛的批量大小下测量的峰值性能。可以看到,对于最大模型和最高分辨率,ViT的理论双二次方尺度与图像尺寸的关系才开始显现。

另一个感兴趣的量是每个模型能够适应核心的最大批量大小,越大越好,这有助于扩展到大型数据集。图12(右图)显示了相同一组模型的这个量。这表明大型ViT模型在内存效率方面明显优于ResNet模型。

D.6 Axial Attention

轴向注意力(Axial Attention)是一种简单但有效的技术,用于在组织为多维张量的大输入上运行自注意力。轴向注意力的一般思想是执行多个注意力操作,每个操作沿着输入张量的一个轴进行,而不是将1维注意力应用于输入的扁平化版本。在轴向注意力中,每个注意力沿着特定轴混合信息,同时保持其他轴上的信息独立。沿着这个方向,Wang等人(2020b)提出了AxialResNet模型,其中ResNet50中所有卷积核大小为3×3的卷积被轴向自注意力替换,即行列注意力,并增加了相对位置编码。我们已经实现了AxialResNet作为基线模型。

此外,我们修改了ViT以处理二维形状的输入,而不是一维的图像块序列,并融入了轴向Transformer块,在这些块中,不是自注意力后跟一个MLP,而是先进行行自注意力加MLP,然后进行列自注意力加MLP。

图13展示了在JFT数据集上预训练时,Axial ResNet、Axial-ViT-B/32和Axial-ViT-B/16在ImageNet 5shot线性上的性能,与预训练计算量(无论是FLOPs的数量还是推理时间,即每秒示例数)的关系。正如我们所见,无论是Axial-ViT-B/32还是Axial-ViT-B/16在性能上都优于它们的ViT-B对应物,但代价是更多的计算。这是因为在Axial-ViT模型中,每个具有全局自注意力的Transformer块被两个轴向Transformer块替换,一个用于行,一个用于列,尽管自注意力操作的序列长度在轴向情况下更小,但每个Axial-ViT块都有一个额外的MLP。对于AxialResNet,尽管在准确性/计算权衡方面看起来合理(图13,左图),但其在TPUs上的朴素实现速度极慢(图13,右图)。

D.7 注意力距离

为了理解ViT如何使用自注意力机制在图像中整合信息,我们分析了不同层次中注意力权重所跨越的平均距离(见图11)。这种“注意力距离”类似于卷积神经网络(CNNs)中的接受场大小。在较低层次中,平均注意力距离在各个头之间高度可变,有些头关注整个图像的大部分区域,而其他头则只关注查询位置附近的小区域。随着网络深度的增加,所有头的注意力距离都增加了。在网络的后半部分,大多数头在各个标记之间广泛地关注。

D.8 注意力图

为了计算输出标记到输入空间的注意力图(见图6和14),我们使用了注意力展开(Attention Rollout)(Abnar & Zuidema,2020)。简单来说,我们对ViT-L/16的所有头的注意力权重进行了平均,然后递归地乘以所有层的权重矩阵。这样可以考虑到通过所有层的标记的注意力混合。

D.9 ObjectNet结果

我们还在ObjectNet基准测试中评估了我们的旗舰模型ViT-H/14,按照Kolesnikov等人(2020)的评估设置,得到了82.1%的top-5准确率和61.7%的top-1准确率。

D.10 VTAB细目

表9显示了在VTAB-1k任务中每个任务获得的分数。