VisionTransformer(ViT)

Introduction

• ViT 这篇论文提出&＃xff0c;在图像分类任务中&＃xff0c;CNN 并非必须的&＃xff0c; pure Transformer 也能取得很好的效果。特别是在大量数据上预训练后再迁移到中小型数据集上时 (ImageNet, CIFAR-100, VTAB, etc.)&＃xff0c;相比 SOTA CNNs&＃xff0c;ViT 仅需更少的训练资源就能取得更好的效果

Method

Vision Transformer (ViT)

• 标准的 Transformer 接受 token embeddings 的序列&＃xff0c;因此为了能直接将 Transformer Encoder 应用到图像上&＃xff0c;我们需要将 2D 图像处理为 1D 向量。一个很直接的想法就是将图片像素经过 Embedding 层后送入 Transformer&＃xff0c;但一张图片的像素很多&＃xff0c;这样做会导致输入序列长度过长。为此&＃xff0c;ViT 将一张图片划分为若干个 $P\times P\&＃61;16\times 16$ 的 patches&＃xff0c;然后将每个 patch 拉伸为一个向量&＃xff0c;送入线性层得到相应的 Embedding。例如&＃xff0c;输入图像大小为 $224\times 224\times 3$&＃xff0c;可以划分为 N&＃61;2242162&＃61;196N&＃61;\frac{224^2}{16^2}&＃61;196N&＃61;1622242​&＃61;196 个 patches&＃xff0c;每个 patch 又被拉伸为 $16\times 16\times 3\&＃61;768$ 维的向量&＃xff0c;送入线性投影层 (i.e. $D\times D$ 的全连接层) 得到 $N\times D\&＃61;196\times 768$ 的 patch embeddings 序列。除了 patch embeddings&＃xff0c;ViT 也加入了 position embeddings 来表示 patches 之间的位置信息 (learnable 1D position embeddings)
• 为了对图片进行分类&＃xff0c;类似于 BERT&＃xff0c;ViT 也添加了 [class] token (输入序列长度为 $196\&＃43;1\&＃61;197$)&＃xff0c;它的最终输出就对应了图像表示 $y$&＃xff0c;在经过一个 classification head 之后就得到了图片的最终类别 (classification head 在预训练时为一个带有一层隐藏层的 MLP&＃xff0c;微调时为一个线性层&＃xff0c;激活函数为 GELU)

GAP v.s. [CLS]
在 CV 领域&＃xff0c;使用 CNN 进行图像分类时&＃xff0c;常见的做法是对最后的特征图进行全局平均池化 (GAP)&＃xff0c;因此作者也进行了使用 GAP 和 [CLS] 进行分类的消融实验&＃xff0c;当使用 GAP 时&＃xff0c;ViT 只对最终的 image-patch embeddings 进行全局平均池化&＃xff0c;然后使用线性层进行分类 (但注意这种方法需要的学习率与 [CLS] 完全不同)。实验结果表明这两种方法效果是差不多的&＃xff0c;但为了使得 ViT 与标准 Transformer 更接近&＃xff0c;作者最终采用了 [CLS] 进行分类

Inductive bias (归纳偏置)

• ViT 比 CNN 少了很多针对图像的归纳偏置&＃xff0c;例如在 CNN 中&＃xff0c;locality (i.e. 局部性. 图片上相邻区域会有相近特征) 和 translation equivariance (i.e. 平移等变性. 设 $f$ 为平移操作&＃xff0c;$g$ 为卷积操作&＃xff0c;则 $f(g(x))\&＃61;g(f(x))$) 体现在了模型的每一层中&＃xff0c;而 ViT 中只有 MLP 是局部且具有平移等变性的&＃xff0c;自注意力层则是全局的。在整个 ViT 中&＃xff0c;也只有 positional embedding 使用了针对图像的归纳偏置&＃xff0c;除此之外&＃xff0c;所有图像的位置信息都必须由模型从头学得
• 因此由于缺少了归纳偏置&＃xff0c;当直接在中等大小数据集 (ImageNet…) 上训练且不进行强正则化时&＃xff0c;ViT 性能并没有同等大小的 ResNet 模型好。但如果让 ViT 在更大的数据集上预训练&＃xff0c;就可以发现大规模训练胜过了归纳偏置 (88.55% on ImageNet, 90.72% on ImageNet-ReaL, 94.55% on CIFAR-100, and 77.63% on the VTAB suite of 19 tasks)

Hybrid Architecture

• 除了上述直接将图像划分为若干个 $16\times 16$ patches 的方法&＃xff0c;ViT 也提出了一种混合结构&＃xff0c;也就是先让图片经过一个 CNN 网络&＃xff0c;然后对特征图划分 patches

Fine-tuning and Higher Resolution

• 当把预训练的 ViT 模型在下游任务上微调时&＃xff0c;我们需要重新训练 classification head。如果下游任务图像分辨率更高&＃xff0c;模型往往会取得更好的效果&＃xff0c;此时如果保持 patch size 不变&＃xff0c;那么就会得到更大的序列长度&＃xff0c;但此时就无法获取有意义的 positional embedding。为此&＃xff0c;ViT 对预训练的 positional embedding 进行 2D 插值

Experiments

Setup

Datasets

• ILSVRC-2012 ImageNet dataset (1k classes and 1.3M images)
• ImageNet-21k (21k classes and 14M images)
• JFT (18k classes and 303M high-resolution images)

预训练数据集中去除了与下游任务数据集重合的图片

Model Variants

• “Base” 和 “Large” 模型与 BERT 一致&＃xff0c;此外还添加了 “Huge” 模型. 在之后的模型命名中&＃xff0c;我们使用简称&＃xff0c;例如 ViT-L/16 表示 $16\times 16$ input patch size 的 “Large” 模型
  

Baseline - ResNet (BiT)

• 网络使用 ResNet&＃xff0c;但将 Batch Normalization 替换为了 Group Normalization&＃xff0c;并且使用了 standardized convolutions&＃xff0c;这些修改提高了迁移学习的性能。我们将修改后的模型称为 “ResNet (BiT)”.

Comparison to SOTA

• 可以看到&＃xff0c;在 JFT-300M 上预训练的 ViT 在所有 benchmarks 上性能都超越了 ResNet-based baselines&＃xff0c;并且相比 CNN&＃xff0c;ViT 需要的训练资源更少

Pre-training data requirements

• 下面的实验研究了数据集大小对 ViT 性能的影响

在不同大小的数据集上预训练后&＃xff0c;在 ImageNet 上微调

灰色阴影区域代表 BiT 的性能范围 (分别使用了两个不同大小的 BiT 网络)

为了提高模型在小数据集上的性能&＃xff0c;实验时使用了 3 种正则化方法 - weight decay, dropout, 和 label smoothing 并对相应的参数进行了优化

使用 JFT 数据集的不同大小的子集进行预训练&＃xff0c;在 ImageNet 上微调

实验时不再进行额外的正则化来排除正则化带来的影响

Scaling study

在 JFT-300M 上预训练&＃xff0c;迁移到其他数据集上

• (1) ViT 在 performance/compute trade-off 上碾压了 ResNet&＃xff0c;仅需更少的训练资源就可以达到与 ResNet 相同的性能 (假设在足够大的数据集上预训练)
• (2) 在训练资源较少时&＃xff0c;hybrids 模型性能略微优于 ViT&＃xff0c;但差距随着模型参数量的增加而逐渐减小
• (3) 随着模型参数量的增大&＃xff0c;ViT 性能并没有饱和的迹象&＃xff0c;这表明继续增大模型大小可以进一步提高模型性能

Inspecting ViT

Patch embedding 层

• ViT 的 patch embedding 层将 flattened patches 进行了线性投影&＃xff0c;下图显示了学得的 embedding filters 的主成分。这些主成分似乎形成了一组能表示 patch 纹理结构的基函数
  

Position embeddings

• 下图展示了每个 patch 与其余 patch 的余弦相似度。可以看到&＃xff0c;空间上更相近的 patches 对应的 position embeddings 也更相似&＃xff0c;除此之外&＃xff0c;在同一行或同一列的 patches 对应的 position embeddings 也更相似 (a row-column structure)
  

Self-attention

• Self-attention 使得 ViT 能从最底层开始就整合图片的全局信息&＃xff0c;下面的实验探索了 ViT 在多大程度上利用了这种能力。为此&＃xff0c;我们计算了每个 Transformer 块的不同 heads 的 mean attention distance (根据每个 head 中不同 patches 之间的空间距离和 attention weight 计算&＃xff0c;类似于 CNN 中的感受野)&＃xff0c;可以发现&＃xff0c;有一些 heads 在最底层就已经开始整合图像的全局信息&＃xff0c;而另一些 heads 在最底层只负责整合局部信息 (这种局部性在 hybrids 模型中表现得更少&＃xff0c;这表明它们的功能可能与 CNN 的较低层相近)&＃xff0c;并且总体而言&＃xff0c;attention distances 随着网络深度的增加而增加&＃xff0c;表明随着层数的增加&＃xff0c;ViT 逐渐关注更高层次的信息
  

Attention map

• 这里使用 Attention Rollout 进行可视化。简单地说&＃xff0c;就是对 ViT-L/16 的所有 heads 中的 attention weights 计算平均值&＃xff0c;然后对于所有层&＃xff0c;递归地乘上权重矩阵

Self-Supervision

• Transformers 在 NLP 领域的成功不仅来源于它强大的可扩展性&＃xff0c;更在于它大规模自监督的预训练方式。因此&＃xff0c;类似 BERT&＃xff0c;我们也尝试了 masked patch prediction&＃xff0c;最终的 ViT-B/16 模型在 ImageNet 上取得了 79.9% 的精度&＃xff0c;这比从头训练得到的精度高了 2%&＃xff0c;但也比有监督学习精度低了 4%。这表明 ViT 的自监督学习仍是一个值得探索的方向

References

• paper: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Fine-tuning code and pre-trained models: https://github.com/google-research/vision_transformer
• ViT 论文逐段精读【论文精读】
• Intriguing Properties of Vision Transformers