Diffusion Transformer为什么生成模型开始从 U-Net 转向 TransformerDiffusion TransformerDiT最容易被误解成一句话把扩散模型里的 U-Net 换成 Transformer。这个说法没错但太粗。真正值得看的是它把图像生成问题重新写成了一个更像语言模型的扩展问题输入被切成 token主干网络按宽度、深度、token 数量扩展质量随计算量可预测地上升。这也是 DiT 后来被很多模型路线吸收的原因。DiT 本身不是一个完整的文生图系统它更像一次架构实验如果扩散过程还在VAE 还在采样器还在只把去噪网络从卷积 U-Net 换成 ViT 风格的 Transformer会发生什么Peebles 和 Xie 在 ICCV 2023 的答案很直接在 ImageNet 类条件生成上DiT-XL/2 达到 256×256 FID 2.27并且模型 Gflops 和 FID 呈明显负相关。这篇文章按机制讲不把 DiT 写成一堆名词。我们先看扩散模型里主干网络到底在做什么再看 DiT 怎么把 latent patch 变成 token为什么 adaLN-Zero 是关键改动最后讨论它和 Stable Diffusion 3 / PixArt / Sora 这条路线的关系。扩散模型里的主干网络到底负责什么扩散模型训练时会把干净样本x0x_0x0​逐步加噪得到xtx_txt​然后让神经网络学习从带噪样本里预测噪声或反向过程参数。常见的简化训练目标可以写成Lsimple(θ)∥ϵθ(xt,t,c)−ϵ∥22 L_{simple}(\theta)\|\epsilon_\theta(x_t,t,c)-\epsilon\|_2^2Lsimple​(θ)∥ϵθ​(xt​,t,c)−ϵ∥22​这里ϵθ\epsilon_\thetaϵθ​就是主干网络。它接收三类东西带噪图像或 latent、时间步ttt、条件ccc。条件可以是类别标签也可以是文本、图像、视频帧或其他控制信号。早期高质量扩散模型几乎都用 U-Net 做这个网络。原因很现实U-Net 天生适合图像网格有多尺度下采样/上采样有 skip connection卷积的局部归纳偏置也很强。ADM、LDM、Stable Diffusion 早期版本都沿着这条路走。DiT 的问题不是「U-Net 能不能用」。U-Net 当然能用而且非常强。DiT 问的是另一个问题当生成模型的质量开始被规模支配时卷积 U-Net 是不是最容易扩展的主干Transformer 在语言、视觉识别、多模态上都已经证明过 scaling 的价值那么扩散模型的去噪网络能不能吃到同样的红利这个问题一旦成立架构选择就从「哪个模块更像图像」变成了「哪个模块在计算量变大时更可预测」。这也是 DiT 的技术张力。DiT 没有在像素空间硬做 TransformerDiT 用的是 latent diffusion 框架而不是直接在像素空间训练 Transformer。先用预训练 VAE 把图像压到 latent 空间再在 latent 上做扩散。论文里 256×256×3 图像经过 VAE encoder 后变成 32×32×4 的 latent。去噪网络处理的是这个 latent不是原图像素。流程大致是image x ↓ VAE encoder E latent z: 32×32×4 ↓ add noise at timestep t noised latent z_t ↓ patchify into latent tokens Transformer denoiser DiT ↓ predict noise and covariance sampled latent z ↓ VAE decoder D image这个选择很重要。Transformer 的注意力成本随 token 数增长很快如果直接对高分辨率像素切 patch成本会很难看。latent diffusion 先把空间尺寸压小DiT 再把 latent 切成 patch整体才有实验可行性。DiT 的 patchify 层和 ViT 很像把I×I×CI \times I \times CI×I×C的 latent 按p×pp \times pp×p切块每个块线性映射成一个 token。token 数是T(I/p)2 T(I/p)^2T(I/p)2如果I32I32I32patch sizep4p4p4就有 64 个 token如果p2p2p2token 数变成 256。论文中特别强调把 patch size 减半会让 token 数变成 4 倍Gflops 至少变成 4 倍但参数量几乎不变。这解释了 DiT 里一个容易被忽略的点模型质量不只由参数量决定。你可以不怎么增加参数只通过更小 patch 让模型处理更多 token计算量上去FID 也会下降。对生成模型来说这比单纯报参数量更有信息量。DiT block 真正改的是条件注入方式如果只看主干DiT 尽量贴近标准 ViTpatch embedding、位置编码、多层 Transformer block、最后把 token decode 回 latent 的空间布局。扩散模型麻烦的地方在条件。去噪网络必须知道当前时间步也要知道类别或文本条件。论文比较了几种注入方式条件注入方式做法代价与效果in-context conditioning把时间步和类别 embedding 当成额外 token 拼到序列里改动最小额外 Gflops 很少但效果较弱cross-attention图像 token 自注意力后再对条件 token 做 cross-attention表达力强但论文估计带来约 15% Gflops 开销adaLN用时间步和类别 embedding 回归 LayerNorm 的 scale / shift计算最省条件作用在归一化层上adaLN-Zero在 adaLN 基础上增加残差前的缩放并把残差初始化为 0论文中效果最好也是最终采用的 blockadaLN-Zero 是 DiT 论文里最有价值的小改动。它不是「Transformer 加个条件」这么简单而是把每个 DiT block 初始成近似恒等函数。具体说block 会从条件 embedding 回归 LayerNorm 的调制参数同时回归残差分支前的缩放参数α\alphaα这些α\alphaα初始输出为 0所以训练开始时残差分支基本不扰动输入。这个设计和扩散 U-Net 里的零初始化残差有相似直觉扩散训练一开始已经够不稳定了尤其是大模型。让每个 block 初始时接近 identity可以让优化更顺。论文在 DiT-XL/2 上比较四种 blockadaLN-Zero 在训练过程中一直优于 cross-attention 和 in-context conditioning并且计算开销很低。我觉得这里有个判断挺关键DiT 的成功不是「Transformer 可以直接替代一切」。如果直接把条件随便塞进 token 序列结果并不好。它真正证明的是扩散模型需要一套适合 timestep/condition 的 Transformer block而 adaLN-Zero 是当时最干净的一种实现。模型大小和 patch size 是两根独立的旋钮DiT 的命名方式是DiT-XL/2这种格式。XL表示 Transformer 配置大小/2表示 latent patch size。论文扫了四种模型大小和三种 patch size一共 12 个模型。模型层数hidden sizeheadsp4 时 GflopsDiT-S1238461.4DiT-B12768125.6DiT-L2410241619.7DiT-XL2811521629.1这张表只对应I32,p4I32,p4I32,p4。如果改成p2p2p2Gflops 会大幅上升。官方 README 给出的 DiT-XL/2 256×256 版本是 119 Gflops512×512 版本是 525 Gflops。两个旋钮的含义不同增大 S/B/L/XL增加深度、宽度、heads参数量和每 token 表达能力都上升。减小 patch size增加 token 数让模型看到更细粒度的 latent patch参数量变化不大但注意力计算变多。论文里最有意思的结论是FID 和 Gflops 的关系比 FID 和参数量的关系更清楚。保持模型大小不变减小 patch size参数几乎不变但质量会上升不同配置如果 Gflops 接近FID 也会接近。这对后来的视频/多模态生成很有启发当数据能被统一成 token模型扩展就可以围绕 token 数和计算量设计。为什么 DiT-XL/2 的结果当时很有说服力DiT-XL/2 在 ImageNet 256×256 类条件生成上继续训练 7M steps 后使用 classifier-free guidance 达到 FID-50K 2.27。作为对比论文表 2 中 LDM-4-G 在 cfg1.50 时是 3.60ADM-GADM-U 是 3.94。512×512 上DiT-XL/2-G cfg1.50 达到 FID 3.04而 ADM-GADM-U 是 3.85。更值得看的其实是计算量对比。项目页给出 256×256 下 LDM-4 是 103 GflopsADM-U 是 742 GflopsDiT-XL/2 是 119 Gflops512×512 下 ADM-U 是 2813 Gflops而 DiT-XL/2 是 525 Gflops。DiT 不是低成本模型但它在 latent 空间里用 Transformer 做到了相当清楚的 compute-quality tradeoff。官方 README 里的预训练权重表也给了两个常用 checkpointcheckpoint分辨率FID-50KInception ScoreGflopsDiT-XL/2256×2562.27278.24119DiT-XL/2512×5123.04240.82525这些数字的使用要小心。DiT 论文是在 class-conditional ImageNet 上比较不等价于今天的开放域文生图体验FID 也受采样器、VAE、评估实现影响。论文为了可比性用 ADM 的 TensorFlow evaluation suite 计算 FID并注明主要结果采用 250 DDPM sampling steps。如果把这些限制拿掉只说「DiT 比 U-Net 强」就会过度简化。更准确的说法是在这个基准和训练设置下DiT 证明了 Transformer 主干可以替代 U-Net并且随计算量扩展时表现很稳。classifier-free guidance 仍然在起作用DiT 不改变 classifier-free guidanceCFG的基本逻辑。训练时随机丢弃条件学习一个 null condition采样时同时跑有条件和无条件预测再按 guidance scale 拉开ϵ^θ(xt,c)ϵθ(xt,∅)s⋅(ϵθ(xt,c)−ϵθ(xt,∅)) \hat{\epsilon}_\theta(x_t,c)\epsilon_\theta(x_t,\varnothing)s\cdot(\epsilon_\theta(x_t,c)-\epsilon_\theta(x_t,\varnothing))ϵ^θ​(xt​,c)ϵθ​(xt​,∅)s⋅(ϵθ​(xt​,c)−ϵθ​(xt​,∅))s1s1s1时基本回到普通条件采样s1s1s1会更强调条件一致性。DiT 论文中的最佳 ImageNet 数字依赖 CFG。例如 256×256 的 DiT-XL/2 不带 guidance 时 FID 是 9.62cfg1.25 时是 3.22cfg1.50 时是 2.27。这说明 DiT 改的是 denoiser 架构不是把扩散模型的其他组件全部推翻。VAE、噪声日程、CFG、采样步数这些变量仍然重要。后来的很多系统也是这样Transformer 主干成为中心但整体质量来自 tokenizer / VAE、训练数据、条件编码器、采样公式和后训练技巧的组合。从 DiT 到文生图问题从类别条件变成多模态条件原始 DiT 主要做 ImageNet 类条件生成条件只是一个类别 label。真正进入文生图系统后条件会变成文本 token甚至多个文本编码器的输出。这里就出现了两条演化路线。第一条是 PixArt-α 这种路线。PixArt-α 明确把 DiT 用到 text-to-image并强调训练成本。它不是只换主干还结合了高质量 caption、预训练 T5 文本编码器、弱到强训练策略等工程选择。它的价值在于证明 DiT 不只适合 ImageNet 类条件任务也能进入开放文本条件生成。第二条是 Stable Diffusion 3 的 MMDiT。SD3 论文标题是「Scaling Rectified Flow Transformers for High-Resolution Image Synthesis」核心组合是 rectified flow transformer backbone。它提出 Multimodal Diffusion TransformerMMDiT用不同权重处理图像和文本表示再让它们通过注意力交互。这个设计比原始 DiT 更适合文生图因为文本不是一个简单 label而是一串语义 token。如果说 DiT 的问题是「图像 latent token 怎么被 Transformer 去噪」MMDiT 的问题就是「图像 token 和文本 token 怎么在同一个生成过程里对齐」。这也是后续 FLUX、SD3 类模型常被放在 DiT 家族里讨论的原因。Sora 的公开技术报告也把 diffusion transformer 放在中心位置。OpenAI 的表述是通过统一视觉数据表示可以在不同持续时间、分辨率和宽高比的数据上训练 diffusion transformers。这里的重点已经从 2D 图像扩展到了时空 patch视频可以被压缩、切块、token 化然后用 Transformer 做大规模生成建模。DiT 对视频生成尤其自然视频生成比图像生成更像 DiT 的主场。原因不是 Transformer 有什么魔法而是视频天然需要长程依赖同一个物体要跨帧一致人物动作要连续镜头运动要合理短时间纹理闪烁也不能太严重。U-Net 可以通过 3D convolution、temporal attention、motion module 去补时间建模但这些设计经常显得拼装。DiT 风格的路线更直接把视觉内容编码成时空 token让注意力在空间和时间上建模。代价是 token 数爆炸所以系统必须在 VAE 压缩率、patch size、attention pattern、训练分辨率之间做非常精细的权衡。这也解释了为什么「patch size」在 DiT 里不是小超参。图像里p2p2p2已经会把 token 数推上去视频里还多了时间维度。token 变多会提高表达粒度也会迅速吃掉训练和推理预算。视频 DiT 的很多工程问题本质上是在问哪些 token 必须全局交互哪些可以局部或分块处理从这个角度看DiT 的意义不只是替代 U-Net。它提供了一套把生成问题写成 token scaling 的语言。语言一旦统一图像、视频、多视角、3D 表示就更容易被放进同一个架构框架里。复现 DiT 时最容易踩的坑官方 PyTorch 仓库可以直接采样预训练模型gitclone https://github.com/facebookresearch/DiT.gitcdDiT condaenvcreate-fenvironment.yml conda activate DiT python sample.py --image-size512--seed1训练 ImageNet 256×256 的 DiT-XL/2 示例命令是torchrun--nnodes1--nproc_per_nodeN train.py--modelDiT-XL/2 --data-path /path/to/imagenet/train但复现论文数字不是跑通脚本这么简单。几个细节会影响结果论文原始训练用 JAX/TPUPyTorch 仓库提供的是复现训练脚本和移植权重。README 里注明 PyTorch 训练结果和 JAX 结果在 400K steps 附近接近但会有随机波动。FID 评估使用 ADM 的 TensorFlow evaluation suite换评估实现可能得到不完全一致的数。A100 上 TF32 会明显影响训练/采样速度也可能带来轻微数值差异。官方仓库 README 提到 FlashAttention、torch.compile、AMP/bfloat16、gradient checkpointing 等都不是最初实现里的完整默认特性。如果只是学习架构最建议先读models.py看 patchify、timestep embedding、label embedding、adaLN-Zero block 和 final layer 如何拼起来。先理解 forward pass再考虑训练。直接从大规模训练开始很容易把问题误判成「模型不行」其实只是数据、评估或混合精度设置没对齐。DiT 的局限也要说清楚DiT 不是免费的扩展。Transformer 的注意力成本让它对 token 数非常敏感。patch size 越小质量可能越好但推理成本也会迅速增加。文生图和视频模型还要引入长文本 token、时序 token、多分辨率训练这些都会让系统设计更复杂。DiT 也没有自动解决数据和对齐问题。原始 DiT 在 ImageNet 类条件生成上很漂亮但开放域文生图需要文本-图像对、caption 质量、偏好对齐、安全过滤、风格泛化。PixArt-α、SD3、Sora 这些系统之所以强不只是因为用了 Transformer 主干。还有一个容易被忽略的问题U-Net 的局部偏置并非坏事。对小模型、小数据、低预算任务卷积结构可能更省、更稳。DiT 的优势通常在规模上显现当训练预算不足时Transformer 未必是最划算的选择。一个更实际的判断如果你把 DiT 当成「U-Net 换 Transformer」会错过它真正的贡献。DiT 把扩散模型主干带进了一个更标准的 scaling 叙事latent 被 token 化模型宽深可扩patch size 控制 token 数Gflops 成为比参数量更直接的质量指标。这就是为什么它影响了后来的文生图和视频生成。不是每个系统都照搬 DiT-XL/2也不是每个模型都用同一种 block。但越来越多生成模型开始围绕 token、attention、flow/diffusion objective 和多模态条件编码来设计。DiT 是这条路线上很清楚的起点。如果要用一句工程化的话收尾DiT 让扩散模型的主干从「图像网络」变成了「可扩展 token 处理器」。这个变化没有替你解决所有问题但它把生成模型带到了 Transformer 最擅长的地盘上。参考资料William Peebles, Saining Xie,Scalable Diffusion Models with Transformers, ICCV 2023 / arXiv:2212.09748https://arxiv.org/abs/2212.09748DiT official project pagehttps://www.wpeebles.com/DiTfacebookresearch/DiT official PyTorch implementationhttps://github.com/facebookresearch/DiTHugging Face Diffusers DiT pipeline docshttps://huggingface.co/docs/diffusers/api/pipelines/ditPatrick Esser et al.,Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, arXiv:2403.03206https://arxiv.org/abs/2403.03206Junsong Chen et al.,PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, arXiv:2310.00426https://arxiv.org/abs/2310.00426OpenAI,Sora: Creating video from texthttps://openai.com/index/sora/