Diffusion Transformers (DiT) 原理总结

创作时间:

作者:

@小白创作中心

Diffusion Transformers (DiT) 原理总结

引用

CSDN

https://m.blog.csdn.net/qq_43069203/article/details/145940359

Diffusion Transformers (DiT) 是一种基于Transformer架构的扩散模型，旨在替代传统的U-Net结构，解决其计算冗余、参数效率低和长程依赖建模不足的问题。DiT通过将VAE输出的潜在特征分块并展平成序列，结合ViT的标准架构来处理这些序列，从而提升了生成质量和计算效率。此外，DiT引入了自适应条件注入机制，通过改进的adaLN-Zero块有效融合时间步和类别信息，使得模型能更好地处理不同条件下的生成任务。实验结果表明，DiT在图像生成任务中展现了强大的性能，尤其在ImageNet数据集上刷新了SOTA指标。

DiT（扩散变换器）是基于Transformer架构的扩散模型，它由UC Berkeley和纽约大学的谢赛宁团队提出，作为对传统基于U-Net的扩散模型的一种创新性替代方案。DiT在图像生成领域展现了强大的性能，尤其是在ImageNet任务上刷新了SOTA（state-of-the-art，最先进）指标，标志着生成模型架构的统一化和发展。其核心创新点如下：

1. U-Net局限性

计算冗余：U-Net的跳跃连接和跨分辨率操作较为低效。
参数效率低：U-Net需要在每个分辨率阶段使用独立的卷积核，导致参数量随深度增加。
长程依赖建模不足：注意力层仅在高分辨率层（如SD v1.5的8x8层）使用，未能有效捕捉全局信息。

2. DiT架构创新

潜在空间分块（Patchify）：将VAE编码的潜在空间特征 Latent 表示分割成小块 patches，并展平成序列，类似ViT（Vision Transformer）的做法。这种方式可以控制计算复杂度，通过调整patch大小来优化性能。

输入：VAE编码的潜在空间特征（尺寸H×W×C，如32×32×4）。

分块：按p×p大小切分，生成N =[ (H×W)/p] ²个块。

投影：线性映射至高维空间（维度d）, 得到DiT输入的Latent特征维度
是 N x d。

Transformer主干：DiT采用标准的ViT架构处理这些序列，完全替代了U-Net中的卷积操作。多头注意力 + MLP

3. 条件注入机制

自适应条件注入：通过改进的adaLN-Zero块，将时间步（timestep）和类别标签等条件信息与输入特征融合。该过程通过调整归一化参数，使得模型能够更好地处理不同条件下的生成任务。

4. DiT设计中的关键块

上下文条件（Context Cond）：将时间步
t
和类别标签
c
作为额外的tokens，并与输入特征拼接，提升了模型的表达能力，同时计算开销较低。
交叉注意力（Cross-Attn）：将额外的条件信息（如时间步和类别标签）视为一个单独的序列，加入额外的多头交叉注意力机制，增加了一些计算负担，但提升了生成质量。
自适应层归一化（adaLN）：用自适应归一化替代标准归一化层，从
t
和
c
的嵌入向量中回归出维度缩放参数γ \gammaγ和移位参数β \betaβ
最小化计算复杂度并优化训练。
adaLN-Zero：进一步优化了自适应归一化机制，进一步在任何残差连接前增加一个维度缩放参数α \alphaα, 并将这些参数初始化为零，优化了块的初始化。提高了模型稳定性。

5. 实验设计与效果

DiT在多个模型规模和patch大小的配置下进行实验，结果表明，增大模型规模和减小patch大小能显著提升扩散模型的生成性能。
在ImageNet数据集上的实验中，DiT-XL/2模型在256x256和512x512分辨率下，分别取得了2.27和3.04的FID（Fréchet Inception Distance）得分，超越了所有现有的扩散模型。