【ViT】【ICLR 2021】用Transformer做大规模图像识别

收录在ICLR2021的《AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》《一张图片等价于多个16X16的单词：用Transformer做大规模图像识别》
原文连接：https://arxiv.org/pdf/2010.11929.pdf?ref=https://githubhelp.com
代码连接：https://github.com/google-research/vision_transformer
#CV #ICLR2021 #Vision Transformer

摘要

Transformer架构已经成为自然语言处理任务中的一个标准，尤其是BERT和GPT系列模型在复杂的自然语言处理任务中取得非常惊艳的效果。而利用Transformer解决视觉任务时，自注意力机制或是和卷积神经网络相结合，或是替代其中某些模块，原有的卷积神经网络结构并没有改变，使得Transformer在视觉领域的应用十分有限。本文作者提出了Vision Transformer（ViT）模型，表明上述对于卷积神经网络的依赖是不必要的，可以利用一个基本的Transformer直接作用于一系列图像块，在图片分类任务中可以取得良好的效果。在大规模数据集上训练后，迁移到中小型数据集进行使用，Vision Transformer可以获得与最好的卷积神经网络相媲美的结果，同时需要的训练资源更少。

1介绍

1.1Transformer应用于CV的难点

在计算机视觉领域，要应用Transformer的难点在于如何将 $2d$ 的图片转化为 $1d$ 的序列或集合。最直观的方法便是将图片的每个像素点当做输入的元素，直接将 $2d$ 的图片拉直，这样做的缺陷在于复杂度过高。如图1所示，Transformer中最主要的是自注意力操作，即输入的 $n$ 个元素两两进行互动，计算得到自注意力的图 $Attention$ ，再通过加权平均得到最后的输出。因为两两之间都要进行计算，因此复杂度为 $O(n^2)$ ，常见的BERT模型中支持的序列长度 $n=512$ ；在基本的图片分类任务中，按上述思想得到的图片像素序列长度 $n$ 为 $224\ast 224=50176$ ，是BERT中序列长度 $n$ 的近百倍，用于分割和检测任务的图片尺寸更大，计算复杂度相应更高。

图1 Transformer示意图

1.2相关尝试

为了解决因输入序列过长导致计算复杂度过高的问题，一些工作尝试从以下角度降低序列长度

• 利用卷积神经网络中的特征图取代原始图片作为Transformer的输入，使得输入序列长度大大减少，但还是和卷积神经网络联系紧密
• 利用轴自注意力或孤立自注意力此类特殊设计的自注意力机制降低输入序列长度，但是由于结构的设计特殊，无法训练一个大模型

1.3本文贡献

作者受到Transformer在NLP领域的可扩展性的启发，在尽量做少的修改的前提下，即不针对视觉任务做特定的改变，提出了Vision Transformer（ViT）模型，应用标准的Transformer直接作用于图片，在图片分类任务中取得了很好的效果。

2Vision Transformer（ViT）

2.1模型总览

ViT模型结构简洁，包含前处理、Transformer Encoder和分类三部分内容。通过前处理阶段将一张图片处理为适合Transformer Encoder输入的形式，其中借鉴了Transformer在自然语言处理领域中常用的Position Embedding和cls分类头解决图像块序列位置信息嵌入问题和Transformer Encoder多个输出的选择问题，同时解决了因输入序列长度太长导致的高计算复杂度问题，Transformer Encoder由循环L次、不针对视觉任务做特定改变的标准Transformer block构成，其输出接一个常用MLP分类头完成图片分类任务。

图2 ViT模型示意图

2.2前处理

前处理阶段将输入的图片 $x$ 处理为适合Transformer Encoder的输入 $Z_0$，如图2中左下侧红色边框部分所示。

图片预处理 对输入图片 $x\in {\mathbb{R}}^{H\ast W\ast C}$ 进行预处理，将 $x\in {\mathbb{R}}^{H\ast W\ast C}$ 划分为 $x_p\in {\mathbb{R}}^{N\ast (P^{2}C)}$ 的图像块序列， (H,W)为原始图片的分辨率，C 为通道数，(P,P)为划分后的图像块分辨率，那么划分后的图像块数目为 $N=\frac{HW}{P^2}$。接下来将划分后的图像块展平为向量，每个向量大小为 $P\ast P\ast C$ ，此时输入图片变化为 $N\ast P^2\ast C$ 向量，表示为 $[x_p^1;x_p^2;...;x_p^N]$ 。在图片预处理阶段解决了输入序列长度太长导致的高计算复杂度问题，使输入图像变化为长度为N、Transformer可处理的向量序列。

Patch Embedding 对图片预处理后得到的每个向量做线性变换，通过全连接层E完成Patch Embedding，结果表示为 $[x_p^{1}E;x_p^{2}E;...;x_p^{N}E]$ ，其中全连接层E的压缩维度为D，D的大小选择取决于Transformer模型的大小。

Position Embedding 划分后的图像块序列是有顺序的，顺序的颠倒意味着输入图片的改变，而自注意力操作是两两之间进行交互，本身并没有顺序的区别，因此借鉴NLP中的思想，在Patch Embedding的过程中加入Position Embedding，用 $E_{pos}$ 表示。加入位置编码后，整体的特征既包含了原本有的图像信息，同时包含了各个图像块所在位置信息。作者尝试了 1d、2d和相对位置编码三种编码方式，实质上2d编码方式更能体现原始图片的结构信息，但实验表明不同的编码方式在本实验中对结果的影响很小，原因在于划分后的图像块序列信息是相对容易探寻到的。

Extra learnable [class] embedding 借鉴BERT中的cls分类字符，用 $\ast$ 表示，其位置编码固定为 $0$ ，表示为 $x_{class}$ 。由于每个特征之间都要做交互，认为它可以从别的特征中学到相应信息，因此最终在Transformer Encoder的若干个输出中，选择它的输出做最后的分类。同时也可以在前处理阶段不加入分类字符cls，借鉴卷积神经网络中对特征图进行全局平均池化的操作，通过对Transformer Encoder的多个输出进行全局平均池化得到特征向量，进而进行分类，经消融实验证明这两种方法可以取得相差不大的结果，从另一角度说明了Transformer结构的有效性。本文中为了对Transformer进行尽可能少的改动而采用cls分类字符。完成前处理步骤后，得到Transformer Encoder的初始输入，表示为：
$$Z_0=[x_{class};x_p^{1}E;x_p^{2}E;...;x_p^{N}E]+E_{pos}$$

2.3Transformer Encoder

Transformer Encoder由循环L次的标准Transformer block构成，一个Transformer block的结构如图二中右侧绿色边框部分所示。

每次循环分两阶段进行，第一阶段中输入 $Z_{l-1}$依次通过线性Norm层和多头自注意力操作，分别用LN和MSA表示，然后通过残差连接得到第一阶段的输出 $Z_l^{\prime }$ ，第一阶段过程表示为：
$$Z_l^{\prime }=MSA(LN(Z_{l-1}))+Z_{l-1}$$
第二阶段将 $Z_l^{\prime }$ 作为输入，将第一阶段的多头自注意力操作替换为多层感知器MLP，此外与第一阶段步骤相同，两阶段完成后得到每一次Transformer block循环的结果 $Z_l$ ，第二阶段过程表示为：
$$Z_l=MLP(LN(Z_l^{\prime }))+Z_l^{\prime }$$
经过L次循环后，取第L次循环输出的第一个位置 $Z_L^0$ 作为Transformer Encoder的输出，即就是上文提到的 $x_{class}$ 所对应的输出。

2.4完成分类任务

将Transformer Encoder的输出 $Z_L^0$ 当做整体输入图像 $x$ 的特征，通过常用的MLP分类头完成分类任务 $Y$，表示为：
$$Y=LN（Z_L^0）$$

2.5ViT模型的形式化描述

整理2.2-2.5所描述步骤，得到ViT模型的形式化描述

3实验

3.1数据集

• ImageNet/ImageNet-1k：包含1,000个类
• ImageNet 21k ——21,000个类别 14,000,000张图片
• JFT：18,000个类别三亿张图片
• CIFAR-10/100
• Oxford-IIIT Pets
• Oxford Flowers-102

3.2实验结果

作者在图3展示了在经过大数据集训练后的效果，可以看到ViT-H/14模型全面优于别的模型。BiT-L是一个大的卷积模型，和Nosiy Student模型同样是Google的工作。Nosiy Student是此前在ImageNet数据集上效果最好的模型。从训练时间的角度来看，ViT-H/14模型的训练时间是2500天TPU v3天数，BiT-L模型需要9900天，Nosiy Student模型花费时间更长。

图3：实验结果1

究竟需要多大的数据集才可以训练好ViT模型呢？图4中灰色折线所包含的区域表示现有卷积神经网络可以达到的效果，不同色彩和大小的圆形表示不同大小的ViT模型。在ImageNet和ImageNet-21k数据集上，ViT的效果其实是没有卷积神经网络好的，因为Transformer都是从两两的交互中去学习信息的，没有卷积神经网络中的归纳偏置（局部性和平移等变性）。在更大的数据集JFT上，ViT的效果是全面优于卷积神经网络的，即便是最小的ViT也比卷积神经网络的效果要好一些，而且没有饱和的趋势，意味着如果数据集一直增大，模型效果会越来越好。

图四：实验结果2

图五中横坐标表示网络深度，纵坐标表示平均注意力距离。自注意力机制使得ViT模型在刚开始就能注意到全局的信息，随着网络深度的增加，逐渐学习到全局的语义信息。

图五：实验结果3

总结

ViT挑战了自从2012年Alexnet提出以来，卷积神经网络在计算机视觉领域的绝对统治地位，表明如果在足够大的数据集上做训练，那么不需要卷积神经网络，直接用标准的Transformer也能很好的处理视觉任务。更重要的是拓宽了研究思路，从任务角度来说，本文利用ViT只做了分类任务，可以继续用ViT解决检测和分割任务；从改变模型结构的角度来说，既可以改变特征的嵌入方式，也可以改变Transformer block的结构；从训练方式来说，可以继续沿用有监督的训练方式，也可以尝试自监督的训练方式。近期热门的MAE和Swin Transformer基于ViT模型，针对上述研究角度做了不同的改进实现，取得了良好的效果，进一步说明了Transformer应用在视觉任务领域的广泛性和可行性。