摘要
虽然Transformer架构已成为自然语言处理任务的实际标准,但其在计算机视觉领域的应用仍然有限。在计算机视觉中,注意力机制要么与卷积网络结合使用,要么用于替换卷积网络的某些组件,同时保持其整体结构不变。我们证明了这种对卷积神经网络(CNN)的依赖并不是必需的,并且直接应用于图像块序列的纯Transformer在图像分类任务上可以表现得非常好。当在大量数据上进行预训练并迁移到多个中等规模或小规模图像识别基准(如ImageNet、CIFAR-100、VTAB等)时,与最先进的卷积网络相比,视觉Transformer(ViT)取得了卓越的结果,同时所需的训练计算资源也大幅减少。
1 引言
基于自注意力的架构,尤其是Transformer(Vaswani等人,2017),已成为自然语言处理(NLP)的首选模型。主要方法是先在大规模文本语料库上进行预训练,然后在较小的特定任务数据集上进行微调(Devlin等人,2019)。得益于Transformer的计算效率和可扩展性,训练前所未有的大规模模型成为可能,参数超过1000亿(Brown等人,2020;Lepikhin等人,2020)。随着模型和数据集的增长,性能仍未出现饱和迹象。
然而,在计算机视觉领域,卷积架构仍然占据主导地位(LeCun等人,1989;Krizhevsky等人,2012;He等人,2016)。受NLP成功的启发,多项工作尝试将类似CNN的架构与自注意力相结合(Wang等人,2018;Carion等人,2020),有些则完全取代了卷积(Ramachandran等人,2019;Wang等人,2020a)。尽管后一类模型在理论上具有效率,但由于使用了专门的注意力模式,尚未在现代硬件加速器上实现有效扩展。因此,在大规模图像识别方面,经典的ResNet类架构仍然是最先进的(Mahajan等人,2018;Xie等人,2020;Kolesnikov等人,2020)。
受NLP中Transformer扩展成功的启发,我们尝试将标准Transformer直接应用于图像,尽可能少地进行修改。为此,我们将图像分割成块,并将这些块的线性嵌入序列作为Transformer的输入。图像块的处理方式与NLP应用中的标记(单词)相同。我们以监督的方式训练模型进行图像分类。
当在中等规模的数据集(如ImageNet)上进行训练且不使用强正则化时,这些模型的准确率比同等规模的ResNet低几个百分点。这种看似令人沮丧的结果是可以预料的:与卷积神经网络(CNN)固有的归纳偏置(如平移等变性和局部性)相比,Transformer缺乏其中一些归纳偏置,因此在数据量不足的情况下训练时泛化能力不强。
然而,如果模型在更大的数据集(1400万至3亿张图像)上进行训练,情况就会发生变化。我们发现,大规模训练胜过了归纳偏置。我们的视觉Transformer(ViT)在足够规模的预训练后,转移到数据量较少的任务上时,取得了出色的结果。当在公开的ImageNet-21k数据集或内部的JFT-300M数据集上进行预训练时,ViT在多个图像识别基准测试中接近或超越了最先进的水平。特别是,最佳模型在ImageNet上的准确率达到88.55%,在ImageNet-ReaL上的准确率达到90.72%,在CIFAR-100上的准确率达到94.55%,在包含19个任务的VTAB套件上的准确率达到77.63%。
2 相关工作
Transformer由Vaswani等人(2017)提出,用于机器翻译,并已成为许多自然语言处理(NLP)任务的最先进方法。基于Transformer的大型模型通常首先在大型语料库上进行预训练,然后再针对手头任务进行微调:BERT(Devlin等人,2019)使用去噪自监督预训练任务,而GPT系列工作则使用语言建模作为其预训练任务(Radford等人,2018;2019;Brown等人,2020)。
将自注意力机制直接应用于图像会要求每个像素都关注其他所有像素。由于像素数量的平方级成本,这种方法无法扩展到实际的输入大小。因此,为了在图像处理中应用Transformer,过去已经尝试了几种近似方法。Parmar等人(2018)仅在每个查询像素的局部邻域内应用自注意力,而不是全局应用。这样的局部多头点积自注意力块可以完全替代卷积(Hu等人,2019;Ramachandran等人,2019;Zhao等人,2020)。在另一项工作中,Sparse Transformer(Child等人,2019)对全局自注意力进行了可扩展的近似,以便能够应用于图像。扩展注意力的另一种方法是在不同大小的块中应用它(Weissenborn等人,2019),在极端情况下仅沿单个轴应用(Ho等人,2019;Wang等人,2020a)。许多这种专门的注意力架构在计算机视觉任务上展示了有希望的结果,但需要在硬件加速器上高效实现复杂的工程。
与我们的工作最相关的是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等人,2020c;Lu等人,2019;Li等人,2019)。
另一个最近的相关模型是图像GPT(iGPT)(Chen等人,2020a),它在降低图像分辨率和颜色空间后将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 视觉Transformer(ViT)
模型的概述如图1所示。标准的Transformer接收一维(1D)的标记嵌入序列作为输入。为了处理二维(2D)图像,我们将图像x重塑为二维(2D)块序列xp,其中x的维度是RH×W×C,xp的维度是RN×(P2·C)。(H,W)是原始图像的分辨率,C是通道数,(P,P)是每个图像块的分辨率,N=HW/P2是块的数量,也作为Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的潜在向量大小D,因此我们将块展平,并通过可训练的线性投影(公式1)映射到D维。我们将此投影的输出称为块嵌入。
图1: Model overview
与BERT的[class]标记类似,我们在嵌入块序列之前添加了一个可学习的嵌入(z0 0 = xclass),其在Transformer编码器输出处的状态(z0 L)作为图像表示y(公式4)。在预训练和微调期间,都会将分类头附加到z0 L。预训练时的分类头是通过具有一个隐藏层的多层感知机(MLP)实现的,而微调时则是通过单个线性层实现的。
为了保留位置信息,将位置嵌入添加到块嵌入中。我们使用标准的可学习一维(1D)位置嵌入,因为我们没有观察到使用更高级的二维(2D)感知位置嵌入会带来显著的性能提升(附录D.4)。所得的嵌入向量序列作为编码器的输入。
Transformer编码器(Vaswani等人,2017)由多头自注意力(MSA,见附录A)和MLP块(公式2、3)的交替层组成。在每个块之前应用层归一化(LN),在每个块之后应用残差连接(Wang等人,2019;Baevski和Auli,2019)。
多层感知机(MLP)包含两层,并使用了高斯误差线性单元(GELU)作为非线性激活函数。
归纳偏置。我们注意到,与卷积神经网络(CNNs)相比,Vision Transformer(Vision Transformer,简称ViT)具有更少的与图像相关的归纳偏置。在CNNs中,局部性、二维邻域结构和平移等变性被内置于整个模型的每一层中。而在ViT中,只有MLP层是局部的和平移等变的,而自注意力层是全局的。二维邻域结构的使用非常有限:仅在模型开始时通过将图像切割成块,以及在微调时调整不同分辨率图像的位置嵌入(如下所述)。除此之外,初始化时的位置嵌入不包含关于块二维位置的信息,所有块之间的空间关系都必须从头开始学习。
混合架构。作为原始图像块的替代方案,输入序列可以由卷积神经网络(CNN,LeCun等人,1989)的特征图形成。在这个混合模型中,将块嵌入投影E(公式1)应用于从CNN特征图中提取的块。作为一种特殊情况,块的空间大小可以为1x1,这意味着输入序列是通过简单地展平特征图的空间维度并投影到Transformer维度上而获得的。分类输入嵌入和位置嵌入的添加方式如上所述。
3.2 微调和更高分辨率
通常,我们在大型数据集上对ViT进行预训练,并针对(较小的)下游任务进行微调。为此,我们移除预训练的预测头,并附加一个零初始化的D×K前馈层,其中K是下游任务的类别数。在比预训练时更高的分辨率下进行微调通常是有益的(Touvron等人,2019;Kolesnikov等人,2020)。当输入更高分辨率的图像时,我们保持块大小不变,这会导致更长的有效序列长度。Vision Transformer可以处理任意长度的序列(受内存限制),然而,预训练的位置嵌入可能不再有意义。因此,我们根据原始图像中的位置对预训练的位置嵌入进行二维插值。请注意,这种分辨率调整和块提取是仅有的两个点,在这两点上,关于图像二维结构的归纳偏置被手动注入到Vision Transformer中。
4 实验
我们评估了ResNet、Vision Transformer(ViT)和混合模型的表征学习能力。为了了解每个模型的数据需求,我们在不同大小的数据集上进行预训练,并评估了许多基准任务。在考虑模型预训练的计算成本时,ViT表现非常出色,以较低的预训练成本在大多数识别基准上达到了最先进的水平。最后,我们进行了一个小规模的自监督实验,并表明自监督的ViT在未来充满希望。
4.1 设置
数据集。为了探索模型的可扩展性,我们使用了ILSVRC-2012 ImageNet数据集,该数据集包含1000个类别和130万张图像(以下简称ImageNet),其超集ImageNet-21k包含21000个类别和1400万张图像(Deng等人,2009),以及JFT(Sun等人,2017)数据集,包含18000个类别和3.03亿张高分辨率图像。我们按照Kolesnikov等人(2020)的方法,对预训练数据集进行去重,以避免与下游任务的测试集重复。我们将在这些数据集上训练的模型迁移到几个基准任务上:ImageNet(使用原始验证标签和清理后的ReaL标签(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所示。我们直接采用BERT中的“Base”和“Large”模型,并增加了更大的“Huge”模型。在下文中,我们使用简短符号来表示模型大小和输入块大小:例如,ViT-L/16表示“Large”变体,其输入块大小为16×16。请注意,Transformer的序列长度与块大小的平方成反比,因此,块大小较小的模型计算成本更高。
对于基线CNN,我们使用ResNet(He等人,2016),但将批量归一化层(Ioffe & Szegedy,2015)替换为组归一化(Wu & He,2018),并使用标准化卷积(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)。
指标。我们通过少量样本或微调准确率来报告下游数据集上的结果。微调准确率反映了在每个数据集上微调后每个模型的性能。少量样本准确率是通过解决一个正则化最小二乘回归问题获得的,该问题将(冻结的)训练图像子集的表示映射到f-1;1gK目标向量。这种表述允许我们以封闭形式恢复精确解。尽管我们主要关注微调性能,但有时我们会使用线性少量样本准确率进行快速即时评估,因为微调的成本太高。
4.2 与最先进技术的比较
我们首先将我们最大的模型——ViT-H/14和ViT-L/16——与文献中的最先进CNN进行比较。第一个比较点是Big Transfer(BiT)(Kolesnikov等人,2020),它使用大型ResNet进行有监督的迁移学习。第二个是Noisy Student(Xie等人,2020),它是一个大型EfficientNet,使用ImageNet和JFT-300M上的半监督学习进行训练,其中移除了标签。目前,Noisy Student是ImageNet上的最先进技术,而BiT-L是此处报告的其他数据集上的最先进技术。所有模型都在TPUv3硬件上进行训练,我们报告了预训练每个模型所需的TPUv3核心天数,即用于训练的TPUv3核心数(每个芯片2个)乘以训练天数。
表2显示了结果。在JFT-300M上预训练的较小模型ViT-L/16在所有任务上的表现都优于BiT-L(它也在同一数据集上进行预训练),同时所需的计算资源大大减少。更大的模型ViT-H/14进一步提高了性能,尤其是在更具挑战性的数据集上——ImageNet、CIFAR-100和VTAB套件。
表2: Comparison with state of the art on popular image classification benchmarks.
模型预训练所需的计算量仍然远低于之前的最优水平。然而,我们注意到预训练效率不仅受架构选择的影响,还受其他参数的影响,如训练计划、优化器、权重衰减等。我们在第4.4节中对不同架构的性能与计算量进行了受控研究。最后,在公开的ImageNet-21k数据集上预训练的ViT-L/16模型在大多数数据集上表现良好,同时预训练所需的资源更少:使用标准的8核云TPUv3大约需要30天即可训练完成。
图2将VTAB任务分解为各自的组,并与该基准上的先前最优方法进行了比较:BiT、VIVI(一种在ImageNet和Youtube上共同训练的ResNet(Tschannen等,2020年))以及S4L(在ImageNet上进行的有监督和半监督学习(Zhai等,2019a年))。在自然和结构化任务上,ViT-H/14的表现优于BiT-R152x4和其他方法。在专门化任务上,前两个模型的性能相似。
图2. Breakdown of VTAB performance in Natural, Specialized, and Structured task groups.
4.3 预训练数据需求
Vision Transformer在大型JFT-300M数据集上预训练时表现良好。与具有更多视觉归纳偏置的ResNet相比,数据集的大小对Vision Transformer有多重要?我们进行了两组实验。
首先,我们在不同大小的数据集上预训练ViT模型:ImageNet、ImageNet-21k和JFT-300M。为了提升在较小数据集上的性能,我们优化了三个基本的正则化参数:权重衰减、dropout和标签平滑。图3显示了微调至ImageNet后的结果(其他数据集的结果见表5)。在最小的数据集ImageNet上预训练时,尽管进行了(适度的)正则化,但ViT-Large模型的表现仍不如ViT-Base模型。在ImageNet-21k上预训练后,它们的性能相似。只有在JFT-300M上,我们才能看到更大模型的全部优势。图3还显示了不同大小的BiT模型的性能范围。在ImageNet上,BiT CNN的表现优于ViT,但在更大的数据集上,ViT超越了BiT。
图3: Transfer to ImageNet
其次,我们在9M、30M、90M以及完整的JFT-300M数据集的随机子集上训练我们的模型。我们没有在较小的子集上进行额外的正则化,并对所有设置使用相同的超参数。这样,我们可以评估模型的固有属性,而不是正则化的影响。然而,我们使用了早停法,并报告了训练过程中达到的最佳验证准确率。为了节省计算量,我们报告了少样本线性准确率,而不是完整的微调准确率。图4包含了结果。在较小的数据集上,与具有相似计算成本的ResNet相比,Vision Transformer更容易过拟合。例如,ViT-B/32的计算速度略快于ResNet50;它在9M子集上的表现要差得多,但在90M+子集上的表现更好。ResNet152x2和ViT-L/16也是如此。这一结果强化了直觉:对于较小的数据集,卷积归纳偏置是有用的,但对于较大的数据集,直接从数据中学习相关模式是足够的,甚至是有益的。
图4: Linear few-shot evaluation on ImageNet versus pre-training size
4.4 缩放研究
我们通过评估从JFT-300M迁移的性能,对不同模型进行了受控的缩放研究。在这种设置中,数据大小不会限制模型的性能,我们评估了每个模型的性能与预训练成本之间的关系。模型集包括:7个ResNet(R50x1、R50x2、R101x1、R152x1、R152x2,预训练7个周期),以及R152x2和R200x3(预训练14个周期);6个Vision 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个周期)(对于混合模型,模型名称末尾的数字不代表块大小,而是ResNet主干中的总下采样率)。
图5显示了迁移性能与总预训练计算量之间的关系(有关计算成本的详细信息,请参见附录D.5)。表6提供了每个模型的详细结果,见附录。可以观察到一些规律。首先,在性能与计算量的权衡上,Vision Transformer优于ResNet。ViT使用大约2-4倍的计算量即可达到相同的性能(5个数据集的平均值)。其次,在较小的计算预算下,混合模型的表现略优于ViT,但对于更大的模型,这种差异消失了。这个结果有些令人惊讶,因为人们可能会期望卷积局部特征处理在任何大小的ViT中都有所帮助。第三,在所尝试的范围内,Vision Transformer似乎并未达到饱和,这激发了未来进行缩放研究的动力。
图5. Performance versus pre-training compute for different architectures: Vision Transformers,
ResNets, and hybrids
4.5 Vision Transformer
为了开始理解视觉Transformer如何处理图像数据,我们分析了其内部表示。Vision Transformer的第一层将扁平化的图像块线性投影到一个低维空间(公式1)。图7(左)显示了所学嵌入滤波器的主要成分。这些成分类似于每个图像块内部精细结构的低维表示的基函数。在投影之后,将所学的位置嵌入添加到图像块表示中。图7(中)显示,模型学会了在位置嵌入的相似性中编码图像内的距离,即更近的图像块往往具有更相似的位置嵌入。此外,还出现了行-列结构,即同一行/列中的图像块具有相似的嵌入。最后,对于更大的网格,有时会出现正弦结构(附录D)。位置嵌入学会了表示二维图像拓扑,这解释了为什么手工制作的二维感知嵌入变体没有带来改进(附录D.4)。自注意力使视觉Transformer(ViT)即使在最低层也能整合整个图像的信息。我们研究了网络在多大程度上利用了这个能力。具体来说,我们基于注意力权重计算了跨图像空间整合信息的平均距离(图7,右)。这个“注意力距离”类似于卷积神经网络(CNN)中的感受野大小。
我们发现,一些注意力头已经在最低层关注到了图像的大部分,这表明模型确实使用了全局整合信息的能力。其他注意力头在低层中始终具有较小的注意力距离。这种高度局部化的注意力在混合模型中不那么明显,混合模型在应用Transformer之前应用了ResNet(图7,右),这表明它可能具有与CNN中早期卷积层相似的功能。此外,注意力距离随着网络深度的增加而增加。总体而言,我们发现模型关注到了对分类具有语义相关性的图像区域(图6)。
图6: Representative examples of attention from the output token to the input space.
4.6 自监督
Transformer 在自然语言处理任务上表现出色。然而,它们的成功不仅源于其出色的可扩展性,还源于大规模的自监督预训练(Devlin等人,2019;Radford等人,2018)。我们还对掩码块预测进行了初步的探索,以实现自监督,模仿BERT中使用的掩码语言建模任务。通过自监督预训练,我们较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率,比从头开始训练显著提高了2%,但仍比监督预训练低4%。附录B.1.2包含更多细节。我们将对比式预训练(Chen等人,2020b;He等人,2020;Bachman等人,2019;Henaff等人,2020)的探索留给未来的工作。
5 结论
我们探索了Transformer在图像识别中的直接应用。与先前在计算机视觉中使用自注意力的工作不同,我们没有在架构中引入除初始块提取步骤之外的特定于图像的归纳偏置。相反,我们将图像解释为一系列块,并使用标准的Transformer编码器进行处理,就像在自然语言处理中一样。这种简单但可扩展的策略在结合大型数据集上的预训练时,效果出奇地好。因此,Vision Transformer在许多图像分类数据集上的表现与或超过了最先进的水平,同时预训练成本相对较低。
尽管这些初步结果令人鼓舞,但仍有许多挑战需要解决。其中之一是将ViT应用于其他计算机视觉任务,如检测和分割。我们的结果与Carion等人(2020)的结果表明,这种方法很有前景。另一个挑战是继续探索自监督预训练方法。我们的初步实验表明,自监督预训练有所改善,但自监督与大规模监督预训练之间仍存在较大差距。最后,进一步扩展ViT可能会带来性能的提升。