← 全部论文 / 2020

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

2020 TPU、加速器与 Transformer 时代
下载 PDF

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale - 中文验证版

英文原文卡片：vit_2020.md

状态：已翻译。

元数据

  • Slug： vit_2020
  • 年份： 2020
  • 会议： ICLR
  • 作者： Alexey Dosovitskiy et al.
  • 阅读状态： read complete
  • 计算范式： TPU、加速器与 Transformer 时代
  • 主要来源： PDF、抽取文本

计算设置

论文明确说明，主要最新水平比较中的所有模型都在 TPUv3 硬件上训练，并以 TPUv3-core-days 报告预训练成本。它将该度量定义为 TPU v3 cores（每芯片两个）乘以训练时间（天）。在 ImageNet-21k 上预训练的 ViT-L/16 据报告可在标准云 TPUv3 的 8 cores 上训练约 30 天。

模型和数据规模也是具体的。ViT-Base 有 12 层、隐藏维度 768、MLP 大小 3072、12 个 head 和 86M 参数。ViT-Large 有 24 层、隐藏维度 1024、MLP 大小 4096、16 个 head 和 307M 参数。ViT-Huge 有 32 层、隐藏维度 1280、MLP 大小 5120、16 个 head 和 632M 参数。预训练数据集为包含 1.3M 图像的 ImageNet、包含 14M 图像的 ImageNet-21k 和包含 303M 高分辨率图像的 JFT-300M。预训练使用 Adam、batch size 4096、高权重衰减 0.1 和线性学习率预热/衰减；微调使用带动量的 SGD 和 batch size 512。

瓶颈

计算瓶颈是如何使图像识别看起来像密集的加速器友好 Transformer 计算，而不使序列长度爆炸。朴素的像素级自注意力模型需要每个像素关注每个其他像素。ViT 取而代之使序列长度 N = HW / P^2，其中 P 是 patch 大小。这意味着更小的 patch 改善空间细节，但二次增加序列长度和注意力成本。

内存是更高分辨率微调期间的实际约束。论文在输入分辨率升高时保持 patch 大小固定，这会增加有效序列长度；它指出 ViT 可以在内存约束范围内处理任意序列长度。数据是另一个瓶颈。由于内置的图像先验少于 CNN，ViT-Large 仅在 ImageNet 上预训练时表现不如 ViT-Base。只有在更大的预训练数据——特别是 JFT-300M——下，更大的模型才展示其收益。

方法适配

ViT 通过将图像转换为规则的 token 序列，使视觉任务适配加速器结构。图像拆分为固定大小的 patch，每个 patch 展平后线性投影到 Transformer 宽度，前置学习到的类别 token，并添加学习到的 1D 位置嵌入。编码器在其他方面是标准 Transformer：多头自注意力、MLP 块、layernorm 和残差连接。这保留了大型密集矩阵乘法，即 TPU 擅长的负载。

该架构有意注入非常少的图像特定结构。局部性和平移等变性内建于 CNN 层，但 ViT 主要在 patch 提取时和在新分辨率的位置嵌入插值期间使用 2D 结构。这一选择与计算相关：它避免了其实现可能更慢的专门注意力模式，即使它们的理论 FLOPs 看起来有吸引力。

论文用轴向注意力变体测试了这一点。Axial-ViT 在某些图中可以改善准确率/计算，但其开放实现据描述在 TPU 上非常慢，作者称广泛的大规模实验可能需要仔细优化的实现。这支持了卡片的主要计算观点：有用的方法不仅是较低的渐近复杂度，而且是与实际加速器软件栈的对齐。

证据

主要比较同时涉及准确率和计算。在 JFT-300M 上预训练的 ViT-H/14 达到 ImageNet 88.55%、ImageNet-ReaL 90.72%、CIFAR-100 94.55% 和 VTAB 77.63，使用约 2.5K TPUv3-core-days。ViT-L/16 在 JFT 上达到 ImageNet 87.76% 和 VTAB 76.28，使用约 0.68K TPUv3-core-days。ViT-L/16 在公开 ImageNet-21k 上达到 ImageNet 85.30% 和 VTAB 72.72，使用约 0.23K TPUv3-core-days。对比的 BiT-L ResNet152x4 使用 9.9K TPUv3-core-days，Noisy Student EfficientNet-L2 使用 12.3K。

受控计算研究指出，在性能/预训练计算权衡上，Vision Transformer 超越 ResNet，使用约少 2-4 倍的计算即可达到五个数据集上的相同平均性能。附录计时部分检查 TPUv3 上的真实硬件速度，因为 FLOPs 并不总能预测挂钟速度。它报告 ViT 模型速度与相似 ResNet 相当，且大型 ViT 模型具有明显的内存效率优势，在多种输入大小上适配更大的每核 batch。

历史影响

ViT 将视觉推进到 Transformer 加速器范式中。它表明，有了足够的数据和 TPU 规模预训练，一个基本 vanilla 的密集 Transformer 可以与高度工程化的 CNN 基线竞争或击败它们。历史影响不仅是"注意力对图像有效"；而是图像模型可以用与 NLP Transformer 相同的矩阵乘法重、大 batch、大数据集的计算语言来表达。

局限

ViT 严重依赖预训练数据规模。在仅 ImageNet 预训练下，较大的 ViT 可能不如较小的 ViT，尽管有正则化。较小的 patch 大小增加序列长度和计算。更高分辨率微调受内存限制。效率比较受优化器、调度、权重衰减和实现质量的影响。专门注意力可能提供准确率/计算收益，但论文的 TPU 实验表明未优化的变体对于大规模使用可能太慢。

链接