基本信息

标题

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

发表期刊

ICLR2021 link

作者

google brain

日期

202010

code

github

colab

概览

研究背景

目的是transformer architecture 在视觉领域的研究
之前的模型通常是在cnn的基础上加入attention机制；本文想要排除对cnn的依赖，直接使用pure transformer来做分类

方法和性质

模型概述

将图像分成固定大小的patches，线性的patch embedding，并添加position embedding，得到一个sequence
将sequence送入一个标准的transformer encoder
为了分类任务，在sequence加入一个额外的classification token

patch embedding & position embeddings & class token

\mathbf{z}_{0} = [\mathbf{x}_ {\text{class}}; \mathbf{x}_{p}^{1} \mathbf{E} ; \mathbf{x}_{p}^{2} \mathbf{E} ;\cdots ; \mathbf{x}_{p}^{N} \mathbf{E} ] + \mathbf{E}_{\text {pos }} , \mathbf{E} \in \mathbb{R}^{\left(P^{2} \cdot C\right) \times D}, \mathbf{E}_ {\text {pos}} \in \mathbb{R}^{(N+1) \times D}

将2D image $\mathbf{x} \in \mathbb{R}^{h \times w\times c}$ 变为适合transformer输入的1D sequence $\mathbf{x_p} \in \mathbb{R}^{n \times (p^{2} \cdot c)}$ .

$(h,w)$ 是原始图像的分辨率， $C$ 是通道数；
$(p,p)$ 是每个patch的分辨率；
$n = hw/p^{2}$ 是一个图片的patch个数，同时也是transformer输入的序列长度
使用可训练的linear projection将压平后的patches映射到D维

class token

与bert的token类似，在sequence上预置可训练的embedding $z^0_0=\mathbf{x}_ \text {class}$
其在encoder输出的状态作为y 及 $y = LN(z^0_L)$
在预训练和fine-turining时，分类头附加到 $z^0_L$
分类头由一个隐含层的MLP实现，在预训练时只有一个线性层

使用标准的可训练的1D 位置嵌入

作者提到更新的 2D-aware位置嵌入，并没有显著的提升

transformer encoder

\mathbf{z}_ {\ell}^ {\prime} = MSA(LN(\mathbf{z}_ {\ell}^ {\prime} ))+\mathbf{z}_{\ell-1} , \ell = 1 \ldots L \\ \mathbf{z}_{\ell}^{\prime} = MLP (LN (\mathbf{z}_ {\ell}^ {\prime }))+\mathbf{z}_{\ell}^{\prime} , \ell = 1 \ldots L \\ y = LN(z^0_L)

由L个MSA（multiheaded self-attention） block和L个MLP block交替构成
- 每个block前使用layernorm，每个block后残差连接
- 两层的MLP，激活函数为GELU

使用更大的分辨率fine-tune

删除了预训练的预测头，并添加了一个D*K维的初始化为0的层，K为下游类的数量。
微调时，使用比预训练时更高的分别率，作者选择不改变patch的大小，但会得到更长的序列

对于transformer来说，只要内存足够，它可以处理任意长度的序列；但更长的序列会导致预训练的position embedding失去意义
对预训练的position embedding，作者根据他们在原始图片中的位置进行2D插值（# 什么操作？？？看代码。。。待补充。。。）

实验

比对模型：ViT, ResNet, 混合模型
- ViT
  - 配置同bert，在base和large之外增加了更大的huge模型
- ResNet
  - group norm代替group norm
  - weight standardization
- 混合模型
  - 将中间层的特征图送入ViT
pre-train 数据集
- ImageNet：1k classes，1.3M images
- ImageNet-21k：21k classes，14M images
- JFT：18k classes， 303M high-resolution images
- 将预训练数据集与下游任务的测试集进行去重（与Big transfer中的方式相同）
training & fine-turning （记录作者的超参，方便后续代码阅读及实践）
- training
  - Adam（ $\beta_1 = 0.9$ $\beta_2 = 0.999$ ）, batch_size=4096, weight_decay=0.1
- fine-turning
  - SGD-momentum， batch_size=512
self-supervision实验尝试

实验结果与SOTA比对

与BiT-L，noisy student在分类任务上的对比
- 1，2，4列可以看出，两个模型在相同的数据集上pre-training，较小的ViT-L/16（patchsize=16 * 16）较优于BiT，而采用较大模型（patchsize=14*14）的ViT-H/14效果更好。在预训练上的计算量也更少。
与BiT-L，VIVE-Ex，S4L在VTAB任务上的对比
- BiT和ViT在specialized领域的表现几乎相同
- 该任务只展示了ViT huge模型，L和B没有列出，猜想是否效果差到不忍直视？

预训练模型实验

pre-train

该部分作者通过一系列实验证明，ViT在大规模高分辨率数据上预训练，要好于其他对比模型，且需要更小的计算量。

attention

文中给出模型所关注的是图片中与类别语义相关的部分

SELF-SUPERVISION

此外，作者受bert启发，对ViT在自监督预训练做了尝试。在imagenet上使用vit-b/16上准确率达到79.9%。效果落后有监督预训练4%

研究结果

将图像处理为sequence of patches，并通过标准transformer的encoder处理它
提供pre-train on large datasets

研究展望

将ViT应用于分类外的其他图像任务
探索模型自监督预训练方法
性能

思考及启发

本文是纯transformer架构在视觉领域的尝试。模型结构与原版transformer变化不大。个人认为比较创新的点在于图像的序列化处理，此外，作者验证并提供了大量预训练模型。

关于模型的思考

作者提到，相比cnn，transformer缺乏图像相关的归纳偏置。这使得在样本较小时，ViT的效果并不好。是否可以将归纳偏置编码到模型中？
关于patch的划分。虽然作者做了多个patch大小的对比实验，但未能证明patch大小的合理性。另外patch大小决定了模型的复杂程度，那么关于patch的划分是否有更合理的方法？
ViT在其他领域的应用

部分细节及表达不明确之处将在读完源代码后进行补充

能力一般水平有限，欢迎批评指正

ViT阅读笔记