论文笔记DiT：Scalable Diffusion Models with Transformers（含transformer的可扩展扩散模型 ）

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。
• 结果：经过 "patchify" 操作，一个二维的 z 就变成了一个一维的 token 序列，序列长度为 T = (H/p) * (W/p)。
• p 直接决定了序列的长度 T，在几乎不增加模型参数量的情况下，可以通过减小 p 来急剧增加模型的计算密度和复杂度。

核心模块：

除了噪声图像输入，扩散模型有时还会处理额外的条件信息，如噪声时间步长t、类标号c、自然语言等。我们探索了4种不同处理条件输入的变压器模块变体。该设计对标准的ViT块设计进行了微小但重要的修改。所有模块的设计如图3所示。

1.In-context Conditioning (上下文条件)

    方法：把 t 和 c 的嵌入向量当作两个额外的 token，直接拼接到图像 token 序列的前面。

    优缺点：非常简单，可以直接使用标准的 ViT 模块。但实验表明效果最差。

2. Cross-attention Block (交叉注意力模块)

    方法：在标准的自注意力模块之后，增加一个交叉注意力层。图像 token 作为 Query，t 和 c 的嵌入作为 Key 和 Value。

    优缺点：这是 LDM（Stable Diffusion v1）中使用的方法，比较强大和灵活。但它增加了最多的计算量（约 15% 的开销），且在 DiT 的实验中并非最佳选择。

3. Adaptive Layer Norm (adaLN) Block (自适应层归一化模块)

    方法：受到 GAN 和 U-Net 扩散模型的启发，作者将 Transformer 模块中的标准 LayerNorm 替换为自适应的 LayerNorm。模型的 LayerNorm 不再学习固定的缩放（γ）和偏移（β）参数，而是通过一个小型网络，根据 t 和 c 的嵌入动态地预测出 γ 和 β。

    优缺点：计算上非常高效，增加的 Gflops 几乎可以忽略不计。它将条件信息全局地作用于所有 token。

4. adaLN-Zero Block (零初始化的 adaLN)

    方法：这是对 adaLN 的一个重要改进。除了像 adaLN 一样预测 γ 和 β 之外，它还为每个残差连接（residual connection）之前添加了一个由 t 和 c 预测的缩放因子 α。

    关键的初始化：在训练开始时，这个 α 被初始化为输出零。这意味着整个 DiT  block 在初始状态下就是一个恒等函数（identity function），输入什么就输出什么。

    灵感来源：这种“零初始化残差分支”的策略在 ResNet 和 U-Net 扩散模型中都被证明有助于稳定大型模型的训练。

    最终选择：实验证明 adaLN-Zero 是效果最好且计算高效的方案。

Model Size：

为了系统地研究可扩展性，作者定义了一系列不同尺寸的模型。

他们沿用了 ViT 的标准配置，通过同时扩展模型的深度 N (层数)、宽度 d (隐藏层维度) 和注意力头的数量，创建了四种尺寸：DiT-S (Small), B (Base), L (Large), XL (XLarge)。

例如 DiT-XL/2 指的是 XL 尺寸的模型，使用 p=2 的 patch size。

Transformer Decoder：

    经过一系列 DiT block 处理后，得到一个更新后的 token 序列。现在需要将这个序列解码回一个二维的预测噪声图。

    对最终的 token 序列应用最后一次 LayerNorm。

    使用一个简单的线性层，将每个 token 解码回它对应的 p x p 大小的 patch 形状。

    将所有解码后的 patch 重新组合，恢复成与输入 z 相同尺寸的二维特征图，作为最终的噪声预测。