An image is worth 16x16 words: Transformers for image recognition at scale
郑重声明:原文参见标题,如有侵权,请联系作者,将会撤销发布!
Published as a conference paper at ICLR 2021
ABSTRACT
虽然Transformer结构已成为自然语言处理任务的事实标准,但其在计算机视觉中的应用仍然有限。在视觉中,注意力要么与卷积网络结合使用,要么用于替换卷积网络的某些组件,同时保持其整体结构。我们表明,这种对CNN的依赖是不必要的,并且直接应用于图像块序列的纯Transformer可以在图像分类任务中表现得很好。当在大量数据上进行预训练并转移到多个中型或小型图像识别基准(ImageNet、CIFAR-100、VTAB等)时,与最先进的卷积网络相比,视觉Transformer (ViT)获得了优异的结果,同时需要更少的计算资源进行训练。1
1 微调代码和预训练模型可在https://github.com/google-research/vision_transformer中获取
1 INTRODUCTION
基于自注意的架构,特别是Transformers(Vaswani et al., 2017),已经成为自然语言处理(NLP)中的首选模型。主要的方法是在大型文本语料库上进行预训练,然后在较小的特定任务数据集上进行微调(Devlin et al., 2019)。由于Transformer的计算效率和可扩展性,训练具有超过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应用中的token(单词)相同。我们以监督的方式对模型进行图像分类训练。
当在中等大小的数据集(如ImageNet)上训练时,如果没有强大的正则化,这些模型的精度会比同等大小的ResNet低几个百分点。这种看似令人沮丧的结果可能是意料之中的:Transformer缺乏CNN固有的一些归纳偏差,如平移等变性和局部性,因此在数据量不足的情况下训练时不能很好地泛化。
然而,如果在较大的数据集(14M-300M图像)上训练模型,则图片会发生变化。我们发现大规模训练胜过归纳偏差。当以足够的规模进行预训练并迁移到数据点较少的任务中时,我们的视觉Transformer(ViT)会获得优异的结果。当在公共ImageNet-21k数据集或内部JFT-300M数据集上进行预训练时,ViT在多个图像识别基准上接近或优于现有技术。特别是,最佳模型在ImageNet上达到88.55%的准确率,在ImageNet-ReaL上达到90.72%,在CIFAR-100上达到94.55%,在由19个任务组成的VTAB套件上达到77.63%。
2 RELATED WORK
3 METHOD
在模型设计中,我们尽可能地遵循原始Transformer (Vaswani et al., 2017)。这种有意简化的设置的一个优点是,可扩展的NLP Transformer架构及其高效实现几乎可以开箱即用。
3.1 VISION TRANSFORMER (VIT)
该模型的概述如图1所示。标准Transformer接收token嵌入的1D序列作为输入。为了处理2D图像,我们将图像重塑为平坦的2D块(patch)序列
,其中(H, W)是原始图像的分辨率,C是通道数,(P, P)是每个图像块的分辨率,N=HW/P2是得到的块数,这也是Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的隐矢量大小D,因此我们将块压平,并使用可训练的线性投影映射到D维(公式1)。我们将此投影的输出称为块嵌入。
类似于BERT的[class] token,我们为嵌入块序列预先准备了可学习嵌入(),其在Transformer编码器的输出处的状态(
)用作图像表征y (公式4)。在预训练和微调期间,分类头都连接到
。分类头在预训练时由具有一个隐藏层的MLP实现,在微调时由单个线性层实现。
位置嵌入被添加到块嵌入以保留位置信息。我们使用标准的可学习1D位置嵌入,因为我们没有观察到使用更先进的2D感知位置嵌入带来的显著性能增益(附录D.4)。所得到的嵌入矢量序列用作编码器的输入。
Transformer编码器(Vaswani et al., 2017)由多头自注意力(MSA,见附录A)和MLP块(公式2, 3)的交替层组成。在每个块之前应用层归一化(LN),在每个块之后应用残差连接(Wang et al., 2019; Baevski & Auli, 2019)。
MLP包含两个具有GELU非线性的层。
感应偏差。我们注意到,视觉Transformer比CNN具有小得多的图像特异性感应偏差。在CNN中,局部性、二维邻域结构和平移同变性被融入整个模型的每一层。在ViT中,只有MLP层是局部且平移同变性的,而自注意层是全局的。二维邻域结构的使用非常谨慎:在模型开始时,通过将图像切割成块,并在微调时调整不同分辨率图像的位置嵌入(如下所述)。除此之外,在初始化时的位置嵌入不携带关于块的2D位置的信息,并且块之间的所有空间关系必须从头开始学习。
混合架构。作为原始图像块的替代方案,输入序列可以由CNN (LeCun et al., 1989)的特征图形成。在该混合模型中,将块嵌入投影E (公式1)应用于从CNN特征图提取的块。作为一种特殊情况,块可以具有1x1的空间大小,这意味着通过简单地将特征图的空间维度平坦化并投影到Transformer维度来获得输入序列。如上所述,添加了分类输入嵌入和位置嵌入。
3.2 FINE-TUNING AND HIGHER RESOLUTION
通常,我们在大型数据集上预训练ViT,并微调到(较小的)下游任务。为此,我们移除预训练的预测头,并添加零初始化的D×K前馈层,其中K是下游类的数量。以比预训练更高的分辨率进行微调通常是有益的(Touvron et al., 2019; Kolesnikov et al., 2020)。当馈送更高分辨率的图像时,我们保持块大小相同,这导致更大的有效序列长度。视觉Transformer可以处理任意序列长度(最高可达内存约束),然而,预先训练的位置嵌入可能不再有意义。因此,我们根据预训练的位置嵌入在原始图像中的位置对其进行2D插值。注意,这种分辨率调整和块提取是关于图像的2D结构的感应偏差被手动注入到视觉Transformer中的唯一点。
4 EXPERIMENTS
我们评估了ResNet、视觉Transformer (ViT)和混合体的表征学习能力。为了了解每个模型的数据需求,我们在不同大小的数据集上进行了预训练,并评估了许多基准任务。当考虑到预训练模型的计算成本时,ViT表现非常好,以较低的预训练成本在大多数识别基准上达到了最先进的水平。最后,我们使用自监督进行了一个小实验,并表明自监督ViT在未来充满希望。
4.1 SETUP
数据集。为了探索模型的可扩展性,我们使用了ILSVRC-2012 ImageNet数据集,该数据集具有1k个类和1.3M个图像(我们在下文中将其称为ImageNet),其超集ImageNet-21k具有21k个类和14M个图像(Deng et al., 2009),以及具有18k个类,303M个高分辨率图像的JFT (Sun et al., 2017)。我们根据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个训练示例。任务分为三组:Natural——任务(如上述),宠物、CIFAR等。Specialized——任务(如医疗和卫星图像)以及Structured——需要几何理解的任务(如定位)。
模型变体。如表1所示,我们基于用于BERT的ViT配置(Devlin et al., 2019)。“Base”和“Large”模型直接取自BERT,我们增加了更大的“Huge”模型。在下文中,我们使用简短的符号来表示模型大小和输入块大小:例如,ViT-L/16表示输入块大小为16×16的“Large”变体。请注意,Transformer的序列长度与块大小的平方成反比,因此具有较小块大小的模型在计算上更昂贵。
对于基线CNN,我们使用ResNet (He et al., 2016),但将批归一化层(Ioffe & Szegedy, 2015)替换为组归一化(Wu & He, 2018),并使用标准化卷积(Qiao et al., 2019)。这些修改改善了传输(Kolesnikov et al., 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在我们的设置中比用于ResNet的SGD表现得略好)。我们使用线性学习率预热和衰减,详见附录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)。
度量。我们通过小样本或微调精度在下游数据集上报告结果。微调精度捕捉每个模型在各自数据集上进行微调后的性能。通过解决将训练图像子集的(冻结的)表征映射到{−1, 1}K目标矢量的正则化最小二乘回归问题,可以获得小样本精度。这种公式使我们能够以闭合项恢复精确解。虽然我们主要关注微调性能,但有时我们会使用线性小样本精度进行快速实时评估,因为微调成本太高。
4.2 COMPARISON TO STATE OF THE ART
4.3 PRE-TRAINING DATA REQUIREMENTS
视觉Transformer在大型JFT-300M数据集上进行预训练时表现良好。与ResNet相比,视觉的归纳偏差更少,数据集的大小有多重要?我们进行了两个系列的实验。
首先,我们在越来越大的数据集上预训练ViT模型:ImageNet、ImageNet-21k和JFT300M。为了提高较小数据集的性能,我们优化了三个基本的正则化参数——权重衰减、随机失活和标签平滑。图3显示了微调到ImageNet后的结果(其他数据集的结果如表5所示)2。当在最小的数据集ImageNet上预训练时,尽管(适度)正则化,但与ViT-Base模型相比,ViT-Large模型表现不佳。通过ImageNet-21k的预训练,他们的表现是相似的。只有JFT-300M,我们才能看到更大模型的全部好处。图3还显示了性能由不同大小的BiT模型跨越的区域。BiT CNN在ImageNet上的表现优于ViT,但在更大的数据集中,ViT的表现更好。
其次,我们在9M、30M和90M的随机子集以及完整的JFT300M数据集上训练我们的模型。我们不对较小的子集执行额外的正则化,并对所有设置使用相同的超参数。通过这种方式,我们评估了固有的模型属性,而不是正则化的影响。然而,我们确实使用了早期停止,并报告了在训练过程中实现的最佳验证准确性。为了节省计算,我们报告了很少的线性精度,而不是完全的微调精度。图4包含了结果。在较小的数据集上,视觉Transformer比具有可比计算成本的ResNet过拟合更多。例如,ViT-B/32比ResNet50稍快;它在9M子集上表现得更差,但在90M+子集上表现更好。ResNet152x2和ViT-L/16也是如此。这一结果强化了一种直觉,即卷积归纳偏差对于较小的数据集是有用的,但对于较大的数据集,直接从数据中学习相关模式是足够的,甚至是有益的。
总体而言,ImageNet上的小样本结果(图4)以及VTAB上的低数据结果(表2)似乎有望实现非常低的数据传输。进一步分析ViT的小样本特性是未来工作的一个令人兴奋的方向。
2 请注意,ImageNet预训练的模型也经过了微调,但还是在ImageNet上。这是因为在微调过程中分辨率的提高也提高了性能。
4.4 SCALING STUDY
4.5 INSPECTING VISION TRANSFORMER
为了开始了解视觉Transformer是如何处理图像数据的,我们分析了它的内部表征。视觉Transformer的第一层将压扁的块线性投影到较低维空间中(公式1)。图7(左)显示了学到的嵌入滤波器的顶部主要组件。这些成分类似于每个块内精细结构的低维表征的合理基函数。
在投影之后,学到的位置嵌入被添加到块表征中。图7(中)显示,该模型学习在位置嵌入的相似性中对图像内的距离进行编码,即更接近的块往往具有更相似的位置嵌入。此外,出现行-列结构;同一行/列中的块具有相似的嵌入。最后,正弦结构有时对于较大的网格是明显的(附录D)。位置嵌入学会了表示2D图像拓扑,这解释了为什么手工制作的2D感知嵌入变体不能产生改进(附录D.4)。
自注意力允许ViT集成整个图像中的信息,即使是在最底层。我们调查网络在多大程度上利用了这种能力。具体来说,我们根据注意力权重计算图像空间中信息整合的平均距离(图7,右)。这种“注意距离”类似于CNN中的感受野大小。我们发现,最底层的一些头已经处理了大部分图像,这表明模型确实使用了全局集成信息的能力。其他注意力头在低层中的注意力距离一直很小。这种高度局部的注意力在Transformer之前应用ResNet的混合模型中不太明显(图7,右),这表明它可能与CNN中的早期卷积层具有类似的功能。此外,注意力距离随着网络深度的增加而增加。在全局范围内,我们发现该模型关注与分类在语义上相关的图像区域(图6)。
4.6 SELF-SUPERVISION
Transformer在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; Henaff et al., 2020)。
5 CONCLUSION
我们已经探索了Transformer在图像识别中的直接应用。与先前在计算机视觉中使用自注意力的工作不同,除了最初的块提取步骤外,我们没有在架构中引入图像特定的归纳偏差。相反,我们将图像解释为一系列块,并通过NLP中使用的标准Transformer编码器进行处理。这种简单但可扩展的策略与大型数据集上的预训练相结合,效果出奇地好。因此,视觉Transformer在许多图像分类数据集上达到或超过了现有技术,同时预训练相对便宜。
虽然这些初步成果令人鼓舞,但仍存在许多挑战。一种是将ViT应用于其他计算机视觉任务,如检测和分割。我们的结果,加上Carion等人(2020)的结果表明了这种方法的前景。另一个挑战是继续探索自监督的预训练方法。我们的初步实验表明,自监督预训练有所改进,但与大规模监督预训练相比仍有很大差距。最后,ViT的进一步扩展可能会提高性能。
APPENDIX
A MULTIHEAD SELF-ATTENTION
B EXPERIMENT DETAILS
B.1 TRAINING
表3总结了我们针对不同模型的训练设置。我们发现,当在ImageNet上从头开始训练模型时,强正则化是关键。当使用时,Dropout应用于除qkv投影外的每个密集层之后,并直接应用于添加位置嵌入到块嵌入之后。混合模型的训练与ViT模型完全相同。最后,所有的训练都是在224分辨率上完成的。
B.1.1 FINE-TUNING
我们使用动量为0.9的SGD对所有ViT模型进行微调。我们对学习率进行了小型网格搜索,见表4中的学习率范围。为此,我们使用来自训练集的小型子分割(10%的Pets and Flowers,2%的CIFAR,1%的ImageNet)作为开发集,并对剩余数据进行训练。对于最终结果,我们对整个训练集进行训练,并对各自的测试数据进行评估。对于微调ResNet和混合模型,我们使用完全相同的设置,唯一的例外是ImageNet,我们在学习率扫描中添加了另一个值0.06。此外,对于ResNet,我们还运行了Kolesnikov等人(2020)的设置,并在这次运行和我们的横扫中选择最佳结果。最后,如果没有另行提及,所有微调实验都以384分辨率运行(以与训练不同的分辨率运行微调是常见的做法(Kolesnikov et al., 2020))。
当将ViT模型转移到另一个数据集时,我们移除整个头(两个线性层),并用一个零初始化的线性层替换它,该线性层输出目标数据集所需的类数。我们发现这比简单地重新初始化最后一层要稳健一些。
对于VTAB,我们遵循Kolesnikov等人(2020)的协议,并对所有任务使用相同的超参数设置。我们使用0.01的学习率,并训练2500步(表4)。我们通过对两个学习率和两个时间表进行小范围扫描,并在200个示例验证集中选择VTAB得分最高的设置,来选择该设置。我们遵循Kolesnikov等人(2020)中使用的预处理,只是我们不使用特定于任务的输入分辨率。相反,我们发现视觉Transformer从所有任务的高分辨率(384×384)中获益最多。
B.1.2 SELF-SUPERVISION
C ADDITIONAL RESULTS
D ADDITIONAL ANALYSES
D.1 SGD VS. ADAM FOR RESNETS
D.2 TRANSFORMER SHAPE
我们对Transformer架构的不同维度进行了消融,以找出哪些最适合扩展到非常大的模型。图8显示了不同配置下ImageNet上的5-样本性能。所有配置都基于具有8层的ViT模型,D=1024,DMLP=2048,块大小为32,所有线的交点。我们可以看到,缩放深度会带来最大的改进,这些改进在64层之前都是清晰可见的。然而,在16层之后已经可以看到收益递减。有趣的是,缩放网络的宽度似乎会导致最小的变化。在不引入参数的情况下,减小块大小并因此增加有效序列长度显示出令人惊讶的稳健改进。这些发现表明,计算可能比参数的数量更能预测性能,而且缩放应该强调深度而不是宽度(如果有的话)。总的来说,我们发现按比例缩放所有维度会带来稳健的改进。
D.3 HEAD TYPE AND CLASS TOKEN
为了尽可能接近原始Transformer模型(类似于BERT),我们使用了一个额外的[class] token,该token被用作图像表征(位于最前)。然后,通过小型多层感知器(MLP)将该token的输出转换为类预测,其中tanh作为单个隐藏层中的非线性。
这个设计是从文本的Transformer模型中继承的,我们在整个主论文中都使用它。最初尝试只使用图像块嵌入,全局平均池化(GAP),然后使用线性分类器——就像ResNet的最终特征图一样——效果非常差。然而,我们发现这既不是由于额外的token,也不是由于GAP操作。相反,性能上的差异完全可以用不同学习率的要求来解释,见图9。
D.4 POSITIONAL EMBEDDING
D.5 EMPIRICAL COMPUTATIONAL COSTS
D.6 AXIAL ATTENTION
D.7 ATTENTION DISTANCE
D.8 ATTENTION MAPS
D.9 OBJECTNET RESULTS
D.10 VTAB BREAKDOWN