论文Scalable Diffusion Models with TransformersICCV 2023作者William Peebles、Saining Xie官方仓库https://github.com/facebookresearch/DiT一、DiT 诞生背景与核心定位1. 传统扩散模型痛点过往潜空间扩散模型SD、DDPM全部采用卷积U-Net作为去噪骨干缩放能力有限单纯增大参数量图像生成质量提升微弱存在性能天花板强卷积归纳偏置依赖局部卷积先验对多模态、超大分辨率场景适配性差多层跳跃连接、多尺度卷积堆叠架构复杂分布式训练优化难度高。2. DiT核心创新DiT 完全抛弃卷积U-Net用纯Vision TransformerViT作为扩散模型去噪骨干证明卷积局部先验并非图像生成的必要条件同时具备完美的缩放定律模型深度、宽度、Patch数量越大算力越高生成FID指标持续下降画质稳定提升无性能瓶颈。3. 基础运行框架纯分类条件无文本PromptDiT 基于潜扩散LDM框架仅替换U-Net去噪网络VAE编解码完全保留完整数据流原图 → VAE Encoder → 4通道潜图 z₀ → 扩散前向加噪 → 带噪潜变量 z_t → DiT Transformer 预测噪声 → 反向还原干净潜图 z₀ → VAE Decoder 解码 → 输出RGB像素图像条件说明原版DiT为类别条件生成条件仅包含两项时间步ttt代表当前潜图噪声强度类别标签yyyImageNet数据集数字类别ID0~999无文本、无CLIP编码器不支持自然语言Prompt。二、DiT 完整整体架构2.1 五大核心组件PatchEmbedPatch分块嵌入输入是VAE输出的二维潜图如32×32×4使用卷积等价线性层将固定尺寸Patch论文标准patch_size2映射为固定维度Token序列随后叠加固定2D正弦余弦位置编码不可学习给每个Patch注入空间位置信息。示例32×32潜图、patch2 → 1024个图像Token。条件编码模块TimestepEmbedder LabelEmbedder分别处理时间步ttt、类别标签yyy输出同维度向量后相加融合得到全局一维条件向量ctembyembc t_{emb} y_{emb}ctemb​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​自注意力支路流程自适应归一化公式hLN(x)⋅(1γ)βh \text{LN}(x) \cdot (1\gamma) \betahLN(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​)xoutxinα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\gamma0,\beta0γ0,β0代入公式hLN(x)⋅(10)0h\text{LN}(x)\cdot(10)0hLN(x)⋅(10)0自适应LN退化为标准无参数LayerNormα0\alpha0α0残差支路输出全部归零整层DiT Block输入特征输出特征等价恒等映射。深层网络梯度可无衰减逐层传递彻底解决深层Transformer梯度消失、训练震荡、难以收敛的问题。渐进式学习逻辑训练迭代中参数逐步更新门控系数α\alphaα从0缓慢激活残差支路模型从简单恒等映射逐步学习复杂图像特征映射关系收敛速度大幅提升。四、DiT 模型规格变体论文提供4种标准尺寸命名规则DiT-XX/p后缀/p代表Patch尺寸模型Transformer层数隐藏层维度注意力头数总参数量DiT-S/212384633MDiT-B/21276812130MDiT-L/224102416458MDiT-XL/228115216675MPatch尺寸消融实验Patch越小图像Token数量越多模型捕捉细节能力越强FID越低但算力、显存占用指数级上升Patch尺寸Token总数单轮推理GFLOPsImageNet FID8640.924.842562.320.1210245.719.5五、训练与推理流程5.1 训练阶段DDPM扩散训练目标图像送入VAE编码器压缩为4通道潜图z0z_0z0​随机采样时间步ttt根据扩散公式给z0z_0z0​添加噪声得到带噪潜图ztz_tzt​采样类别标签yyy编码融合得到全局条件cccDiT网络输入zt、t、yz_t、t、yzt​、t、y预测潜图中添加的噪声损失函数为预测噪声与真实噪声的MSE反向传播更新DiT全部权重支持CFGClassifier-Free Guidance训练10%概率随机丢弃类别标签学习无条件生成推理时提升画面细节。5.2 推理生成阶段采样初始化纯高斯噪声潜图设定目标生成类别yyy全程固定不变仅计算一次条件向量ccc从最大时间步逐步反向采样每一步送入DiT预测当前噪声结合CFG权重融合条件/无条件噪声预测逐步去除潜图噪声采样完成后得到干净潜图送入VAE解码器输出最终图像。六、DiT 对比传统U-Net扩散模型对比维度U-NetSD/DDPMDiT原版分类条件骨干架构多层卷积、上下采样、跳跃连接纯TransformerPatch Token序列多头自注意力无卷积、无跳跃连接空间归纳偏置极强卷积天然捕捉局部像素关系极弱无内置空间先验完全依靠数据学习图像结构条件注入方式时间嵌入叠加、交叉注意力融合条件adaLN-Zero自适应归一化调制无额外注意力层算力开销更低模型缩放能力缩放收益有限参数量提升画质增长缓慢存在性能瓶颈完美遵循缩放定律模型越大算力越高画质持续稳定提升条件类型可扩展文本、类别等多类条件原版仅支持ImageNet数字类别标签无文本Prompt、无CLIP编码器训练稳定性深层网络易震荡、梯度消失adaLN-Zero零初始化实现恒等映射深层训练收敛稳定七、工程落地关键规范Patch尺寸选型论文标准折中配置patch_size2兼顾画质与显存占用位置编码规范必须使用固定2D正弦余弦位置编码requires_gradFalse不参与梯度更新区别于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架构。