DiT：Diffusion Transformer原理简介
论文：Scalable Diffusion Models with Transformers（ICCV 2023）
作者：William Peebles、Saining Xie
官方仓库：https://github.com/facebookresearch/DiT
一、DiT 诞生背景与核心定位
1. 传统扩散模型痛点
过往潜空间扩散模型（SD、DDPM）全部采用卷积U-Net作为去噪骨干：
- 缩放能力有限：单纯增大参数量，图像生成质量提升微弱，存在性能天花板；
- 强卷积归纳偏置：依赖局部卷积先验，对多模态、超大分辨率场景适配性差；
- 多层跳跃连接、多尺度卷积堆叠，架构复杂，分布式训练优化难度高。
2. DiT核心创新
DiT 完全抛弃卷积U-Net，用纯Vision Transformer（ViT）作为扩散模型去噪骨干，证明卷积局部先验并非图像生成的必要条件，同时具备完美的缩放定律：模型深度、宽度、Patch数量越大，算力越高，生成FID指标持续下降，画质稳定提升，无性能瓶颈。
3. 基础运行框架（纯分类条件，无文本Prompt）
DiT 基于潜扩散LDM框架，仅替换U-Net去噪网络，VAE编解码完全保留，完整数据流：
原图 → VAE Encoder → 4通道潜图 z₀ → 扩散前向加噪 → 带噪潜变量 z_t → DiT Transformer 预测噪声 → 反向还原干净潜图 z₀ → VAE Decoder 解码 → 输出RGB像素图像条件说明：原版DiT为类别条件生成，条件仅包含两项：
- 时间步ttt：代表当前潜图噪声强度；
- 类别标签yyy：ImageNet数据集数字类别ID（0~999），无文本、无CLIP编码器，不支持自然语言Prompt。
二、DiT 完整整体架构
2.1 五大核心组件
PatchEmbed（Patch分块嵌入）
输入是VAE输出的二维潜图（如32×32×4），使用卷积等价线性层，将固定尺寸Patch（论文标准patch_size=2）映射为固定维度Token序列；
随后叠加固定2D正弦余弦位置编码（不可学习），给每个Patch注入空间位置信息。
示例：32×32潜图、patch=2 → 1024个图像Token。条件编码模块（TimestepEmbedder + LabelEmbedder）
分别处理时间步ttt、类别标签yyy，输出同维度向量后相加融合，得到全局一维条件向量c=temb+yembc = t_{emb} + y_{emb}c=temb+yemb：- TimestepEmbedder：标量时间步t → 正弦频率升维256维 → 双层MLP非线性变换，输出与模型隐藏层同维度的时间嵌入；
- LabelEmbedder：数字类别ID y → Embedding查表，输出等维度类别嵌入；
关键特性：单次推理全程yyy固定不变，条件向量ccc仅计算1次，所有DiT Block共用同一份ccc。
堆叠多层DiT Block（核心计算单元）
由N层完全相同的Transformer块串联组成，所有特征调制、注意力计算均在块内完成，块内部核心创新为adaLN-Zero自适应归一化条件注入机制。FinalLayer 输出调制层
复用全局条件ccc做自适应归一化调制，将Token序列映射回Patch维度，支持同时预测噪声与方差（learn_sigma）。Unpatchify 无参数复原层
仅通过张量reshape、维度置换，无任何可训练参数，将一维Token序列还原为二维潜图，送入VAE解码器。
2.2 四种DiT Block 条件注入方案消融
论文对比4种给Transformer注入条件（t、y）的方案，最终选择性能最优、开销最低的adaLN-Zero：
| 方案 | 实现逻辑 | 计算开销 | FID效果 | 深层训练稳定性 |
|---|---|---|---|---|
| In-Context 拼接Token | 将条件编码为独立Token，拼接到图像Token序列前端，统一自注意力 | 最低 | 最差 | 一般 |
| Cross-Attention 交叉注意力 | 自注意力后新增交叉注意力层，图像Token为Q，条件向量为KV | 最高（+15%算力） | 中等 | 较差，深层易震荡 |
| 基础adaLN | 条件向量回归LN缩放γ、偏移β，无残差门控α | 低 | 中等 | 不稳定，易梯度消失 |
| adaLN-Zero（论文选用） | 基础adaLN基础上新增残差门控α，调制MLP末层全零初始化 | 低 | 最优 | 极强 |
三、核心创新：adaLN-Zero 完整原理
3.1 模块作用
摒弃交叉注意力、Token拼接等高开销方案，通过自适应归一化调制将全局一维条件ccc注入每一层图像特征；全程不新增额外注意力层，算力开销极低。
3.2 数学与数据流
- 输入：图像Token特征xxx、全局共享条件向量ccc；
- 条件投影MLP：输入ccc，一次性输出6组独立调制参数：
shiftattn, scaleattn, gateattn, shiftmlp, scalemlp, gatemlpshift_{attn},\ scale_{attn},\ gate_{attn},\ shift_{mlp},\ scale_{mlp},\ gate_{mlp}shiftattn,scaleattn,gateattn,shiftmlp,scalemlp,gatemlp - 自注意力支路流程
自适应归一化公式：h=LN(x)⋅(1+γ)+βh = \text{LN}(x) \cdot (1+\gamma) + \betah=LN(x)⋅(1+γ)+β- 使用scaleattn(γ1)scale_{attn}(\gamma_1)scaleattn(γ1)、shiftattn(β1)shift_{attn}(\beta_1)shiftattn(β1)对无参数LayerNorm输出做缩放偏移；
- 多头自注意力计算图像Patch之间的空间关联；
- 残差连接，乘门控系数gateattn(α1)gate_{attn}(\alpha_1)gateattn(α1)：xout=xin+α1⋅Attnoutx_{out}=x_{in}+\alpha_1 \cdot Attn_{out}xout=xin+α1⋅Attnout。
- FFN前馈支路：逻辑与自注意力支路完全一致，使用配套shiftmlp/scalemlp/gatemlpshift_{mlp}/scale_{mlp}/gate_{mlp}shiftmlp/scalemlp/gatemlp调制。
3.3 Zero零初始化（最关键创新点）
调制支路adaLN_modulation双层MLP的最后一层线性层权重、偏置全部初始化为0，带来两大训练增益：
- 训练初期等价恒等映射
初始化阶段6组调制参数全部输出0：- γ=0,β=0\gamma=0,\beta=0γ=0,β=0，代入公式h=LN(x)⋅(1+0)+0h=\text{LN}(x)\cdot(1+0)+0h=LN(x)⋅(1+0)+0，自适应LN退化为标准无参数LayerNorm；
- α=0\alpha=0α=0，残差支路输出全部归零；
整层DiT Block输入特征=输出特征，等价恒等映射。深层网络梯度可无衰减逐层传递，彻底解决深层Transformer梯度消失、训练震荡、难以收敛的问题。
- 渐进式学习逻辑
训练迭代中参数逐步更新，门控系数α\alphaα从0缓慢激活残差支路；模型从简单恒等映射逐步学习复杂图像特征映射关系，收敛速度大幅提升。
四、DiT 模型规格变体
论文提供4种标准尺寸，命名规则DiT-XX/p，后缀/p代表Patch尺寸：
| 模型 | Transformer层数 | 隐藏层维度 | 注意力头数 | 总参数量 |
|---|---|---|---|---|
| DiT-S/2 | 12 | 384 | 6 | 33M |
| DiT-B/2 | 12 | 768 | 12 | 130M |
| DiT-L/2 | 24 | 1024 | 16 | 458M |
| DiT-XL/2 | 28 | 1152 | 16 | 675M |
Patch尺寸消融实验
Patch越小，图像Token数量越多，模型捕捉细节能力越强，FID越低，但算力、显存占用指数级上升：
| Patch尺寸 | Token总数 | 单轮推理GFLOPs | ImageNet FID |
|---|---|---|---|
| 8 | 64 | 0.9 | 24.8 |
| 4 | 256 | 2.3 | 20.1 |
| 2 | 1024 | 5.7 | 19.5 |
五、训练与推理流程
5.1 训练阶段（DDPM扩散训练目标）
- 图像送入VAE编码器，压缩为4通道潜图z0z_0z0；
- 随机采样时间步ttt，根据扩散公式给z0z_0z0添加噪声，得到带噪潜图ztz_tzt；
- 采样类别标签yyy，编码融合得到全局条件ccc；
- DiT网络输入zt、t、yz_t、t、yzt、t、y，预测潜图中添加的噪声；
- 损失函数为预测噪声与真实噪声的MSE，反向传播更新DiT全部权重；
- 支持CFG（Classifier-Free Guidance）训练：10%概率随机丢弃类别标签，学习无条件生成，推理时提升画面细节。
5.2 推理生成阶段（采样）
- 初始化纯高斯噪声潜图；
- 设定目标生成类别yyy，全程固定不变，仅计算一次条件向量ccc；
- 从最大时间步逐步反向采样，每一步送入DiT预测当前噪声；
- 结合CFG权重融合条件/无条件噪声预测，逐步去除潜图噪声；
- 采样完成后得到干净潜图，送入VAE解码器输出最终图像。
六、DiT 对比传统U-Net扩散模型
| 对比维度 | U-Net（SD/DDPM） | DiT（原版分类条件） |
|---|---|---|
| 骨干架构 | 多层卷积、上下采样、跳跃连接 | 纯Transformer，Patch Token序列+多头自注意力，无卷积、无跳跃连接 |
| 空间归纳偏置 | 极强，卷积天然捕捉局部像素关系 | 极弱，无内置空间先验，完全依靠数据学习图像结构 |
| 条件注入方式 | 时间嵌入叠加、交叉注意力融合条件 | adaLN-Zero自适应归一化调制，无额外注意力层，算力开销更低 |
| 模型缩放能力 | 缩放收益有限，参数量提升画质增长缓慢，存在性能瓶颈 | 完美遵循缩放定律，模型越大算力越高，画质持续稳定提升 |
| 条件类型 | 可扩展文本、类别等多类条件 | 原版仅支持ImageNet数字类别标签，无文本Prompt、无CLIP编码器 |
| 训练稳定性 | 深层网络易震荡、梯度消失 | adaLN-Zero零初始化实现恒等映射，深层训练收敛稳定 |
七、工程落地关键规范
- Patch尺寸选型：论文标准折中配置
patch_size=2，兼顾画质与显存占用； - 位置编码规范：必须使用固定2D正弦余弦位置编码，
requires_grad=False不参与梯度更新，区别于ViT可学习位置编码； - 零初始化强制约束：adaLN-Zero调制MLP末层、FinalLayer输出线性层，权重、偏置必须全零初始化，否则深层模型极易训练震荡、无法收敛；
- Unpatchify无参数约束：仅做张量维度变换，不使用卷积层，无任何可训练参数，避免引入额外训练偏差；
- 数据集规范：原版DiT基于ImageNet分类数据集训练，评测标准样本量50000张；样本量低于10000时FID指标波动极大，无法用于横向性能对比。
八、总结
- 架构革新：DiT是首个完整使用纯Transformer替代扩散模型U-Net去噪骨干的框架，保留LDM潜空间压缩方案，仅替换中间去噪网络；
- 核心突破 adaLN-Zero：自适应归一化+零初始化残差门控实现轻量化条件注入，解决深层Transformer训练不稳定、梯度消失问题，综合性能优于交叉注意力、Token拼接等方案；
- 缩放定律价值：实验证明扩散模型具备明确的缩放定律，模型参数量、算力提升可稳定降低FID、提升图像生成质量，为超大尺度图像、视频生成模型奠定理论基础；
- 原版DiT边界：原版仅支持ImageNet数字类别条件生成，无文本编码器、不支持自然语言Prompt；支持文字Prompt的是后续多模态扩展MMDiT，不属于论文原版DiT架构。