导语
- 会议:ICLR 2021
- 链接:arxiv.org/pdf/2010.11...
虽然Transformer架构已成为NLP任务的事实标准,但其在计算机视觉领域的应用仍然有限。在CV领域中,注意力要么与卷积网络结合使用,要么用来替代卷积网络的某些组件。本文研究表明,这种对卷积网络(CNNs)的依赖并非必要。一个直接应用于图像块序列的纯粹Transformer在图像分类任务上可以表现得非常出色。当在大量数据上预训练,并转移到多个中等或小型图像识别基准(如ImageNet、CIFAR-100、VTAB等)后,视觉Transformer(ViT)与最先进的卷积网络相比取得了卓越的成果,同时在训练时所需的计算资源大大减少。
1 简介
基于自注意力的架构,特别是Transformer,已成为自然语言处理(NLP)的首选模型。主流的方法在大型文本语料库上进行预训练,然后在较小的特定任务数据集上进行微调。得益于Transformer的计算效率和可扩展性,研究者们成功训练了具有超过1000亿参数的大型模型。随着模型和数据集的不断增长,性能仍未显示出饱和迹象。
在计算机视觉领域,卷积架构仍然占据主导地位。受到NLP成功的启发,多项研究尝试将类似CNN的架构与自注意力结合,其中一些完全替代了卷积。后者模型在理论上效率高,但由于使用了特殊的注意力模式,尚未在现代硬件加速器上有效扩展。因此,在大规模图像识别中,经典的ResNet式架构仍然是最先进的技术。本文尝试将标准Transformer直接应用于图像,尽量少地进行修改。为此,本文将图像分割成块,并将这些块的线性嵌入序列作为输入提供给Transformer。图像块被视为与NLP应用中的token相同,以监督方式对模型进行图像分类训练。
当在中等规模的数据集(如ImageNet)上训练而不强化正则化时,这些模型的准确率略低于相当大小的ResNets几个百分点。这种看似令人沮丧的结果可能是预期的:Transformer缺乏CNN所固有的一些归纳偏差,如平移等变性和局部性,因此当在数据量不足时训练的情况下,它们的泛化能力并不佳。然而,如果在更大的数据集上进行训练,情况就会改变。本文发现大规模训练超过归纳偏差。当ViT在足够规模上预训练并转移到数据点较少的任务时,获得了优异的结果。当在公共的ImageNet-21k数据集或内部的JFT-300M数据集上预训练时,ViT在多个图像识别基准上达到或超过了最先进的水平。特别是,最佳模型在ImageNet上的准确率达到88.55%,在ImageNet-ReaL上为90.72%,在CIFAR-100上为94.55%,在VTAB的19个任务上为77.63%。
2 相关工作
略
3 方法
ViT在模型设计中尽可能遵循原始Transformer,这种设置的优点是可以方便的扩展到NLP Transformer架构及其高效实现。
3.1 视觉Transformer(ViT)
ViT如图1所示。标准Transformer输入是一维的token嵌入序列。为处理二维图像,作者将图像 <math xmlns="http://www.w3.org/1998/Math/MathML"> x ∈ R H × W × C x\in \mathbb{R}^{H×W×C} </math>x∈RH×W×C 重塑为一系列扁平化的2D图像块序列 <math xmlns="http://www.w3.org/1998/Math/MathML"> x p ∈ R N × ( P 2 × C ) \mathbf{x}_p \in \mathbb{R}^{N \times\left(P^2 \times C\right)} </math>xp∈RN×(P2×C), 其中(H, W)是原始图像的分辨率,C是通道数,(P, P)是每个图像块的分辨率, <math xmlns="http://www.w3.org/1998/Math/MathML"> N = H W P 2 N = \frac{HW}{P^2} </math>N=P2HW 是最终图像块的数量,也是Transformer的有效输入序列长度。 Transformer在其所有层中使用恒定的潜在向量维度D,因此作者将图像块扁平化并通过可训练的线性投影映射到D维(公式1)。这个投影的输出称为图像块嵌入(patch embeddings)。
类似于BERT的[CLS] token,ViT在嵌入图像块的序列前添加一个可学习的嵌入( <math xmlns="http://www.w3.org/1998/Math/MathML"> z 0 0 = x c l a s s z^0_0 = x_{class} </math>z00=xclass),其在Transformer编码器的输出( <math xmlns="http://www.w3.org/1998/Math/MathML"> z L 0 z_L^0 </math>zL0)处的状态作为图像表示y(公式4)。无论是在预训练还是微调期间,都会在 <math xmlns="http://www.w3.org/1998/Math/MathML"> z L 0 z_L^0 </math>zL0上附加一个分类头。预训练时分类头由一个含有单隐藏层的多层感知器(MLP)实现,而在微调时则由单层线性层实现。
位置嵌入被添加到图像块嵌入中以保留位置信息。ViT使用标准的可学习的一维位置嵌入,因为作者没有观察到使用更先进的二维感知位置嵌入带来显著的性能提升(附录D.4)。最终的嵌入向量序列作为编码器的输入。 Transformer编码器由多头自注意力层和MLP块交替组成(公式2,3)。每个块之前都应用层归一化(LN),每个块之后都有残差连接。MLP包含两层,使用GELU非线性激活函数。
归纳偏置
注意到,与CNNs相比,ViT具有更少的图像特定归纳偏置。在CNNs中,局部性、二维邻域结构和平移等变性贯穿整个模型的每一层。在ViT中,只有MLP层是局部的和平移等变的,而自注意力层则是全局的。二维邻域结构的使用非常有限:在模型开始时通过将图像切成块,在微调时调整不同分辨率图像的位置嵌入(如下所述)。除此之外,初始时的位置嵌入不携带有关图像块二维位置的信息,图像块间的所有空间关系都必须从头开始学习。
混合架构
作为原始图像块的替代,输入序列可以从CNN的特征图中形成。在这种混合模型中,图像块嵌入投影E(公式1)应用于从CNN特征图中提取的块。在特殊情况下,块可以具有1x1的空间大小,这意味着输入序列通过简单地扁平化特征图的空间维度并投影到Transformer维度来获得。分类输入嵌入和位置嵌入的添加如上所述。
3.2 微调和更高分辨率
通常在大型数据集上预训练ViT,并微调到(较小的)下游任务。为此,微调时移除预训练的预测头,并附加一个零初始化的D × K前馈层,其中K是下游类别的数量。通常在比预训练时更高的分辨率下微调是有益的。当输入更高分辨率的图像时,如果保持图像块大小不变,将导致更大的有效序列长度。ViT可以处理任意序列长度(受内存限制),但预训练的位置嵌入可能不再有意义。因此,作者根据它们在原始图像中的位置进行二维插值预训练位置嵌入。注意,这种分辨率调整和图像块提取是将关于图像二维结构的归纳偏置手动注入到ViT中的唯一机会。
4 实验
4.1 实验设置
这里首先介绍一下ViT的模型命名规则,主要涉及两个关键参数:模型的大小和输入图像块的大小。
- 模型大小:ViT的模型大小通常遵循BERT模型的命名约定,主要有三种类型:Base、Large、Huge
- 输入图像块的大小 :命名中的数字表示将输入图像划分为多大的块。图像块的大小直接影响模型的输入序列长度,因为每个图像块都被视为一个输入单元(类似于NLP中的"token")。数字越小,意味着图像块越小,从而导致序列长度增加,计算成本相应提高。例如:
- ViT-L/16:表示这是一个大型(Large)ViT模型,输入图像被划分为16x16像素的块。
- ViT-H/14:表示这是一个巨型(Huge)ViT模型,输入图像被划分为14x14像素的块。
实验部分主要探讨了ViT、ResNet以及混合模型在多个数据集上的表示学习能力。这些数据集包括ILSVRC-2012的ImageNet(1k类和1.3M图像)、其超集ImageNet-21k(21k类和14M图像)以及JFT(18k类和303M高分辨率图像)。研究中还去除了预训练数据集与下游任务测试集之间的重复数据。为了改进迁移学习效果,作者还对基线CNN模型进行了调整,如替换批归一化(Batch Normalization)层为组归一化(Group Normalization),并使用标准化卷积。
在训练和微调阶段,所有模型均使用Adam优化器,并采用Batch=4096和权重衰减0.1。此外,微调阶段使用带动量的SGD,特别地,对于一些模型在ImageNet上的微调,采用了更高的分辨率。评估模型性能时,使用了微调准确率和少数样本准确率两种指标,后者通过解决正则化最小二乘回归问题获得。这一设置部分为理解模型的性能及其在不同任务上的应用提供了详细的背景。
4.2 与SOTA方法对比
作者对比了最大的ViT模型(ViT-H/14和ViT-L/16)和当前文献中的最先进的CNN模型。具体对比了两个模型:一是使用大型ResNets进行监督迁移学习的Big Transfer (BiT),二是在去标签的ImageNet和JFT300M数据集上使用半监督学习方法训练的大型EfficientNet,即Noisy Student。当前,Noisy Student在ImageNet上表现优异,而BiT在其他数据集上表现突出。
实验显示,ViT模型在预训练效率和性能方面都表现卓越。较小的ViT-L/16模型在所有任务上都超越了BiT-L,并且在计算资源需求上大大低于后者。ViT-H/14模型在更具挑战性的数据集上,如ImageNet和CIFAR-100,表现更为出色,同时其预训练的计算需求也远低于之前的顶尖模型。研究也指出,预训练效率不仅受到模型架构的影响,还受到训练策略和优化器选择等因素的影响。
此外,ViT-L/16模型在公开的ImageNet-21k数据集上也展示了良好的性能,并且预训练所需的资源更少,可以在标准的云TPUv3(8核)上在大约30天内完成训练。这些结果在VTAB分类任务中得到了进一步的验证,ViT-H/14在自然和结构任务中优于其他方法,而在专业任务中与顶尖模型相当。这些发现强调了ViT模型作为图像识别领域的强大候选者的潜力。
4.3 预训练数据需求
作者探讨了ViT对预训练数据集大小的依赖性,并进行了两组实验以评估其表现。
首先,在不同规模的数据集上预训练ViT模型:从ImageNet到ImageNet-21k,再到庞大的JFT300M。结果(图3)显示,尽管进行了适度的正则化,ViT在最小的数据集ImageNet上的表现依旧不如在更大的数据集上。特别是在JFT300M上预训练时,大型ViT模型才能显示出其真正的潜力。
其次,在从9M到300M不等的数据集规模上训练模型,这些实验没有额外的正则化措施,以便准确评估模型的内在特性。分析(图4)表明,在小数据集上,ViT比计算成本相当的ResNets更容易过拟合。例如,ViT-B/32在9M数据集上表现不佳,但在90M以上的数据集上则表现出色。
这些发现提示,虽然卷积的归纳偏见对小数据集有利,但在处理大规模数据集时,直接从数据中学习模式可能更为有效。这些结果为未来的研究提供了新的方向,尤其是在探索ViT在少数样本学习方面的潜能。
4.4 规模研究
作者通过从JFT-300M数据集对不同规模的模型进行对比。本研究中,数据大小并未限制模型的性能,使得作者得以评估每个模型的性能与预训练成本之间的关系。这些模型包括7种ResNet配置、6种ViT配置,以及5种混合模型,其中混合模型的命名末尾数字代表在ResNet主干中的总降采样比率。
研究结果(图5)显示,ViT在性能与计算成本的权衡方面优于ResNet,平均而言,ViT在达到相似性能时所需的计算量约为ResNet的四分之一至二分之一。此外,尽管混合模型在较小的计算预算下略显优势,但在更大规模的模型中这种差异逐渐消失。这一发现略显意外,因为通常预期卷积的局部特征处理能在任何尺寸的ViT中提供辅助。此外,研究中未观察到视觉Transformer在测试范围内出现性能饱和现象,这表明未来扩展ViT规模是一个值得探索的方向。
4.5 审视ViT
为理解ViT是如何处理图像数据的,作者分析了其内部表示。ViT的第一层将展平的图像块线性投影到一个低维空间(公式1)。图7(左)显示了学习到的嵌入滤波器的主要成分。这些成分看起来像是对每个块内精细结构的低维表示的合理基函数。
在投影之后,学习到的位置嵌入被添加到图像块的表示中。图7(中)显示,模型学会在位置嵌入的相似性中编码图像内的距离,即相邻的图像块往往具有更相似的位置嵌入。此外,行列结构变得明显;同一行或列的块具有相似的嵌入。最后,对于更大的网格,有时可以看到正弦结构(附录D)。位置嵌入学会代表二维图像拓扑的事实解释了为什么手工制作的二维感知嵌入变体没有带来改进(附录D.4)。
自注意力机制允许视觉Transformer甚至在最低层也能整合整个图像的信息。作者研究了网络在多大程度上利用了这一能力。具体来说,根据注意力权重计算信息整合的平均图像空间距离(图7,右)。这种"注意力距离"类似于CNN中的感受野大小。
4.6 自监督
Transformer在自然语言处理任务上展现出了令人印象深刻的性能。然而,它们的成功不仅源于出色的可扩展性,还源于大规模的自监督预训练。作者也进行了一项初步探索,模仿BERT中使用的掩蔽语言建模任务,对掩蔽图像块进行预测以实施自监督。通过自监督预训练,较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率,比从头开始训练提高了2%,但仍比监督预训练低4%。附录B.1.2包含了更多细节。本文将探索对比预训练留给未来的工作。
5 总结
本文探索了将Transformer直接应用于图像识别的方法。与先前在计算机视觉中使用自注意力的工作不同,本文没有在架构中引入除初始图像块提取步骤外的图像特定归纳偏置。相反,将图像解释为一系列图像块,并使用在自然语言处理中应用的标准Transformer编码器对其进行处理。这种简单但可扩展的策略,在与大规模数据集预训练结合时,表现出意外的良好效果。因此,视觉Transformer在许多图像分类数据集上达到或超过了最先进的水平,同时预训练的成本相对较低。
尽管这些初步结果令人鼓舞,但仍有许多挑战存在。一个挑战是将ViT应用于其他计算机视觉任务,如检测和分割。本文的结果,结合Carion等人(2020年)的研究,显示了这种方法的潜力。另一个挑战是继续探索自监督预训练方法。初步实验表明自监督预训练有所改进,但自监督和大规模监督预训练之间仍有很大的差距。最后,进一步扩展ViT很可能会带来性能的提升。