2021-ICLR-An Image is Worth 16x16 Words:Transformers for Image Recognition at Scale

1. 摘要

这篇文章[1]主要提出如何将 Transformer 用在计算机视学领域(即 ViT, Vision Transformer),用一句话概率就是这篇文章的标题:An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,即将图像划分成许多 16×1616 \times 16 的 patch,这样每一个 patch 就相当于自然语言中的一个单词,然后将这些 patch 组成一个序列,并使用一个线性层对每个 patch 进行编码(类似于自然语言处理中的词嵌入),最后就可以将编码得到的序列数据直接输入原版的 Transformer 进行训练,而不需要对架构上做任何针对 CV 的改动。作者提出,虽然注 Transformer 早已被提出并在 NLP 领域引发了一阵热潮,但在 CV 领域 CNN 仍然占主导地位。过去一些尝试将注意力机制用在 CV 中的工作,要么是将注意力和 CNN 结合,要么是将注意力机制作为层嵌入一般的 CNN 网络架构(比如 ResNet),而作者提出其实可以完全不依赖于 CNN,纯 Transformer 也可以在 CV 任务上表现得很好。特别是,在大数据集上进行训练得到预训练模型后,在小数据集上进行微调测试的效果能够媲美 SOTA 的 CNN 网络。

2. 引言

自 Transformer 出现在 NLP 领域以来,CV 领域也有一些工作尝试将其迁移过来。但由于图像数据过高的维度,直接按像素展开得到序列数据会导致序列过长,无法使用现有的硬件设备进行处理。过去的有关这方面的一些工作,主要都是在处理如何得到降低这个序列数据的长度。

  • Wang et al. 2018[2]就提出,可以先使用一个 CNN 来提取图像的特征,然后将特征展开作为图像的序列数据。比如 ResNet50,最后提取得到的特征 Map 只有 14×1414 \times 14,展开成一维就得到 196196 长度的序列数据。

作者希望能够直接将将原来的 Transformer 架构用在 CV 领域,并且和在 NLP 领域类似,可以通过在大数据集上预训练后用于各种 CV 任务,以达到和 NLP 领域的大统一。于是作者就提出,将图像划分成 16×1616 \times 16 的 patch,每个 patch 相当于自然语言中的一个单词,这样一张 224×224224 \times 224 的图像就只有 14×1414 \times 14 个 patch,展开成序列长度只有 196196,解决了图像序列长度过大的问题。

不过作者也提到,ViT 在一些 CV 任务上表现并不如现有的最好的 CNN,很大原因是 ViT 缺乏 CNN 的那种归纳偏置。

对于 CNN,一般认为其有两点归置偏置:

  1. locality:即图片上相邻的区域会有相邻的特征;
  2. translation equivariance:即对图片先做平移后做卷积,和先做卷积后做平移,最后的结果都是一样的,即平移等变性。

但是当模型和数据集大小都提上来时,ViT 的性能也就上来了,能够媲美 SOTA 的 CNN 网络。

3. 架构

ViT 的整体流程:首先先将图片划分为许多 16×1616 \times 16 的 patch,然后经过一个线性投射层得到每个 patch 的编码向量,再将位置编码加入编码向量,同时引入一个全局 patch * 及其位置编码(即类似 BERT 中的 [cls] token),一同输入 Transformer 编码器中,最后只取 patch * 对应的输出,过一个 MLP Head 后用于分类,使用交叉熵损失。从上图右边可以看到,ViT 使用的就是标准的 Transformer 编码器。具体来说,如果图片大小为 224×224224 \times 224,则

  • patch 个数:(224/16)2=196(224 / 16)^2 = 196;
  • patch 大小:16×16×3=76816 \times 16 \times 3 = 768
  • 输入向量:196×768196 \times 768
  • 线性投射层输出:196×768196 \times 768
  • 加上位置编码向量:196×768196 \times 768(位置编码向量长度为 768768);
  • Transformer 编码器输入:197×768197 \times 768(加上了 patch *);

由于作者在本文中只做了分类任务,因此只用到了 Transformer 的编码器部分。

4. 消融

4.1 [cls] token

在经典的 CNN 中,得到最终的分类向量都是都是通过在最后的 Feature Map 上使用一个 GAP(全局池化)操作;而在 ViT 中,作者为了尽可能和原版的 Transformer 保持一致,因此借鉴了 BERT 中的 [cls] token 的思想。作者也提到,如果不使用 [cls] token,对 Transformer 输出的所有序列数据使用一个 GAP 也是可以的,二者效果差不多(不过要好好调参),如下图所示:

4.2 位置编码

作者提出,使用 11-D 和 22-D 或者相对的位置编码效果差不多,但没有位置编码就会导致性能下降。

  • 11-D 位置编码:对图像 patch 使用光栅扫描展成一维序列;
  • 22-D 位置编码:将图像 patch 看作是栅格上的一点,如果位置编码的长度为 DD,就分别使用 D2\frac{D}{2} 的向量来表示每个栅格点的横纵坐标,最后再拼接得到 DD 长的位置编码;
  • 相对位置编码:在进行 Transformer 编码器前,做一次额外的注意力:对 Query pqp_q 和 Key pkp_k 使用使用相对注意力,即用 pqpkp_q - p_k 表示这两个 patch 的相对位置信息,然后对 pqpkp_q - p_k 做 Embedding 后再计算注意力,这样就将相对位置信息加入到了输入数据中。

但最后作发现,这三种方式都差不多:

不过,不同的超参数会导致学到的位置编码不同,不过总体而言,ViT 还是学到了真实的位置编码,如下图所示:

注意,作者在本文中使用的 11-D 和 22-D 位置编码都是可学习的位置编码。设位置编码的长度为 DD

  1. 对于 11-D 的位置编码,作者构造了一个 196196 维的位置编码表,表中第一行是一个可学习的位置编码,长度为 DD
  2. 对于 22-D 的位置编码,作者分别构造了两个 196196 维的位置编码表。一个是行位置编码表,其中每一行长度为 D2\frac{D}{2};另一个是列位置编码表,其中每一行长度为 D2\frac{D}{2}。求每个 patch 的位置编码时,用其对应的行的位置编码和列的位置编码拼接在一起得到最终的位置编码。

5. 实验

作者在大数据集 JFT-300M 上预训练后,在测试数据集上进行微调后,实验结果对比如下图所示:

可以看到,ViT 基本取得最好的结果。不过提升也比较,所以作者对比了另一个指标,即所需的训练天数。更详细地,作者和经典的 ResNet 进行了对比,分别从数据集大小、所需的训练资源等方面进行了对比:

可以发现,当训练数据集规模上来时,ViT 的强在性能也就展现了出来。此外,作者也做了一些可视化解释,如下图所示:

从左到右分别表示线性 Embedding 层的可视化,位置编码相似度的可视化和多头注意力距离可视化。

  • Embedding 层的可视化:可以看到,类似于 CNN 一样,Embedding 层学到了一些 gabor filter ,这些 filter 是可以作为基函数来表示真实的图像 patch;
  • 位置编码相似度的可视化:和前面说到的一样,ViT 的确学到了 patch 的位置编码信息;
  • 多头注意力距离可视化:可以看到,在比较浅层的网络,多头注意力的确注意到不同距离的图像 patch。

6. 结论

这篇文章探索了将原版的 Transformer 扩展到了 CV 任务。除了位置编码和划分 patch 有一些归纳偏置外,ViT 没有其它 CV 任务特有的归纳偏置,做到了和 NLP 的大一统,也为后来多模态任务发展奠定了基础。ViT 不仅可扩展性好,而且相比于一些明星的 CNN(比如 ResNet),所要的训练资源还更少些。而且,ViT 训练时使用的图像是 224×224224 \times 224 的,因此如果在测试时换成其它分辨率的图片,就会造成不兼容,因此训练时学习到的位置编码只有 16×1616 \times 16 个,无法直接用在不同分辨率的图片上。一种解决方法是对位置编码进行插值,但当分辨率差异很大时,插值的方法也很难奏效。

不过,ViT 只在分类任务上做了实验,在检测和分割任务上还没有涉及,因此 ViT 在其它各种 CV 任务的扩展还有待发展。此外,在 NLP 领域,自从 Transformer 的提出,诸如 BERT 等各式各样的模型都是使用自监督学习的方式进行训练,而在 CV 领域还是以有监督为主。因此将 ViT 用于自监督学习上是未来需要探索的方向。作者虽然有作自监督的实验,但效果不太好,不过后续的 MAE 等方法的确说明了 Transformer 真的也可以在 CV 自监督任务上做到很好。

在 ViT 出现以后,也涌现了很多工作来分析 Transformer 为什么这么有效。有人发现,将自注意力换成 MLP,甚至将自注意力换成不能学习的池化操作,仍然可以工作的很好。因此,有人提出,可能 Transformer 的整体结构才是其如此有效的原因,而不是其中的自注意力等算子。

附录

  1. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., ... & Houlsby, N. (2020, September). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations.
  2. Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In CVPR, 2018