Vision Transformer 这两年
作者｜Maximilian Schambach
OneFlow 编译
翻译｜胡燕君、杨婷
在 NLP 领域取得巨大成功后，Transformer 架构在计算机视觉方面的作用日渐凸显，成为越来越普遍的 CV 工具。自 2020 年 10 月 Vision Transformer 模型推出以来，人们开始高度关注 Transformer 模型在计算机视觉上的应用。
图 1：各类 Vision Transformer 模型的推出时间（此处以论文在 arXiv 平台上的发表时间为准）
恰逢 Vision Transformer 推出两周年之际，借此机会我们对其稍作介绍，并讨论这两年来发展出的多种 Vision Transformer 模型变体以及 Transformer 在计算机视觉应用方面面临的各种挑战。本文由 OneFlow 社区编译。
1、自注意力机制和 Transformer 架构
从 NLP 说起，2017 年，Attention is all you need 一文提出了 Transformer 架构(Vaswani et al. 2017)。Transformer 架构的本质是一个序列到序列模型：输入的是一种称为 token 的序列，token 在 NLP 中是指对句子的数学表示。将句子转化为序列，要将句子中的每个词（或子词）映射为一个向量表示，该向量表示称为对应词的 embedding（嵌入）。
生成的 embedding token 序列由 Transformer 编码器在几个自注意力层和全连接层中处理，然后输出一个与输入长度相同的高级 token 表示序列。翻译等任务需要使用完整的输出序列，而分类等任务则只需用到单个表示。在分类任务中，要为整个句子或整个段落生成单个表示用于分类，通常会为序列附加一个特殊的可学习嵌入向量 Class_Token（[CLS] token）。可以将[CLS] token 的表示传入 classfier head（分类器头）。
下图是 Transformer 架构的总览，图左展示了 Transformer 编码器。自注意力机制是 Transformer 架构的关键。自注意力层中的所有中间 token 表示会通过计算成对相似性（pairwise similarities）进行互动，这些成对相似性会成为上一层相应 token 表示的权重。
图 2：Transformer 编码器和自注意力机制（图片来自原论文)
为了计算成对相似性，通常会通过简单可学习的线性层从 embedding 中计算 query (Q)、key (K)和 value (V)向量，然后所有 K 和 Q 之间进行双向 softmax 自注意力计算，获得各自的相似性，然后将相似性与 V 相乘。也就是说，在自注意力层中，每个 token 都通过各自 K 和 Q 的简单点积与其他 token 进行比较，随后再运行(scaled) softmax 操作。
如图 2 右边所示，为了能够在每个阶段学习不同形式的自注意力，Transformer 架构使用了多个 self-attention head（自注意力头）。具体来说，每个 head 都映射了来自输入表示的不同 Q、K、V 向量。例如，在 Transformer 的某个固定级别，不同的 attention head 可以关注 token 之间的短、长距离或语义和句法关系。不同的 attention head 的输出由线性层进行连接和处理，以便重新获得通过 skip connection 传递的输入表示的维度。不熟悉 Transformers 基本概念的读者可以参考原论文获取更多详细信息（https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html）。
2、Vision Transformer 模型
从 Transformer 架构被引进 NLP 到它在计算机视觉领域（ImageNet 图像识别任务）取得 SOTA 性能，中间隔了 3 年，这和深度学习研究的迅猛进展相比显得较为缓慢。但从抽象的角度讲，这并不令人意外。毕竟 Transformer 是序列模型，无法直接兼容图像和视频等高维且近似连续的输入数据类型。与序列不同，图像和视频可被视为从底层连续信号中离散采样的数据。
此外，Transformer 模型没有局部性（locality）和平移不变性（translational equivariance）等归纳偏置（inductive bias），但归纳偏置在处理图像等数据时十分有用。要让 Transformer 模型习得这些属性，需要使用大量的训练数据。相反，卷积神经网络（CNN）实行归纳偏置，CNN 过去十年在视觉领域取得的成功很大程度归功于此。
不过，也有观点认为，没有归纳偏置可以让 Transformer 架构更具通用性。例如，Transformer 模型的浅层即可提取全局信息，但 CNN 由于感受野（receptive field）较小，只能在神经网络的深层提取全局信息。
从实践角度讲，Transformer 模型在计算机视觉领域应用进展缓慢的另一个原因是：Transformer 模型的输入序列长度缺乏可扩展性。由于标准的双向自注意力机制中所有 token 表示都是成对比较的，所以标准的 Transformer 模型中，随着 input token 序列变长，计算所需时间和内存会呈二次方增长。因此，即使是分辨率只有 256 × 256 的小型图像，每个像素点为一个 input token，总共就形成含 65,000 个 token 的超长序列，导致操作可行性极低。
为了解决上述难题，Dosovitskiy 等人提出了 Vision Transformer 模型（ViT，图 3），该模型使用的是一种“简单粗暴”的方法：
首先，一张大小为 X × X × C 的输入图像被分成 M × M 个图像块（patch），每个图像块为固定大小 P × P（为了简化问题，此处假定输入图像都是正方形）。每个图像块（P × P × C）被展平（flatten）为一个形状为 P ⋅ P ⋅ C 的向量，然后运用可训练的线性层将该向量映射为一个维度为 D 的 embedding。展平得到的向量为图像块 embedding，即 input token，其中包含了每个图像块的信息。
最后，一个 M × M 的图像块 embedding 阵列被展平为一个长度为 N = M ⋅ M 的 token 序列，然后喂入 Transformer 模型中。可以将分块、展平和线性映射操作视为一个如下的 2D 卷积操作：使用 D 个形状为 P × P × C 的卷积核（kernel），步长（stride）为(P, P)，无填充（padding）。
图 3：ViT 模型架构。仅[CLS] token 的输出表示被用于分类任务和有监督训练
（图源：https://openreview.net/forum?id=YicbFdNTTy）
与 NLP 领域的标准 Transformer 模型相同，embedding token 和附加的[CLS] token 一起传入 Transformer 编码器，编码器（图 3 右）包含几个多头自注意力机制（multi-headed self-attention）与 MLP 层。标准的 ViT 中，潜在的维度 D 从头到尾都是固定的。[CLS] token 的输出被输入到 classification head 中，[CLS] token 的输出表示包含输入图像的所有必要信息，即[CLS] token 的输出表示被用作潜在图像表征，与 CNN 自编码器中的瓶颈表征（bottleneck representation）相似。[CLS] token 在有监督分类任务中训练时需要与高度依赖图像的 token embedding 交互。
标准的 ViT (ViT-B)中，图像块大小为 P = 16，因此一张大小为 64 × 64 × 3 的 RGB 图像会被分为 16 个图像块。每个图像块被展平为一个长度为 16 ⋅ 16 ⋅ 3 = 768 的向量，向量被线性层映射为一个维度 D = 768 的向量。长度为 1 + 16 的 input token 序列被传入包含 12 层组件的编码器中进行处理，每个自注意力模块有 12 个 head。ViT 模型中约有 8600 万参数，更大型的 ViT 模型变体，如 ViT-L 和 ViT-H 分别有 3 亿和 6 亿参数。基础 ViT 模型的大小相当于标准 ResNet-152 模型（6000 万参数），而 SOTA CNN 模型的大小则相当于 ViT-L 和 ViT-H。Dosovitskiy 等人在论文中将 ViT 与 ResNet-152x4 和 Efficientnet-L2 这两种 CNN 模型对比，前者有 9 亿参数，后者有 5 亿参数。
总体而言，ViT 模型，特别是 ViT-L 和 ViT-H，在 JFT-300M 大型数据集上预训练，然后迁移到 ImageNet 和 CIFAR 等中小型数据集上，都取得了 SOTA 性能表现。与基于 CNN 的视觉领域大型主干网络（backbone）相反，ViT 需要的训练资源较少。不过，即便是相对较少的训练资源，中小型研究机构也难以负担。此外，虽然上述方法在原理上能使 Transformer 模型更好地应用于视觉领域，但还存在其他难点，ViT 原始论文和后续论文也提出了相应的解决方法。
3、Transformer 模型在视觉领域的应用难点
第一个难点是，Transformer 模型无法区分 input token 的排列位置变化。也就是说，Transformer 模型不能理解 input token 之间的位置关系。在 NLP 任务中，token 在序列中的 1D 位置信息非常重要，因为位置信息代表了语序，对此，Transformer 的解决方法是对 token embedding 进行位置编码。
将位置信息编码到一个 token embedding，即为每个 token 的位置信息单独生成一个维度为 D 的向量，称为 positional embedding，然后将它添加到 token embedding 中，再传入 Transformer 模型。Positional embedding 可能是可学习的，也可能是固定的，它的生成方式是根据相应 token 的位置在不同频率的正弦曲线上的不同位置采点(Vaswani et al. 2017)。
Positional embedding 也称为傅里叶特征。通过上述方法，可以为任意长度的序列生成 positional embedding，因此，Transformer 模型可以处理任意长度的句子。
然而，ViT 模型更为棘手，因为它需要恢复图像块的 2D 位置信息——即该图像块截取自原图像的哪一个区域。要做到这一点，ViT 对每 N 个 input token 使用可学习的 1D positional embedding，由 N 个可学习的维度为 D 的向量组成。训练时，对 positional embedding 没有显式监督，也就是说，ViT 并不能理解 input token 之间的实际 2D 关系。
不过，Dosovitzky 等人通过计算每对 positional embedding 之间的相似性后发现，ViT 习得的 positional embedding 确实包含相应图像块的 2D 位置信息。这一重大发现证明了 Transformer 架构的强大性能和通用性。
然而，这种方法需要在训练中将输入序列的长度 N 调整为固定值。图像块的大小也需要固定（否则就会影响 linear embedding 层），这意味着要将输入图像调整为固定分辨率才能用于训练，因为小批次训练要求每批次中所有输入图像的分辨率相同，在 CNN 模型中则不需要这样做。在 ViT 模型中，推理时也需要调整图像块大小，即调整输入图像的分辨率，因为分辨率过大会导致 input token 序列过长。
如图 4 中的例子，图像块大小依然是前面提到的 16 × 16，那么一张大小为 128 × 128 的图像就会被分为 64 个图像块。前面提到，训练时一张大小为 64 × 64 的图像被分成 16 个图像块，那么 ViT 模型就只习得 16 个 positional embedding，这时模型需要将这 16 个 positional embedding 泛化从而适应不同的分辨率，同时保证它们所代表的 2D 位置信息依然完整。
好在，代表同一 2D 位置关系的各个 positional embedding 具有相似性，因此可以采用 2D 插值（interpolation）将它们泛化，以适应任意分辨率，同时保留所习得的位置信息（图 4）。这一方法由 Dosovitskiy 等人提出，它看似可行，但原始论文并没有对这种方法进行量化估计。原始论文中，所有用于训练的图像被降低至常用的 224 × 224 分辨率，用于微调（fine-tuning）的图像则调至更高分辨率。
图 4：将 positional embedding 从 4 × 4 个图像块重采样为 8 × 8 个图像块
其次，这种方法会导致每个图像块内像素点之间的位置关系丢失，还会导致密集预测（dense prediction）任务只能在图像块层次完成。例如语义分割中的逐像素点分类任务。这就要求我们在表现力和吞吐量之间做出取舍：图像块尺寸较小（最小尺寸为 1 × 1 px）可以更大程度地保留原图像的底层空间关系，获得密集的潜在表征（latent representation），但会导致吞吐量变小，所需内存空间也会大得难以实现。
为了解决这个问题，一些早期论文提出将注意力机制限制在本地像素邻域中（http://proceedings.mlr.press/v80/parmar18a.html），或者采用下采样（downsampling），将输入图像的分辨率调整得非常小（http://proceedings.mlr.press/v119/chen20s.html；https://openreview.net/forum?id=HJlnC1rKPB）。如果图像块大小达到最高限制，即单个图像块即包含整张原图，则没有可用的位置信息，这时 token embedding 层就会成为模型架构中的瓶颈，严重影响模型性能。
同样，各图像块边缘的邻近像素点之间的位置关系也会丢失。分块操作不受图像内容的影响，也就是说，图像的分块方式是随意的，有可能导致分出的每个图像块中各个像素点之间的情境关联并不大，进而导致从输入图像中习得的重要特征无法直接应用到图像块上，需要通过相应的 token embedding 进行重组。这就增加了下游任务（downstream task）的难度，也对 patch embedding 层提出了更高要求，还导致难以实现平移不变性。
最后，如前所述，标准的 Transformer 架构中，随着输入序列变长，运行所耗费的时间和内存空间都会呈二次方增长。因此，ViT 模型不适合处理高分辨率图像和其他高维数据，也不宜将图像块设置得太小。为了解决以上难题，ViT 模型首度公开不久后就出现了许多变体。
4、Vision Transformer 模型发展现状
接下来我们将介绍一些 ViT 模型变体，它们进一步推动了计算机视觉的 SOTA 性能，并解决了部分上述提到计算机视觉应用难题。
DeiT 模型（Data-Efficient Image Transformer）
Touvron 等人通过数据增强（data augmentation）等精细训练策略，在不依靠大型专有数据集的情况下实现了出色的模型性能（依然以 ImageNet 图像分类任务为衡量标准），这一点与原始 ViT 模型相反，ViT 模型是在 Google 的 JFT-300M 闭源数据集上训练的。
Touvron 等人还使用一种专门针对 Transformer 模型的蒸馏（distillation）技术进一步改进 ViT 模型。将一个性能强大，但可能庞大而难以训练的 teacher 模型“蒸馏”为一个基于 Transformer 架构的 student 模型，做法和 Hinton 等人提出的”知识蒸馏“相似（Hinton et al. 2014）。这种模型称为 DeiT 模型（Data-Efficient Image Transformer）。
与图 5 所示相似，distillation token 的输出表示被输入到额外的 classification head 中用于预测 teacher 模型的输出标签。[CLS] token 按照惯常做法被传入（另一个）classifier head 获得真实分类标签。总损失函数是使用真实标签计算出的常规交叉熵损失的加权平均值和基于 distillation classifier head 的 logit 输出与 teacher 模型 logit 输出的 KL 散度(Kullback Leibler divergence)的损失项。这样，Transformer 中，student 模型就可以借助 teacher 模型来提升自身的训练速度和质量。
图 5: DeiT 的特征蒸馏（图片来自相应论文）
此外，更强的数据增强可能会导致图片不符合其真实标签，例如当裁剪出的图像块没有包含原图标签所指物体（图 6）。为了缓解这个问题，可以使用一个不仅基于真实标签，而且基于 teacher 模型所预测标签的损失。此外，一些图像的分类标签可能不够明确，因为尽管图像可能包含多种物体，但每张图像只关联一个标签。
他们使用强大的 CNN 分类器显著提高了 student 模型的基准表现，这可能是利用了 teacher 模型中包含的归纳偏置，采用这种方式可以在不需要太多的训练数据的情况下从头开始训练 Transformer 模型。事实上，蒸馏得出的模型性能优于 teacher 模型，并在 ImageNet 图像分类任务上取得了 SOTA 性能，缩小了 Transformer 模型和 CNN 模型之间的差距。
DeiT 模型加快了训练速度，用 ImageNet 训练 DeiT 只需 4 块 GPU 训练 3 天。虽然它的性能不如那些在 JFT-300M 等更大型数据集上预训练的 ViT 模型，但它的计算成本要小得多。
图 6：来自 ImageNet 数据集的示例图像（左），对其进行裁剪会产生不同的标签（右）
Swin Transformer 模型 （Shifted Windows Transformer）
为了解决 ViT 模型中图像块边缘空间信息部分丢失和计算复杂度二次方增长的问题，Liu 等人引入了一种基于分层特征图和移位窗口（shifted window, Swin）的 ViT 模型（Liu et al. 2021）。Swin Transformer 有两个核心贡献，如图 7 所示。
图 7：Swin 中的分级特征图（上）与后续神经网络层窗口移位（下）（图源：https://openaccess.thecvf.com/content/ICCV2021/html/Liu_Swin_Transformer_Hierarchical_Vision_Transformer_Using_Shifted_Windows_ICCV_2021_paper.html）
首先，为了降低计算复杂度，自注意力机制的计算范围限制在非重叠层本地窗口内，每个窗口默认包含 7×7 个图像块。（图 3 为了方便图解，将窗口描绘为 4×4 个图像块。）Swin Transformer 模型共有 4 个阶段，为了处理输入图像中的非本地大型空间关系信息，图像块序列经过每个阶段时都会被下采样。下采样由可训练的线性层完成，将 2×2 个图像块连接起来进行下采样，将特征维度扩大 2 倍，使下采样效率提升 2 倍。
由此形成了一个特征金字塔，可以用在需要密集本地计算和全局特征的任务中，与传统 U-Net 架构中生成的特征金字塔相似（Ronneberger et al. 2015）。与特征金字塔类似，SIFT 和 SURF 等经典图像处理算法中会使用图像金字塔来获取尺度不变特征（scale-invariant feature）。默认情况下，Swin Transformer 并不处理[CLS] token，而是将最后一个阶段已平均的特征喂入 classification head，后者用于根据 ImageNet 数据集对模型进行有监督训练。
此外，每个阶段中，基于移位窗口有两种不同的图像块切分方式（图 7）。此处，移位窗口包含之前已分离且不参与自注意力机制计算的图像块。论文显示，在 ImageNet 图像分类任务、COCO 目标检测任务和 ADE20k 语义分割任务中使用移位窗口，都可以实现比静态窗口高得多的性能。总体而言，由于每个窗口的图像块数量固定，因此 Swin Transformer 的内存和计算复杂度从二次方增长降低为线性增长，性能比原始 ViT 模型和 DeiT 模型更加出色。
除了 Swin Transformer，NLP 和计算机视觉领域还有很多方法可以基于标准双向 softmax 自注意力机制降低 Transformer 模型的计算复杂度，包括使用 softmax 自注意力机制的稀疏近似或低秩近似（low-rank approximation）（Wang et al. 2020, Zaheer et al. 2020, Kitaev et al. 2020, Choromanski et al. 2021），调整注意力机制（Ali et al. 2021, Jaegle et al. 2021, Lu et al. 2021, Jaegle et al. 2022），优化标准方法——例如优化 IO 操作等（Dao et al 2022）。为了对不同的 Transformer 架构进行深度比较，我们使用的是 Long Range Arena 这一较新的评估基准（Tay et al. 2021）。
DINO 模型（Self-distillation with no labels）
目前谈及的 ViT 模型都经过有监督的分类任务（预）训练。Caron 等人从另一个方向进行研究，在 DeiT 的基础上构建了一种自监督训练 ViT 模型，也取得了良好成果。这种模型无需使用 DeiT 所用的显式 teacher 模型，而是引入了一种无标签自蒸馏（self-distillation with no labels, DINO，https://openaccess.thecvf.com/content/ICCV2021/html/Caron_Emerging_Properties_in_Self-Supervised_Vision_Transformers_ICCV_2021_paper.html）方法。
图 8：DINO 架构。图源 Meta AI 文章
如图 8 所示，teacher 模型被定义为 student 模型的指数移动平均值（EMA），student 模型是一个标准的 ViT 模型。NLP 领域常见的自监督训练方法使用遮罩（masking）和 input token 序列补全方法来避免模型坍塌（Devlin et al. 2019, Brown et al. 2020），而视觉领域的自监督学习通常使用的是对比损失（contrastive loss）。在对比学习框架中，两个版本的神经网络处理相同（或不同）的增强图像，经过训练努力输出相同（或不同）的表示。
然而，与近来的 BYOL (Grill et al. 2020)、Barlow Twins (Zbontar et al. 2021)和 VICReg (Bardes et al. 2022)等自监督学习方法相似，DINO 并不通过对比损失显式利用负样本。相反，DINO 将同一输入图像的两个增强版本分别传入 student 模型和 teacher 模型，teacher 模型处理一张原图像的全局裁剪（涵盖输入图像的较大区域），student 模型处理一张原图像的局部裁剪（涵盖输入图像的较小区域）。使用简单的 MLP projection head 从[CLS] token 的输出表示中得到预测的 softmax logit，然后用交叉熵损失与 teacher 模型中的 logit 进行比较。梯度只在 student 模型网络中传播。
还有其他工程方法可以避免模型坍塌，这些方法的核心概念都是使 student 模型预测出一个与 teacher 模型的预测结果相似的图像表示——teacher 模型处理的输入图像与 student 模型的略有不同——这样可以提升从局部到全局的关联性。用 ImageNet 图像和自蒸馏方法训练 DINO（不使用标签），结果发现 DINO 可以在无监督的情况下为同一标签下的图像生成相似的表示。DINO 通过简单的 KNN 分类器和潜在图像表征就在无监督 ImageNet 图像分类任务上取得了出色表现。
论文还显示，DINO 可以高效学习包含输入图像语义分割信息的表示（图 9）。为了提供合理分辨率的语义分割映射，DINO 中的图像块尺寸为 4×4，比 ViT 和 DeiT 的都小得多。基于 DINO 表示的自然语义集群，近期一些论文还探讨了自监督语义分割（Hamilton et al. 2022）。
图 9：标准视觉主干网络无监督训练得到的分割图（上）DINO 无监督训练得到的分割图（下）（图片来自相应论文）
5、总结
多年来，Transformer 模型成为 NLP 的首选模型，如今，Transformer 也逐渐取代 CNN 成为计算机视觉领域的 SOTA 模型。ViT-L 和 ViT-H 等大型 Transformer 模型也被视为基础模型，它们在作为通用型视觉主干网络使用时也表现出优异性能，特别是利用大型数据集训练之后。
然而，Transformer 模型应用于图像和视频时还存在不少难点，业内正在进行积极研究。我们还须探索高样本效率的自监督学习策略和使 Transformer 模型更适应特定领域或数据集的方法，相信也会引发更多应用型计算机视觉研究。
许多主流的 Vision Transformer 架构都由 Google (原始 ViT)、Microsoft (Swin)、Huawei (TNT)、 OpenAI (iGPT)、Meta (DINO)等科技巨头研发，需要使用大量计算资源。为了降低从零开始训练 Vision Transformer 的计算复杂度、硬件要求和数据要求，业内进行了大量研究。本文介绍了一部分但未能穷举。相信 Vision Transformer 在视觉领域的运用将会越来越成功。
（本文经授权后由 OneFlow 编译发布，译文转载请联系 OneFlow 获得授权。原文：1. https://medium.com/merantix-labs-insights/a-brief-history-of-vision-transformers-revisiting-two-years-of-vision-research-26a6bd3251f3；2. Merantix Momentum GitHub: https://github.com/merantix-momentum）
其他人都在看
欢迎下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-Inc/oneflow/
不至于成为世界上最快的深度学习框架。 2022-03-23 加入
★ OneFlow深度学习框架：github.com/Oneflow-Inc/oneflow ★ OF云平台：oneflow.cloud
评论