DiT：扩散Transformer架构

可视化演示

展示DiT用Transformer替代U-Net的设计：4个DiT Block内部结构(AdaLN-Zero→MSA→AdaLN-Zero→MLP)，条件注入(timestep+class→MLP→modulation)，以及DiT与U-Net架构对比和缩放曲线（FID vs 模型大小）。

架构演进：从 U-Net 到 DiT

老照片修复工作室曾经有一套非常成熟的流水线：U-Net 像一批资深修复师，先把照片缩小到不同尺度分别观察，再逐步放大补细节，同时还开辟了几条“捷径”，把早期看到的结构信息直接传给后期，保证大轮廓不丢失。这套设计非常精巧，但它有一个问题：每个环节都要专门定制，想扩大规模时，卷积层、跳跃连接、通道数都要重新调，扩展性不够灵活。

DiT 的出现，相当于把这条分工细致的流水线换成了一个开放广场。广场上每个人都能看到其他人，不需要预先规定谁和谁相邻。DiT 保留了扩散模型“逐步去噪”的整体流程，只是把负责去噪的骨干网络从 U-Net 换成了 Transformer。Transformer 不假设图像有固定的空间网格，而是把图像看成一组可以互相通信的 Token，这让模型更容易通过增加深度和宽度来扩展能力。

扩散模型的训练可以这样理解：我们把干净照片 \(x_0\) 一点点加噪声，得到中间状态 \(x_t\)。这个过程可以用一个非常简洁的公式描述：

\[

\displaystyle q(\mathbf{x}_t \mid \mathbf{x}_0) = \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\,\boldsymbol{\varepsilon}

\]

1. 含噪结果：\(\mathbf{x}_t\) 是第 \(t\) 步加噪后的图像，由干净图像 \(\mathbf{x}_0\) 和随机噪声共同生成。

\[

\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\,\boldsymbol{\varepsilon}

\]

2. 原图保留系数：\(\sqrt{\bar{\alpha}_t}\) 控制原图信号还剩多少，越早期的步骤这个值越大，画面越清晰。

\[

\sqrt{\bar{\alpha}_t} \in [0, 1]

\]

3. 噪声强度系数：\(\sqrt{1 - \bar{\alpha}_t}\) 控制噪声影响有多大，\(t\) 越大，这个值越接近 1，画面越接近纯噪声。

\[

\sqrt{1 - \bar{\alpha}_t} \in [0, 1]

\]

如果把扩散生成比作一层层刮掉灰尘，U-Net 是经验丰富但结构固定的老工匠，而 DiT 则是一台可以随时加装模块的现代化设备。DiT 不改变“刮灰尘”的流程，只负责把“刮灰尘”的大脑换得更聪明、更可扩展。

图像分块：Patchify

进入 Transformer 的开放广场之前，我们得先把照片裁成统一大小的拼图块。如果整张图直接塞给 Transformer，它会因为数据量太大而无法处理；如果只塞一个像素，它又看不到足够丰富的语义。于是 DiT 采用了 ViT 的经典做法：Patchify。

Patchify 很像把一张大幅照片裁成一叠小方块。每个小方块被“压平”并映射成一个向量，也就是 Transformer 世界里的 Token。为了让模型知道谁在左上角、谁在右下角，每个 Token 还会额外加上位置编码。这样，原本像素之间的强局部关系，就被转换成序列中 Token 之间的全局关系。

假设图像进入 DiT 前的特征图高度为 \(H\)、宽度为 \(W\)，每个 Patch 的大小是 \(p\)，那么 Patch 的总数也就是 Token 序列长度为：

\[

\displaystyle N = \left\lceil \frac{H}{p} \right\rceil \times \left\lceil \frac{W}{p} \right\rceil

\]

1. 行方向块数：高度 \(H\) 除以每个 Patch 的高度 \(p\)，向上取整得到每一列能放几个 Patch。

\[

\left\lceil \frac{H}{p} \right\rceil

\]

2. 列方向块数：宽度 \(W\) 除以 Patch 的宽度 \(p\)，向上取整得到每一行能放几个 Patch。

\[

\left\lceil \frac{W}{p} \right\rceil

\]

3. 序列总长：把行块数和列块数相乘，就得到送入 Transformer 的 Token 数量。

\[

N = \left\lceil \frac{H}{p} \right\rceil \times \left\lceil \frac{W}{p} \right\rceil

\]

Patchify 是 DiT 与 U-Net 最大的分水岭之一。U-Net 天然把图像当作“有上下左右关系的像素网格”，而 DiT 把图像当作“一组顺序可以被打乱但需位置编码提示的拼图”。拼图切得越小，模型越能关注细节，但要处理的 Token 也越多，计算量会快速增长。

DiT 结构：让 Patch 彼此对话

在 DiT 的 Transformer Block 里，每个拼图块不再只是被局部卷积核扫过，而是进入一个开放大厅：所有拼图块同时发言，每块都能听到其他所有块的声音。这种“所有人对话所有人”的机制叫做自注意力（Self-Attention）。

自注意力会为每个 Token 生成三个角色：Query（我要找什么）、Key（我有什么）、Value（我能提供什么）。Token 之间通过 Query 和 Key 的匹配程度决定应该关注谁，再用对应的 Value 聚合信息。公式如下：

\[

\displaystyle \operatorname{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V}) = \operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_k}}\right)\mathbf{V}

\]

1. 相关度打分：用 Query 矩阵 \(\mathbf{Q}\) 和 Key 矩阵 \(\mathbf{K}\) 做点积，得到每两个 Token 之间的相关分数。

\[

\mathbf{Q}\mathbf{K}^{\top}

\]

2. 缩放稳定：除以 \(\sqrt{d_k}\)，避免维度变大后点积数值过大，导致 softmax 进入饱和区。

\[

\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_k}}

\]

3. 加权聚合：softmax 把分数变成非负且和为 1 的权重，再乘上 Value 矩阵 \(\mathbf{V}\)，完成信息融合。

\[

\operatorname{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_k}}\right)\mathbf{V}

\]

真正实用的 DiT Block 会使用多头注意力：把特征拆成几个子空间，每个头关注不同类型的关系，最后拼接起来。每个 Block 还会接一个前馈网络，对每个 Token 独立做非线性变换，并用残差连接把输入重新加到输出上。整体结构非常规整，像一栋每层结构都相同的办公楼，方便工程上堆叠更多楼层。

条件调制：AdaLN-Zero

老照片修复师不能闷头乱修，他需要知道两件事：现在是修复到第几步了？这张照片属于什么类型？Diffusion 模型也一样，DiT 必须把“当前噪声程度”和“类别标签”注入网络。DiT 使用的注射器叫 AdaLN-Zero。

AdaLN 的全称是 Adaptive Layer Normalization，也就是自适应层归一化。普通 LayerNorm 会把特征统一拉到一个标准范围，而 AdaLN 会根据条件 \(c\) 动态调整每个 Token 的缩放和偏移。AdaLN-Zero 中的 “Zero” 更是关键：每个残差分支的输出层初始化为零，让每个 Block 一开始什么都不做，相当于恒等映射。

\[

\displaystyle \mathbf{h}' = (1 + \gamma(\mathbf{c})) \odot \operatorname{LayerNorm}(\mathbf{h}) + \beta(\mathbf{c})

\]

1. 归一化：对输入特征 \(\mathbf{h}\) 做 LayerNorm，把每个 Token 的特征向量拉成均值为 0、方差为 1 的分布。

\[

\operatorname{LayerNorm}(\mathbf{h})

\]

2. 条件缩放：\(\gamma(\mathbf{c})\) 是由条件 \(c\) 生成的缩放向量，先加 1 再逐元素乘到归一化结果上。

\[

(1 + \gamma(\mathbf{c})) \odot \operatorname{LayerNorm}(\mathbf{h})

\]

3. 条件偏移：\(\beta(\mathbf{c})\) 是由条件 \(c\) 生成的偏移向量，用来让特征按条件平移到合适位置。

\[

\mathbf{h}' = (1 + \gamma(\mathbf{c})) \odot \operatorname{LayerNorm}(\mathbf{h}) + \beta(\mathbf{c})

\]

AdaLN-Zero 让 DiT 的训练非常稳定。DiT 论文实验表明，它比直接把条件拼进序列、或者用 Cross-Attention 注入条件的方式都更高效。它像一个聪明的调味师：先让菜品保持原味，再逐步根据顾客要求加盐加辣。

条件生成：类别引导

顾客来修复照片时往往会说：“这是一张风景照”或“这是一张人物照”。DiT 把这种类别提示编码成一个向量，作为条件加入模型。训练时，模型有时会看到类别条件 \(y\)，有时会故意看不到条件，从而同时学会“有条件生成”和“无条件生成”。

真正采样时，为了得到更符合类别提示的结果，DiT 会把两条路线的预测噪声做外推，这就是无分类器引导（Classifier-Free Guidance）。公式中 \(w\) 是引导强度：

\[

\displaystyle \hat{\boldsymbol{\varepsilon}}_\theta = (1 + w)\,\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,t,y) - w\,\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,\varnothing)

\]

1. 条件预测：\(\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,t,y)\) 表示看到类别 \(y\) 时模型预测的噪声。

\[

\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,t,y)

\]

2. 无条件预测：\(\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,\varnothing)\) 表示不提供类别时模型预测的噪声。

\[

\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,\varnothing)

\]

3. 加权外推：用 \(1+w\) 放大有条件预测，再减去 \(w\) 倍无条件预测，得到最终更强烈的条件信号。

\[

\hat{\boldsymbol{\varepsilon}}_\theta = (1 + w)\,\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,t,y) - w\,\boldsymbol{\varepsilon}_\theta(\mathbf{x}_t,\varnothing)

\]

在 DiT 中，类别标签和时间步会一起生成条件向量 \(c\)，再通过 AdaLN-Zero 注入每一个 Block。这样模型既能知道“现在该修复到什么程度”，也能知道“我该生成哪一类照片”，两者共同控制生成方向。

缩放策略：规模定律

DiT 论文最吸引人的地方，是它系统地研究了“模型变大之后会发生什么”。DiT 给出的答案是：在扩散模型里，Transformer 骨干拥有非常漂亮的缩放曲线——只要算力增加，生成质量就会稳定提高。

为什么 DiT 能做到这一点？因为它把模型能力拆成了几个清晰可控的方向：模型宽度（隐藏维度）、模型深度（Block 数量）、Patch 大小。DiT 通常会先让 VAE 把 256x256 的原图压缩成 32x32 的潜空间特征图，再在这个特征图上做 Patchify。假设潜空间边长为 32，Patch 大小为 \(p\)，那么 Token 数为：

\[

\displaystyle N = \left(\frac{32}{p}\right)^2

\]

1. 单边块数：潜空间边长 32 除以 Patch 大小 \(p\)，得到每边有多少个 Patch。

\[

\frac{32}{p}

\]

2. 二维网格：因为潜空间是正方形，所以 Token 数是单边块数的平方。

\[

\left(\frac{32}{p}\right)^2

\]

3. 代入观察：\(p=8\) 时只有 16 个 Token，\(p=2\) 时则有 256 个 Token，计算量随 Token 数平方级上涨。

\[

N = \left(\frac{32}{2}\right)^2 = 256

\]

DiT 的实验清晰展示了一条趋势：无论增大模型宽度、增加 Block 深度，还是减小 Patch 大小，只要总计算量 GFLOPs 上升，FID 指标就会下降，也就是生成质量更好。这说明 DiT 是“吃算力”的架构，也说明它具备良好的可扩展性。

对比效果：Transformer 骨干的价值

U-Net 像一位经验丰富的本地工匠，非常擅长在小数据集上用卷积的局部归纳偏置快速出活。但它的网络结构高度定制，想把模型从几千万参数扩展到几十亿参数，需要大量手工调整。DiT 则像一支具备标准流程的现代团队，结构简单统一，只要增加 Block 和宽度，就能平滑扩大容量。

DiT 的每个 Block 可以写成下面这种简洁形式：

\[

\displaystyle \mathbf{x}_{l+1} = \mathbf{x}_l + \operatorname{Attn}(\mathbf{h}) + \operatorname{FFN}(\mathbf{h}')

\]

1. 恒等主干：\(\mathbf{x}_{l+1}\) 是在上一层特征 \(\mathbf{x}_l\) 的基础上累加得到的，保证信息可以跨层直通。

\[

\mathbf{x}_{l+1} = \mathbf{x}_l + \cdots

\]

2. 注意力分支：\(\operatorname{Attn}(\mathbf{h})\) 是经过 AdaLN 调制后的多头自注意力结果，负责全局信息交换。

\[

\operatorname{Attn}(\mathbf{h})

\]

3. 前馈分支：\(\operatorname{FFN}(\mathbf{h}')\) 是经过第二次调制的 MLP 结果，负责逐 Token 的非线性变换。

\[

\operatorname{FFN}(\mathbf{h}')

\]

正是这种“主干恒等 + 分支注入”的简洁结构，让 DiT 比 U-Net 更容易扩大规模。最终，DiT 在 ImageNet 上的图像生成 FID 刷新了当时的纪录，并验证了 Transformer 在扩散模型中的骨干价值。后来许多知名生成模型都延续了 DiT 的设计思路，可以说 DiT 把扩散模型带入了一个新的架构时代。

AI 应用场景

Python 演示

import numpy as np def softmax(x):    x = x - x.max(axis=-1, keepdims=True)    e = np.exp(x)    return e / e.sum(axis=-1, keepdims=True) class MiniDiTBlock:    def __init__(self, dim=16, heads=4):        self.dim = dim        self.heads = heads        self.head_dim = dim // heads        self.scale = self.head_dim ** -0.5        self.w_q = np.random.randn(dim, dim) * 0.02        self.w_k = np.random.randn(dim, dim) * 0.02        self.w_v = np.random.randn(dim, dim) * 0.02        self.w_o = np.zeros((dim, dim))  # 零初始化残差输出        self.w_ff1 = np.random.randn(dim, dim * 2) * 0.02        self.w_ff2 = np.zeros((dim * 2, dim))  # 零初始化前馈输出        self.w_mod = np.random.randn(dim, dim * 4) * 0.02     def forward(self, x, c):        B, T, D = x.shape        mod = c @ self.w_mod        shift_msa, scale_msa, shift_mlp, scale_mlp = np.split(mod, 4, axis=-1)         def adaln(h, scale, shift):            mu = h.mean(axis=-1, keepdims=True)            var = h.var(axis=-1, keepdims=True)            norm = (h - mu) / np.sqrt(var + 1e-6)            return norm * (1 + scale[:, None, :]) + shift[:, None, :]         h = adaln(x, scale_msa, shift_msa)        Q = h @ self.w_q        K = h @ self.w_k        V = h @ self.w_v        Q = Q.reshape(B, T, self.heads, self.head_dim).transpose(0, 2, 1, 3)        K = K.reshape(B, T, self.heads, self.head_dim).transpose(0, 2, 1, 3)        V = V.reshape(B, T, self.heads, self.head_dim).transpose(0, 2, 1, 3)         attn = softmax((Q @ K.transpose(0, 1, 3, 2)) * self.scale)        out = (attn @ V).transpose(0, 2, 1, 3).reshape(B, T, D)        x = x + out @ self.w_o.T         h2 = adaln(x, scale_mlp, shift_mlp)        x = x + np.maximum(0, h2 @ self.w_ff1) @ self.w_ff2        return x def demo():    np.random.seed(42)    B, T, D = 1, 16, 16    x = np.random.randn(B, T, D)    c = np.random.randn(B, D)    block = MiniDiTBlock(D, 4)    y = block.forward(x, c)    print('输入:', x.shape, '输出:', y.shape)    print('首 Token 输出均值:', round(float(y[0, 0, :].mean()), 4)) if __name__ == '__main__':    demo()

1. self.w_o = np.zeros((dim, dim))零初始化残差输出，对应 AdaLN-Zero 的核心思想：初始阶段不让新增分支干扰主路径。
2. mod = c @ self.w_mod条件向量 \(c\) 通过线性层生成四组调制信号，包括注意力和前馈网络各自的 scale 与 shift。
3. shift_msa, scale_msa, shift_mlp, scale_mlp = np.split(mod, 4, axis=-1)把调制向量拆成四份，分别控制注意力层和 MLP 层。
4. Q = Q.reshape(B, T, self.heads, self.head_dim).transpose(0, 2, 1, 3)把 Token 特征拆成多个头，每个头在独立子空间计算注意力。
5. attn = softmax((Q @ K.transpose(0, 1, 3, 2)) * self.scale)用缩放点积计算注意力权重，让每个 Patch 决定该关注哪些其他 Patch。
6. x = x + out @ self.w_o.T残差连接加上注意力输出，保证信息能在深层网络中直接流动。
7. x = x + np.maximum(0, h2 @ self.w_ff1) @ self.w_ff2MLP 分支同样通过残差方式叠加，完成逐 Token 非线性变换。

常见误区

面试题