架构演进 · 视觉与扩散

视觉 Transformer

ViT-B/16 · 图像块作为词元

对比
ViT-B/16 vs GPT-2 small
层
12 GPT-2 12
dmodel
768 GPT-2 768
头
12 GPT-2 12
词元
1 + 196 个图像块 GPT-2 ≤ 1,024 文本
输出
1,000 个类别 GPT-2 50,257 个词元
参数
86M GPT-2 124M

全部步骤

  1. GPT-2 vs ViT

    ViT-B/16 的宽度和深度与 GPT-2 small 相同，结构像 BERT（没有掩码），但它的词元是图像块：词表查找变成一次作用于原始像素的矩阵乘法，输出是整张图的一个类别，从 [CLS] 词元读出。点击标签跳过去。

    12 × 768 · 86M · 197 个词元

  2. 把图像当作序列

    一张 224 × 224 的图像被切成 14 × 14 的网格，每格是 16 × 16 的图像块。逐行读下来，这 196 个图像块就成了词元序列；每个是 16 × 16 × 3 = 768 个红、绿、蓝数值。（这里是一张玩具灰度图。）

    224² → 14 × 14 个图像块 → 196 个词元

  3. 图像块嵌入

    每个展平的图像块都乘以同一个学习得到的矩阵。它取代了 GPT-2 的词表查找，等价于步长为 16 的 16 × 16 卷积。右侧是 ViT-B/16 真实的 768 个滤波器的前几个主成分。

    x = patches · Wpatch (768 × 768)

  4. [CLS] 与位置

    最前面放一个学习得到的 [CLS] 词元，197 个位置各得到一个学习得到的位置向量，只是一个普通的一维列表。然而在真实模型中，每个图像块位置的向量都与它的邻居、以及同一行和同一列的最相似：训练重新发现了二维网格。把鼠标停在一张小图上看看。

    x = [CLS; patches] + P (197 × 768)

  5. 整张图一个标签

    全部 197 个词元经过 12 层编码器，每个图像块都关注每个图像块。分类时只用 [CLS] 的输出，映射到 1,000 个 ImageNet 类别。ViT 需要非常大的预训练数据集（多达 3 亿张图像）才能胜过卷积网络。

    class = head(LN(xCLS))

代码

# ViT: an image becomes a sequence of patch tokens
patches = img.unfold(16)                        # (196, 16 · 16 · 3 = 768)
x = patches @ W_patch + b                       # = Conv2d(3, 768, kernel=16, stride=16)
x = torch.cat([cls_token, x]) + pos_emb         # (197, 768), learned positions
for block in blocks:                            # GPT-2 small's shape, pre-LN
    x = x + attn(ln_1(x))                       # no mask: every patch sees every patch
    x = x + mlp(ln_2(x))
logits = head(ln_f(x[0]))                       # [CLS] only → 1,000 classes

延伸阅读