DiT（Diffusion Transformer）：把 U-Net 换成 Transformer，扩散模型迎来架构换代

如果你用过 Stable Diffusion，你大概知道它背后那个叫 U-Net 的东西——一个长得像字母 U 的卷积网络，带一堆跳跃连接，被誉为扩散模型的大脑。过去几年，这个架构几乎成了文生图领域的标配，大家都觉得：搞扩散模型，必须用 U-Net。

但 2023 年，一篇论文把这个共识砸了个粉碎。它把 U-Net 整个拆掉，换成了一个纯 Transformer，然后发现：不仅效果没降，反而在可扩展性上把 U-Net 按在地上摩擦。

这篇论文就是 Scalable Diffusion Models with Transformers，作者 William Peebles 和 Saining Xie 给这个新架构起了个名字——DiT（Diffusion Transformer）。后来 OpenAI 的 Sora 背后，也用到了和 DiT 高度相似的技术路线，直接把扩散模型送进了视频生成时代。

我最早看到这篇论文时，第一反应是：“这不就是 ViT 嘛，换了个马甲？” 但仔细读完后，我才意识到，它不只是换了个 backbone，而是重新定义了扩散模型该怎么设计。下面我就把我理解的过程全部分享给你，包括我踩过的坑、想通的关键点，以及 DiT 到底厉害在哪、什么时候会拉胯。

U-Net 到底做错了什么？

其实 U-Net 没做错什么。它能在扩散模型里称霸，是因为它天生适合处理图像去噪任务：下采样路径提取全局语义，上采样路径恢复细节，跳跃连接保留高分辨率信息。而且它很省参数，计算效率高。

但问题在于：U-Net 很难规模化。你想把模型做大，无非堆更多的卷积层、加更多的通道数。可卷积层是局部的，感受野增长慢，参数利用率低。而且 U-Net 的结构是精心设计的，改动起来很麻烦——你没法像堆 Transformer 层那样，直接说“加 12 层试试”。

更关键的是，扩散模型在潜在空间里工作（比如 LDM），输入已经是压缩后的特征图，不再是像素空间。这时候，卷积的归纳偏置（局部性、平移等变性）的重要性已经大幅下降，反而是 Transformer 的全局建模能力开始变得更有吸引力。

换句话说，U-Net 的强项在潜在空间里被削弱了，而 Transformer 一直被诟病的“缺少归纳偏置”反而成了优势——它不需要硬编码的局部性，可以自由学习任意全局关系。

DiT 是怎么把 Transformer 插进扩散模型的？

DiT 的设计思路简单到让人有点不安：把一个 ViT 结构的 Transformer 直接扔进潜在扩散模型的去噪网络里，然后想办法把时间步和类别条件喂给它。

具体来说，DiT 的输入是一张潜在空间的特征图（比如 32×32×4）。它先做 patch embedding：把特征图切成 2×2 的小块，每个块展平成一个向量，然后加上位置编码，扔进一串 Transformer 块。最后再通过一个线性层把 patch 拼回特征图，输出预测的噪声。

整个过程和 ViT 几乎一模一样，唯一的区别在于：它需要把扩散模型的条件（时间步 t 和类别标签 c）注入 Transformer 块。DiT 论文里试了四种注入方式：

• In-context conditioning：把条件当成额外的 token 拼在 patch 序列里，类似 ViT 的 class token。
• Cross-attention：在 Transformer 块里插入一个交叉注意力层，用条件作为 query 去 attend patch 序列。
• Adaptive layer norm (adaLN)：用条件去回归层归一化的缩放和偏置参数，类似 StyleGAN 的做法。
• adaLN-Zero：在 adaLN 的基础上，所有残差连接初始化为零，让网络一开始相当于恒等映射，训练更稳定。

实验表明，adaLN-Zero 效果最好，而且是所有方案里计算开销最小的。这个结论很重要，因为扩散模型在推理时要做很多步去噪，每一层多一点计算，整体就会慢很多。

我一开始看到 adaLN-Zero 时觉得很奇怪：为什么把残差初始化为零能提升效果？后来想通了：扩散模型去噪是个渐进过程，一开始的噪声很大，模型输出的修正应该很小，让网络慢慢学会“小步修改”。零初始化恰好给了这个归纳偏置，避免模型一上来就乱改。

一张表说清 U-Net 和 DiT 的根本区别

这张表透露了一个关键信息：当模型规模大到一定程度，架构的“可扩展性”比“精巧设计”重要得多。DiT 就是靠这个逻辑，在 ImageNet 256×256 生成上，用 1.2B 参数的模型把 FID 干到了 2.27，刷新了当时的最好成绩。

可扩展性：DiT 真正的杀手锏

DiT 论文里最让我震撼的，不是它的生成质量，而是那张可扩展性曲线。作者把 DiT 的模型规模从 33M 一直拉到 675M 参数，同时控制训练计算量（Gflops），发现：更大的模型在相同计算量下，生成质量更好；而且随着计算量增加，大模型的优势越来越明显。

这意味着，只要你有足够的算力，DiT 可以持续变强，不会像 U-Net 那样碰到某个瓶颈就不再涨了。这种“算力即正义”的特性，正是 Transformer 统治 NLP 的核心原因，现在扩散模型也沾光了。

不过这里有个坑：DiT 在低计算量下并不比 U-Net 强。如果你用的是小模型（比如 33M），U-Net 可能因为卷积的归纳偏置，收敛更快，FID 更低。DiT 的优势要等到模型参数量超过 200M 才会显现。所以，如果你只是想在单个 GPU 上跑个小模型玩玩，DiT 可能不是最优选。

Sora 和 DiT 是什么关系？

很多人以为 Sora 就是 DiT 的视频版，其实不完全对。Sora 的技术报告里提到，它使用了一种基于 Transformer 的扩散模型架构，在时空潜在 patch 上操作，这和 DiT 的思想一脉相承。但 Sora 做了很多视频特有的设计，比如时空 patch 化、更长的上下文、更复杂的条件控制等。

可以这样理解：DiT 证明了 Transformer 可以在扩散模型里替代 U-Net，并且可扩展性极好；Sora 则把这个结论推到了极致，用海量数据和算力，做出了前所未有的视频生成效果。DiT 是那个“点燃引信”的学术工作，而 Sora 是工业界投下的炸弹。

我踩过的两个认知坑

坑一：我以为 patch 大小不重要，结果被打脸。DiT 论文里有一个实验：把 patch size 从 2 改成 4、8，发现 patch 越小，生成细节越丰富，但序列长度平方级增长，计算量暴增。patch 越大，速度快但细节丢失。这其实是一个 计算-质量权衡，和 ViT 在图像分类里的结论一致。但扩散模型对细节更敏感，所以 patch size 的选择比分类任务更关键。

坑二：我一开始以为 DiT 不需要位置编码，因为 U-Net 的全卷积结构自带位置信息。但 DiT 是纯注意力，没有位置编码模型根本不知道 patch 在哪。论文用了正弦位置编码，后来还有工作改用可学习位置编码或 RoPE，效果更好。位置编码是 DiT 的“隐形地基”，做不好直接塌房。

三个最常见的问题

DiT 能直接替代 Stable Diffusion 里的 U-Net 吗？

理论上可以，实际上已经有人做了。比如 Stability AI 的 SD3 就采用了 DiT 架构。但要注意，DiT 对算力要求更高，在小分辨率下不一定比 U-Net 快，而且需要重新训练，不是即插即用。

DiT 为什么不用卷积？能不能加一点卷积？

可以加。已经有工作把卷积和注意力混合，比如把 patch embedding 换成卷积、或者在 Transformer 块里加深度卷积，这样在小模型上也能获得不错的性能。但 DiT 原版坚持纯 Transformer，是为了验证“纯注意力也能做到最好”，是一种极致的学术探索。

DiT 的推理速度是不是比 U-Net 慢很多？

在相同参数规模下，DiT 的推理确实比 U-Net 慢，因为自注意力的计算量随序列长度平方增长。但 DiT 的优势在于大规模下可以通过更少的去噪步数（如 DDIM 采样）找回速度，而且模型越大，生成质量提升越明显，用稍微慢一点换质量飞跃，很多人觉得是划算的。

DiT 的局限，以及什么时候它会让你失望

首先，训练 DiT 真的贵。论文里那个 675M 参数的模型，用了 256 块 V100 训练了 7 天。普通开发者想复现，没个几十万预算很难。即便是 300M 参数，也需要 8 块 A100 跑上好几天。所以当前 DiT 更多是工业界和土豪实验室的玩具。

其次，DiT 对数据质量要求更高。因为没有卷积的归纳偏置，它需要从数据里自己学空间关系，数据量不够或者噪声太多，模型容易学歪。这也是为什么 DiT 论文只在 ImageNet 上做实验，放到更复杂的自然图像上，需要更大的数据集（如 LAION-5B）和更长的训练。

最后，DiT 的推理延迟在长序列上是个硬伤。视频生成把时间维度也拉进 patch 序列后，序列长度轻松上万，自注意力的平方复杂度会让人想哭。Sora 能扛住，是因为 OpenAI 用了大量工程优化（比如 FlashAttention、稀疏注意力等），但这不是每个团队都能做到的。

但话说回来，DiT 把扩散模型从“手工设计卷积架构”的泥潭里拉了出来，推进了“算力驱动”的范式。从此以后，你再也不用纠结 U-Net 的通道数该设多少、下采样倍数怎么调，而是像搭大语言模型一样，直接堆层数、扩宽度，然后让算力说话。这才是 DiT 真正的历史意义。