Abstract
论文的核心思想非常直接:用一个标准的 Transformer 架构替换掉扩散模型中常用的 U-Net 主干网络,并证明这种新架构(称为 DiT, Diffusion Transformer)具有出色的可扩展性(Scalability)。
Background & Motivation
在论文发表前,Transformer 已经在自然语言处理(BERT, GPT)和计算机视觉(ViT)等领域取得了巨大成功,成为了一种'统一'的架构。然而,在图像生成领域,特别是扩散模型中,大家仍然普遍使用 U-Net。U-Net 因其多尺度特征融合和卷积的局部归纳偏置而被广泛采用。
在深度学习中,一个好的架构应该具备良好的'可扩展性'——即投入更多的计算资源(更大的模型、更多的数据),性能应该会持续稳定地提升。ViT 已经证明了 Transformer 在视觉识别任务上具有这种特性。作者们希望验证 DiT 是否也具备这种优良特性,为未来的生成模型发展指明一条清晰的路径。
并且,作者用 Gflops 而不是参数量来衡量和讨论模型复杂度。因为参数量在图像任务中有严重缺陷,无法反映真实计算成本。
Diffusion Transformers
3.1. Preliminaries
Diffusion formulation:把一张真实的、干净的图像 x0,通过 T 个步骤,逐渐地、一点一点地加入高斯噪声,直到它变成一个完全的纯噪声图像 xT。模型(在 DiT 中就是那个 Transformer)的任务是预测噪声。具体来说,给定加噪图像 xt 和时间步 t,模型需要预测出在第 t 步加入的那个噪声 εt。这个预测出的噪声被记为 εθ(xt)。之后主要用简单的 MSE 损失来训练噪声预测网络,同时用一个更完整的损失项(完整的 DKL)来训练协方差的预测。
Classifier-free guidance:这部分介绍了一种非常重要的技术,用于在生成过程中加强条件(比如类别标签 c)对生成结果的影响,从而显著提高生成图像的质量和与条件的匹配度。
在生成每一步时,我们不仅要预测'在有条件 c 的情况下的噪声',还要预测'在没有任何条件(用一个特殊的空 embedding ∅ 表示)的情况下的噪声'。然后,将这两个预测结果进行线性组合。
最终的噪声预测 ε̂θ 是通过一个公式来计算的:
ε̂θ(xt, c) = εθ(xt, ∅) + s * (εθ(xt, c) - εθ(xt, ∅))
其中:
- εθ(xt, ∅):无条件下的噪声预测(模型自由发挥)。
- εθ(xt, c):有条件 c 下的噪声预测(模型听从指令)。
- (εθ(xt, c) - εθ(xt, ∅)):这个差值可以理解为'从自由发挥到听从指令'的方向。
- s:引导尺度(guidance scale),s > 1。这个参数控制了我们要在多大程度上'强调'这个方向。s 越大,生成图像与条件 c 的相关性就越强,通常视觉效果也更锐利,但可能会牺牲多样性。
CFG 被广泛证明能极大提升生成样本的质量,DiT 也不例外。
Latent Diffusion Models, LDM:DiT 论文明确指出,他们采用了 LDM 框架,即在一个卷积 VAE 提供的潜在空间上,应用了他们提出的 Transformer 架构。
这三个预备知识点,层层递进地构建了 DiT 的完整工作流程:
先用 LDM 框架将图像压缩到低维潜在空间,DiT(作为扩散模型的核心)在这个潜在空间中,遵循扩散模型原理,学习如何从噪声中恢复出潜在表示,最后,在生成时,使用无分类器引导(CFG)技术来提升生成质量和可控性。
3.2. Diffusion Transformer Design Space

Patchify:
- 将 z 分割成一个个不重叠的小块(patch),每个小块的大小为 p x p。例如,如果 p=2,那么一个 32x32 的 z 就会被切分成 (32/2) * (32/2) = 16 * 16 = 256 个小块。
- 线性嵌入:每个小块被展平并通过一个线性层,映射成一个维度为 d 的向量,这个向量就叫做一个 token。



