DiT（Diffusion Transformer）：颠覆图像生成的Transformer革新

原创于 2025-09-30 10:25:35 发布 · 1.3k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#transformer #深度学习 #人工智能

部署运行你感兴趣的模型镜像

摘要

Diffusion Transformer（DiT） 是近年来计算机视觉领域最具突破性的架构之一，它成功地将Transformer架构引入扩散模型，取代了传统的U-Net主干网络。DiT不仅继承了Transformer的卓越扩展性和训练稳定性，还在ImageNet 256×256和512×512生成任务上取得了2.27 FID的顶尖结果，显著超越了此前基于U-Net的扩散模型。本文将深入剖析DiT的核心原理、架构设计、关键技术以及未来发展趋势，为读者提供全面理解DiT技术内涵和实践应用的权威参考。

1 引言：从U-Net到Transformer的范式转移

1.1 扩散模型的发展历程

扩散模型作为当前图像生成领域的核心架构，其基本思想是通过前向加噪过程和反向去噪过程学习数据分布。前向过程逐步将噪声应用于真实数据，而反向过程则试图从噪声中重建原始数据。在DiT出现之前，几乎所有扩散模型都采用基于卷积架构的U-Net作为核心骨干网络。

U-Net凭借其编码器-解码器结构中的跳跃连接，能够有效融合多尺度特征，成为扩散模型的默认选择。然而，U-Net固有的归纳偏置和局部性特点，也限制了其在捕获长距离依赖关系和大规模预训练方面的潜力。

1.2 Transformer的跨领域崛起

Transformer架构自2017年在NLP领域提出后，迅速席卷了各个机器学习领域。在自然语言处理中取得巨大成功后，Transformer通过Vision Transformer（ViT）等模型进一步征服了计算机视觉领域。Transformer的自注意力机制能够全局捕捉序列元素间的依赖关系，加上其卓越的扩展特性（模型规模与性能呈幂律关系），使其成为处理大规模数据的理想选择。

然而，在DiT之前，Transformer在生成式模型特别是扩散模型中的应用仍属空白。这一架构鸿沟激发了研究者的探索欲望：能否用Transformer彻底替换U-Net，构建一个纯粹的Transformer扩散模型？

2 DiT核心技术原理

2.1 扩散模型基础

扩散模型的核心数学原理建立在两个过程上：前向噪化过程和反向去噪过程。

前向过程可以表述为：
$q(x_t|x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I)$
其中$ \beta_t $是噪声调度参数，$ x_t $是第t步的噪声图像。

通过重参数化技巧，可以直接从$ x_0 $计算$ x_t $：
$x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon$
其中$ \alpha_t = 1 - \beta_t $，$ \bar{\alpha}t = \prod{i=1}^{t}\alpha_i $，$ \epsilon \sim \mathcal{N}(0, I) $。

反向过程则需要学习一个去噪模型：
$p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$
模型需要预测给定当前噪声图像$ x_t $和时间步t时，上一时间步的图像分布。

2.2 潜在扩散模型（LDM）

直接在高分辨率像素空间训练扩散模型计算成本极高。潜在扩散模型（LDM） 通过两阶段方法解决这个问题：

感知压缩：使用自编码器（VAE）将图像压缩到更小的潜在空间
潜在扩散：在潜在空间而非像素空间训练扩散模型
重建解码：从扩散过程采样后，使用解码器将潜在表示恢复为图像

这种方法大幅降低了计算需求，使训练更高分辨率的模型成为可能。

2.3 DiT架构设计

DiT架构的核心思想是将Transformer应用于扩散过程的潜在空间。其整体架构可以通过以下图表展示：

2.3.1 Patchify与线性嵌入

DiT的第一步与ViT相同，将输入的空间表示转换为一系列token。对于大小为$ H \times W \times C $的潜在图像， D i T 将其分割为$ P \times P $的 p a t c h es ，然后将每个 p a t c h 展平并线性投影到隐藏维度 D 中。这样得到的 t o k e n 数量$ T $为：
$\frac{H}{P} \times \frac{W}{P}$

patch大小$ P $对模型计算量有显著影响。减半$ P $会使$ T $变为 4 倍，从而让计算量也变为 4 倍。尽管$ P $对GFLOPs有显著影响，但参数量没有实质变化。

2.3.2 DiT Blocks设计

在patchify之后，输入token会经过一系列Transformer Block。DiT探索了四种不同类型的条件调节机制，这些设计都对标准ViT Block进行了微小修改：

In-Context Conditioning：将时间步$ t $和类标签$ c $作为额外token附加到输入序列中，视为与图像token无区别。
Cross-Attention Block：将$ t $和$ c $的embedding连接成一个长度为2的序列，与图像token序列分开，在Transformer Block中添加交叉注意力层。
Adaptive Layer Norm (adaLN)：不直接学习缩放和移位参数$ \gamma $和$ \beta $，而是从$ t $和$ c $的embedding中预测它们。
adaLN-Zero Block：在adaLN基础上，对每个Block的第一个LayerNorm的缩放因子进行零初始化，以加速训练。

实验表明，adaLN-Zero Block在计算效率和性能上达到了最佳平衡，成为DiT最终采用的方案。

3 DiT的关键技术优势

3.1 卓越的扩展特性

DiT最重要的优势在于其优异的扩展性。随着模型规模（GFLOPs）的增加，DiT的性能持续提升且未出现饱和迹象。研究者通过控制不同规模配置验证了这一特性：

表：DiT模型规模配置与性能关系

模型配置	Transformer Blocks	Hidden Size	Heads	Params (M)	GFLOPs	FID-50K
DiT-S	12	384	6	33	25.6	68.4
DiT-B	12	768	12	130	103.6	43.5
DiT-L	24	1024	16	458	106.4	19.2
DiT-XL	28	1152	16	675	119.8	17.6

这种可预测的扩展规律使研究者能够通过增加计算预算来系统提升模型性能，为构建更强大的生成模型提供了清晰路径。

3.2 训练效率与稳定性

相比于U-Net架构，DiT展现了更优的训练稳定性和效率。这主要归功于：

Transformer的标准训练方案：DiT能够直接受益于NLP领域多年积累的Transformer优化技术
自适应归一化层：adaLN-Zero机制确保条件信息有效融入前向过程，同时保持训练稳定性
梯度一致性：Transformer的全局计算特性避免了卷积网络中的梯度异常问题

3.3 全局感受野与连贯生成

自注意力机制为DiT提供了全局感受野，使其能够捕捉图像中远距离区域间的依赖关系。这一特性在生成复杂结构化图像时尤为重要，比如：

物体对称性：确保生成物体的结构对称和比例协调
场景布局一致性：维持全局场景布局的合理性和视觉连贯性
纹理连贯性：保持纹理样式在整个图像中的一致性

相比之下，卷积网络的局部感受野需要堆叠多层才能建立远程依赖，且效率较低。

4 DiT的优化与改进方向

4.1 计算效率优化

尽管DiT性能卓越，但其计算成本仍是实际应用中的挑战。以下是几种有效的优化策略：

4.1.1 注意力机制优化

自注意力的计算复杂度随序列长度呈平方增长，这成为处理高分辨率图像的瓶颈。优化方法包括：

窗口注意力：将注意力计算限制在局部窗口内，大幅降低计算量
线性注意力：通过核函数近似将复杂度从$ O(N^2) $降为$ O(N) $
多查询注意力：多个查询头共享相同的键值对，减少内存带宽需求

4.1.2 令牌路由（TREAD）

TREAD方法通过令牌路由机制在训练时使用部分令牌集，跳过部分计算以减少成本，推理时仍使用完整令牌集。这种方法可实现14-37倍的训练速度提升，同时收敛到更好的生成质量。

4.2 架构改进

4.2.1 后层归一化争议

有批评指出DiT全面使用后层归一化（Post-LayerNorm） 可能不太稳定，尤其适用于数值范围变化剧烈的扩散过程。但DiT作者谢赛宁回应称，没有证据表明Post-Norm会带来负面影响，且Lightning DiT（结合SwiGLU、RMSNorm、RoPE、patch size=1）已成为经过验证的稳健升级版。