2021年1月25日

[翻译]AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Abstract

尽管Transformer结构在自然语言处理领域已经成为主流,但在计算机视觉领域仍然受到限制。在视觉领域,注意力结构要么和卷积网络一起应用,要么在保证总结构不变的情况下用于替换卷网络中的某些组件。本文表明,这种对于CNN的依赖是没有必要的,直接在transformer结构中输入一系列图像块也可以达到非常好的效果。在大量图像上预训练后,在中小型数据集(ImageNet,CIFAR-100,VTAB…)上,Vision Transformer(ViT)仅需少量计算资源,即达到了卷积网络的SOTA。

Introductiion

自学习注意力结构,尤其是Transformer结构,被NLP任务选择。主导方法是在一个大型文本语料库上预训练,然后在较小的特定任务的数据集上进行微调。Transformer具有超强的计算效率和扩展性,可以训练前所未有的超过100B的参数,并且随着模型和数据集的增长,仍未出现饱和迹象。
在视觉领域,卷积仍然居于主导。收到NLP领域成功的影响,很多人尝试将自学习注意力机制和卷积结构组合,甚至替代。但由于这些专门的注意力机制尚未在硬件加速上得到支持,类似ResNet的结构仍然是主导。

受到NLP领域中Transformer的scaling(?)成功,我们在尽量最小的修改下,直接把标准Transformer网络应用到图像上。为了达到目的,我们把图像分为多个小块,并将每个小块当成线性embedding,作为Transormer的输入,其地位等同于NLP领域中的词向量,然后训练一个有监督的图像分类任务。

当对中等规模数据集(ImageNet)进行训练时,ViT比同等规模的ResNet第几个百分点,这令人沮丧的表明,Transformer缺少CNN固有的一些归纳偏差,例如平移不变性和局部性,因此在数据量不足的情况下不足以完美收敛。

但是如果在规模较大(14M-300M)的数据集上结果会发生改变,我们发现大规模训练胜过归纳偏差。Vit在足够数据上的预训练后,转移到数据点较少的任务时,可以获得出色的结果。当在ImageNet-21k数据集或室内JTF-300M数据集上预训练后,ViT在多个图像识别基准上胜过SOTA。最好的模型在ImageNet上达到88.55%,在ImageNet-Real上达到90.72%,在VTAB上达到77.63%。

Related Work

Transformer在机器翻译领域由Vaswani等人提出,是目前在很多NLP领域SOTA的方法。基于Large Transformer结构的模型通常在大型语料库上进行预训练,然后针对手头的工作进行微调,例如BERT使用去噪自监督预训练任务,GPT工作线使用语言模型作为其预训练任务。

在图像上自学习注意力机制的简单应用需要计算每个像素和其他像素的相关性,由于像素数量为二次方量级,所以无法使用真实像素大小。因此为了将Transformer应用到图像上,之前的工作尝试了几种方式:Parmar仅对每个点的局部临域使用注意力机制,而不是计算全局信息,这样的局部多头点积式注意力机制完全可以取代卷积;类似Sparse Transformer的工作使用了对全局注意力可缩放的近似值,才能应用在图像上;扩大注意力的另一种方法是将其应用在不同大小的特征块上。这些专门设计的注意力架构在视觉任务上都取得了不错的结果,但需要进行复杂的工程实现才能在硬件加速器上运行。

组合注意力机制和卷积也是一种方式,例如在图像分类任务中增强特征图,或在目标检测、视频处理等任务重对输出增加进一步的注意力机制处理。

我们尚不了解Transformer在全局范围内对全尺寸图像的关注。与我们的模型最接近的是iGPT,该模型在降低图像分辨率和色彩空间后将Transformers应用于图像像素。该模型以无监督方式作为生成模型进行训练,然后可以对所得表示进行微调或线性探查分类性能,从而在ImageNet上实现最好达到72%的准确性。

与标准ImageNet数据集相比,本文相比Imagenet数据集,增加了很多其他识别论文的数据集,使用额外的数据源可以在标准基准上获得SOTA的结果。我们关注到Sun研究了CNN性能如何随数据集大小而扩展,以及Kolesnikov
在大型数据集上对CNN转移学习进行实证探索,我们也将重点放在大型数据集上,但是训练的是Transformers,而不是先前工作中使用的基于ResNet的模型。

Method

在模型设计中,我们尽可能地遵循原始的Transformer。这种可扩展的NLP Transformer结构的优点是,容易实现并且容易应用。

Vision Transformer

如图1,标准Transformer结构接受1D token embedding序列作为输入,为了处理2D图像,我们将图像x(HxWxC)重塑为N个2D块p(PPxC),PP是每个图像块的分辨率,N=HW/PP表示块数以及Transformer的输入长度。Transformer使用D维固定的 latent vector 处理所有的输入,将所有的p压平并使用可学习线性工程映射成D维。我们将这个工程的输出称为patch embedding。

与BERT的token相似,我们在输入序列最初加入一个可学习的embedding,记为z0,它将在Transformer编码器输出的时候作为最终的图像表达。在预训练和微调任务中,都使用z0作为分类头。分类头的实现在预训练时是有一层隐藏层的MLP,在微调时是一个单线性层。

对于每个patch embedding,都要附加标准1D位置embedding来保持位置信息,因为在使用更复杂的2D位置信息后并没有得到明显受益(即生成两个1/2长度的embedding分别表示x、y位置编码,然后拼接)。将嵌入位置信息的序列作为编码器的输入。

Transformer encoder由多头注意力结构和MLP块组成,每个结构之前加入Layer Norm,之后加入残差连接,MLP是夹着非线性GELU的两层结构。

混合结构

在混合模型中,可以使用CNN的特征图作为Transformer的输入序列,特殊情况下,特征图的每块尺度为1×1,这意味着只需将特征图的空间尺寸展平,映射到Transformer需要的维度即可,然后加入分类头embedding和位置embedding。

Fine-tuning and higher resolution

我们在大数据集上预训练ViT,在小数据集上微调。这时,我们删掉预训练的预测头,增加一个初始化为0的DxK的前向层,K为输出类。与预训练相比,以更高的分辨率进行微调受益更大。当输入高分辨率图像时,保持patch尺度不变,那么会有更长的序列长度。Transformer可以接受动态长度(最大到内存限制),然而预训练时使用的位置embedding会失效。因此我们根据其在原图中的位置,对预训练的位置编码进行2D插值。

Experiments

我们评估了ResNet,ViT,混合结构的表现。为了了解每个模型的数据要求,我们在不同大小的数据集上进行了预训练,并评估了许多基准任务。在考虑上预训练的计算成本时,ViT模型表现出色,在大部分识别任务上达到SOTA。最后我们使用自监督进行了一个小实验,证明自监督的ViT潜力很大。

Setup

Datasets

预训练数据集

  • ImageNet 1k classes 1.3M images
  • ImageNet-21k 21k classes 14M images
  • JFT 18k classes 303M high-resolution images
    基准测试任务(去掉预训练中已出现的)
  • Real labels
  • Cifar-10/100
  • Oxford-IIIT Pets
  • Oxford Flowers-102
  • 19-task VTAB

Model Variants

Tia-t1.png

  • BiT(Big Transfer) Resnets152x4 4100M
  • NoisyStudent EfficientNet 480M

Comparison to SOTA

Pretrain on JFT
Tia-t2.png
用不同大小的预训练数据集的表现
Tia-f1.png
计算量与表现作图
Tia-f2.png
计算量与表现作图
Tia-f5.png
网络参数与结果表现
Tia-f3.png
网络视野
Tia-f4.png

Conclusion

展望未来:

  1. 将其应用到其他视觉任务中
  2. 寻找自监督预训练方式
  3. 进一步扩展提高性能
Share

You may also like...

发表评论

您的电子邮箱地址不会被公开。