架构演进 · 视觉与扩散

扩散 Transformer

DiT-XL/2 · 预测噪声，而不是下一个词元

对比
DiT-XL/2 vs GPT-2 small
层
28 GPT-2 12
宽度
1,152 · 16 个头 GPT-2 768 · 12
词元
256 个潜空间图像块 GPT-2 ≤ 1,024 文本
输出
每个图像块的噪声 GPT-2 50,257 个分数
参数
675M · 118.6 GFLOPs GPT-2 124M

全部步骤

  1. GPT-2 vs DiT

    DiT 保留了 Transformer 块，但改变了流经它的东西：带噪声的图像块而不是词元，没有掩码，输出是每个图像块中的噪声。时间步和类别标签不是词元；它们设定每个 LayerNorm 的缩放和平移，并给每个子层加门控（adaLN-Zero）。点击标签跳过去。

    28 × 1,152 · 675M · 256 个图像块

  2. 加噪

    扩散模型训练一个模型来撤销噪声。一张干净的图像与高斯噪声混合，t = 1 时混一点点，到 t = 1,000 时完全是噪声；模型看到带噪图像和 t，预测加入的噪声。这是在一张玩具图像上、按 DiT 的调度精确混合的结果。

    xt = √ᾱt · x0 + √(1 − ᾱt) · ε

  3. 潜空间图像块

    DiT 作用于压缩后的图像：一个预训练的自编码器把 256 × 256 × 3 像素变成 32 × 32 × 4 的潜变量，其中 2 × 2 的图像块组成 256 个各含 16 个数的词元。图像块越小，词元越多，计算越多，图像也越好。

    256² × 3 → 32² × 4 → 256 个词元

  4. adaLN-Zero

    在 GPT-2 中，LayerNorm 的 γ 和 β 训练后就固定了。在 DiT 中，一个小网络根据时间步和类别为每个块分别算出它们，以及每个子层输出上的一个门控。这些是真实的 DiT-XL/2 数值；有些子层几乎被关掉了。

    x + α · f(LN(x) · (1 + γ) + β)

  5. 每张图要跑很多趟

    生成一张图从纯噪声开始，每次去掉一点：250 步，每一步都完整跑一遍全部 28 个块，用无分类器引导时还要跑两遍。这里每一步都用真实噪声，就像一个完美的模型那样；真实模型只能估计它。

    250 趟 × 118.6 GFLOPs

代码

# the DiT block: the timestep and class set LayerNorm's scale and shift, and a gate
c = t_embedder(t) + y_embedder(y)                        # (1152,)
shift1, scale1, gate1, shift2, scale2, gate2 = adaLN_modulation(c).chunk(6)
x = x + gate1 * attn(layer_norm(x) * (1 + scale1) + shift1)
x = x + gate2 * mlp(layer_norm(x) * (1 + scale2) + shift2)
# training: add noise at a random t, predict it
x_t = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * eps
loss = F.mse_loss(model(x_t, t, y), eps)
# sampling: 250 steps, each a full forward pass (twice with guidance)
for t in reversed(steps): x = p_sample(model, x, t, y)

延伸阅读