DIT: Scalable Diffusion Models with Transformers--Sora/SD3相关DIT技术论文阅读_dit中的transformer是否有自回归

OpenAI发布Sora，以及Stability.AI发布的SD3，根据其技术报告，使用了可扩展的transformer扩展模型，《Scalable Diffusion Models with Transformers》是其相关的一篇重要论文。
关于DIT作者进阶的论文SIT《SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers 》介绍，下一篇文章解析！！（知乎：二度；公众号：二度简并）

论文重点

文章提出，使用一种基于Transformer架构的新型扩散模型DiT，并训练了潜在扩散模型，用对潜在patch进行操作的Transformer替换常用的U-Net主干网络。通过以Gflops衡量的前向传递复杂度来分析扩散Transformer(DiT)的可扩展性。通过增加Transformer的深度/宽度或增加输入token的数量，始终具有较低的FID。

DiT架构与标准Vision Transformer (ViT)非常相似，但有一些小但重要的调整。扩散模型需要处理条件输入，例如扩散时间步长或类标签。最有效的是带有adaptive layer norm layers (adaLN)的ViT块。重要的是，这些adaLN层还会在块内任何残差连接之前激活，并进行初始化，以便每个ViT块都是恒等函数。简单地改变注入条件输入的机制就会在FID方面产生巨大的差异，这一更改是获得良好性能所需的唯一更改。除此之外，DiT是一个相当标准的Transformer模型。

论文摘要

1. 探索了一种基于transformer架构的新型扩散模型。训练了图像的潜在扩散模型，将常用的 U-Net 骨干替换为在潜在补丁上操作的transformer。
2. 通过基于前向传播复杂度的 Gflops 测量来分析扩散transformer（DiTs）的可扩展性。
3. 具有更高 Gflops 的 DiTs — 通过增加transformer的深度/宽度或增加输入token的数量会具有较低的 FID。
4. 除了具有良好的可扩展性属性外，最大的 DiT-XL/2 模型在条件 ImageNet 512×512 和 256×256 基准上均优于所有先前的扩散模型，后者实现了 2.27 的最新 FID。
5. 揭示扩散模型中架构选择的重要性，并为未来的生成建模研究提供经验基线。展示了 U-Net 的归纳偏差对扩散模型的性能并不至关重要，它们可以轻松地被标准设计，如Transformer所替代。
6. 研究了transformer相对于网络复杂度与样本质量的扩展行为。通过在 Latent Diffusion Models（LDMs）框架下构建和基准测试 DiT 设计空间来展示，扩散模型在 VAE 的潜在空间内进行训练，可以成功地用transformer替换 U-Net 骨干。
7. 进一步展示了 DiTs 是扩散模型的可扩展架构：网络复杂度（通过 Gflops 测量）与样本质量（通过 FID 测量）之间存在着强烈的相关性。通过简单地扩大 DiT 并训练具有高容量骨干（118.6 Gflops）的 LDM，能够在条件 ImageNet 256×256 生成基准上实现 2.27 FID 的最新结果。

简介

机器学习正在经历由 Transformer 驱动的复兴。在过去的五年里，自然语言处理、视觉和其他几个领域的神经架构很大程度上已被 Transformer 所涵盖。 然而，许多类别的图像级生成模型仍然抵制这一趋势——虽然 Transformer 在自回归模型中得到广泛使用，但在其他生成模型框架中的采用较少。 例如，扩散模型一直处于图像级生成模型最新进展的前沿； 然而，它们都采用卷积 U-Net 架构作为骨干网络的实际选择。U-Net 最初在像素级自回归模型和条件 GAN 中取得了成功，它继承自 PixelCNN++，并进行了一些更改。该模型是卷积模型，主要由ResNet块组成。与标准U-Net相比，额外的空间自注意力块（变压器中的重要组成部分）以较低的分辨率散布。Dhariwal和Nichol取消了 UNet 的几种架构选择，例如使用自适应归一化层为卷积层注入条件信息和通道计数。
通过这项工作，文章的目标是揭开扩散模型中架构选择的重要性，并为未来的生成建模研究提供经验基线。 表明，U-Net 归纳偏差对于扩散模型的性能并不重要，并且可以很容易地用Transformer等标准设计代替。 因此，扩散模型完全可以从最近的架构统一趋势中受益，例如，通过继承其他领域的最佳实践和训练方法，以及保留可扩展性、稳健性和效率等有利的属性。标准化架构也将为跨领域研究开辟新的可能性。
在文章中，关注一类基于Transformer的新型扩散模型。 称它们为Diffusion Transformer，简称 DiT。 DiT 遵循Vision Transformers (ViTs) 的最佳实践，事实证明，与传统的卷积网络（例如 ResNet）相比，它可以更有效地进行视觉识别扩展。更具体地说，文章研究了 Transformer 在网络复杂性与样本质量方面的缩放行为。 表明，通过在潜在扩散模型（LDM）框架下构建和基准化 DiT 设计空间，其中扩散模型在 VAE 的潜在空间中进行训练，可以成功地用Transformer替换 U-Net 主干。 进一步证明 DiT 是扩散模型的可扩展架构：网络复杂性（以 Gflops 衡量）与样本质量（以 FID 衡量）之间存在很强的相关性。通过简单地扩展 DiT 并训练具有高容量主干（118.6 Gflops）的 LDM，能够在类条件 256 × 256 ImageNet 生成基准上实现 2.27 FID 的最先进结果。
文章发现U-Net并非不可替代，并且很容易使用诸如Transformers的结构替代U-Net，使用Transformers可以很好地保持原有的优秀特性，比如可伸缩性、鲁棒性、高效性等，并且使用新的标准化架构可能在跨领域研究上展现出更多的可能。文章从网络复杂度和采样质量两个方面对DiTs方法进行评估。

相关工作

Transformer

Transformers 已经取代了跨语言、视觉、强化学习和元学习的领域特定架构。作为通用自回归模型和ViT，它们在增加模型大小、训练计算和语言领域数据的情况下表现出了显着的扩展特性。除了语言之外，Transformer 还经过训练可以自回归预测像素。 他们还接受了离散codebook本作为自回归模型和掩蔽生成模型的训练； 前者在高达 20B 参数时表现出出色的缩放行为。 最后，在 DDPM 中探索了 Transformer 来合成非空间数据； 例如，在 DALL·E 2 中生成 CLIP 图像嵌入。 文章主要研究了Transformer用作图像扩散模型的骨干时的缩放特性。

DDPMs

扩散和基于分数的生成模型作为图像生成模型特别成功，在许多情况下优于生成对抗网络（GAN）以前是最先进的。过去两年 DDPM 的改进很大程度上是由改进的采样技术推动的，最显着的是无分类器指导、重新制定扩散模型来预测噪声而不是像素以及使用级联 DDPM 管道，其中 低分辨率基础扩散模型与上采样器并行训练 。对于上面列出的所有扩散模型，卷积 U-Nets 是主干架构的事实上的选择。 并行工作引入了一种基于 DDPM 注意力的新颖、高效的架构。
扩散模型是借鉴了物理学上的扩散过程，在生成模型上，分为正向和逆向的过程。正向过程是向信号中逐渐每步加少量噪声，当步数足够大时可以认为信号符合一个高斯分布。所以逆向过程就是从随机噪声出发逐渐的去噪，最终还原成原有的信号。去噪过程一般采用UNet或者ViT，使用t步的结果和条件输入预测t-1步增加的噪声，然后使用DDPM可以得到t-1步的分布，经过多步迭代就可以从随机噪声还原到有实际意义的信号。如果使用原始DDPM速度会慢很多，所以很多工作如DDIM、FastDPM等工作实现了解码加速。在图像的无条件生成任务上，扩散模型的性能已经超过了GANs，并且在有条件生成如文图生成任务上大放异彩。

架构的复杂度

在评估图像生成文献中的架构复杂性时，使用参数计数是相当常见的做法。 一般来说，参数计数可能不能很好地代表图像模型的复杂性，因为它们没有考虑到显着影响性能的图像分辨率等因素。相反，作者的大部分模型复杂性分析都是通过理论Gflops的视角进行的。 这使与架构设计文献保持一致，其中Gflop被广泛用于衡量复杂性。 在实践中，黄金复杂度指标仍然存在争议，因为它经常取决于特定的应用场景。Nichol和Dhariwal改进扩散模型的开创性工作与我们最相关——他们分析了U-Net架构类的可扩展性和Gflop属性。文章重点关注Transformer类。

方法

基础

Diffusion基础

高斯扩散模型假设前向噪声过程逐渐将噪声应用于实际数据，前向过程是一个T步逐渐加噪的马尔科夫链。所以，给定前向扩散过程作为先验，扩散模型训练反转的过程，可以通过去除所加噪声从XT恢复成X0，并且每步的扩散过程都采样自特定的高斯分布。

Classifier-free guidance

条件扩散模型采用额外信息作为输入，例如类标签 c。 在这种情况下，逆过程变为pθ(xt-1|xt, c)，其中θ和Σθ以c为条件。 在这种情况下，可以使用无分类器指导来鼓励采样过程找到 x，使得 log p(c|x) 较高。所以在想要条件的概率较大，就可以将条件的梯度增加到优化目标里。模型在训练时，使用一个网络架构优化两个模型（uncond，cond）。众所周知，与通用采样技术相比，无分类器指导可以产生显着改进的样本，并且这种趋势也适用于DiT模型。

潜在扩散模型

直接在高分辨率像素空间中训练扩散模型在计算上可能会令人望而却步。 潜在扩散模型（LDM）通过两阶段方法解决这个问题：（1）学习一个自动编码器，使用学习的编码器 E 将图像压缩成更小的空间表示； (2) 训练表示 z = E(x) 的扩散模型，而不是图像 x 的扩散模型（E 被冻结）。 然后可以通过从扩散模型中采样表示 z 并随后使用学习的解码器 x = D(z) 将其解码为图像来生成新图像。
简单来说就是，模型使用VAE（固定权重）将图片encoder到隐空间，生成结果同样也是通过VAE解码成原始大小的图片。

DIT架构与设计思路

Patchify

DiT 的输入是通过VAE后的一个稀疏的空间表示z（对于256×256×3图像，z的形状为32×32×4）。DiT 的第一层是“patchify”，它通过将每个 patch 线性嵌入到输入中，将空间输入转换为一系列 T 标记，每个标记的维度为 d。 patchify 之后，将标准 ViT 基于频率的位置嵌入（正弦余弦版本）应用于所有输入token。patchify创建的token T 的数量由patch大小超参数 p 决定。如图 4 所示，将 p 减半将使 T 增加四倍，从而使Transformer总 Gflop 至少增加四倍。 尽管它对 Gflops 有重大影响，但更改 p 对下游参数计数没有有意义的影响。 文章采用将 p = 2, 4, 8 添加到 DiT 设计空间，进行相关对比试验。

DIT模块设计

在Patch化之后，输入token由一系列Transformer块处理。除了噪声图像输入之外，扩散模型有时还会处理其他条件信息，例如噪声时间步 t、类标签 c、自然语言等。文章探索了以不同方式处理条件输入的 Transformer 块的四种变体。 这些设计对标准 ViT 模块设计进行了微小但重要的修改。

In-context 条件

将 t 和 c 的向量emb作为输入序列中的两个附加token（与输入的图像token做拼接），将它们与图像token没有区别。 这类似于 ViT 中的 cls token，它允许无需修改即可使用标准 ViT 块。 在最后一个块之后，从序列中删除条件token。 这种方法向模型引入的新 Gflop 可以忽略不计。

Cross-attention模块

将 t 和 c 的emb连接成一个长度为 2 的序列，与图像token序列分开。Transformer块被修改为在多头自注意力块之后包含一个额外的多头交叉注意力层，类似于 LDM 用于调节类标签的方法（DiT结构与Condition交互的方式，与原来U-Net结构类似）。 交叉注意力为模型增加了最多的 Gflops，大约增加了 15% 的开销。

Adaptive layer norm（adaLM）模块

探索用自适应layer norm (adaLN) 替换Transformer块中的标准layer norm。不是直接学习维度尺度和平移参数 γ 和 β，而是根据 t 和 c 的嵌入向量之和对它们进行回归。 这种方式，adaLN 添加的 Gflops 最少，因此计算效率最高。 它也是唯一限制对所有token应用相同功能的调节机制。（使用adaLN替换原生LayerNorm（NeurIPS2019的文章，LN 模块中的某些参数不起作用，甚至会增加过拟合的风险。所以提出一种没有可学习参数的归一化技术））

adaLN-zero 模块

之前关于 ResNets 的工作发现，将每个残差块初始化为恒等函数是有益的。扩散 U-Net 模型使用类似的初始化策略，在任何残差连接之前对每个块中的最终卷积层进行零初始化。作者探索了 adaLN DiT 块的修改，它具有相同的功能。 除了回归 γ 和 β 之外，还回归维度缩放参数 α，这些参数在 DiT 块内的任何剩余连接之前立即应用。（对残差块的γ, β, α进行衰减）

模型大小

模型应用一系列 N DiT 块，每个块都以隐藏维度大小 d 进行操作。 遵循 ViT，使用标准Transformer配置来联合缩放 N、d 和注意力头。使用四种配置：DiT-S、DiT-B、DiT-L 和 DiT-XL。 它们涵盖了广泛的模型大小和触发器分配，从 0.3 到 118.6 Gflops，使得能够衡量扩展性能。这些模型配置范围从 33M 到 675M 参数

Transformer Decoder

在最终的 DiT 块之后，需要将图像token序列解码为输出噪声预测和输出对角协方差预测。这两个输出的形状都等于原始空间输入。使用标准线性解码器来做到这一点； 应用最终层Layer Norm（如果使用 adaLN 则为自适应）并将每个token线性解码为 p×p×2C 张量，其中 C 是 DiT 空间输入中的通道数。 最后，将解码后的token重新排列为其原始空间布局，以获得预测的噪声和协方差。文章主要探索的完整 DiT 设计空间是patch大小、Transformer块架构和模型大小。

实验设置

文章重点探索 DiT 设计空间并研究模型类的缩放属性。模型根据其配置和潜在patch大小p命名；例如，DiT-XL/2 指的是 XLarge 配置且 p = 2。

训练

1. ImageNet 数据集，以256 × 256 和 512 × 512 图像分辨率训练
2. 零初始化最终的线性层，否则使用 ViT 的标准权重初始化技术。
3. AdamW 优化器，1×10−4 的恒定学习率，没有权重衰减，批量大小为 256
4. 数据增强是水平翻转
5. 与之前的 ViT 工作不同，没有发现学习率预热或正则化对于训练 DiT 达到高性能是必要的。
6. 训练过程中保持 DiT 权重的指数移动平均值 (EMA)，衰减为 0.9999。所有结果均使用 EMA 模型。
7. 在所有 DiT 模型大小和 patch 大小上使用相同的训练超参数。训练超参数几乎完全保留自 ADM。
8. 没有调整学习率、衰减/热身计划、Adam β1/β2 或权重衰减

Diffusion

1. 使用来自Stable Diffusion的现成预训练变分自动编码器（VAE）模型
2. VAE编码器的下采样因子为 8，给定RGB图像x的形状为256×256×3，z=E(x)的形状为32×32×4。
3. 从扩散模型中采样新的潜在变量后，使用 VAE 解码器x=D(z)将其解码为像素。保留来自ADM的扩散超参数

评价指标

1. 使用 Frechet Inception Distance (FID)来衡量缩放性能，这是评估图像生成模型的标准指标。
2. 使用250个DDPM采样步骤报告FID-50K。FID 对小的实施细节很敏感，为了确保准确的比较，文章所有值都是通过导出样本并使用 ADM 的 TensorFlow 评估套件获得的。
3. 文章的FID编号不使用无分类器指导。
4. 另外对比Inception Score、sFID和Precision/Recall 作为次要指标

实验结果

1. DIT结构：adaLN-Zero效果最好

训练了四个最高 Gflop DiT-XL/2 模型，每个模型都使用不同的块设计：

1. in-context（119.4 Gflops）
2. cross-attention（137.6 Gflops）
3. adptive-LN（adaLN，118.6 Gflops）
4. adaLN-zero （118.6 Gflops）
   图5为FID结果。adaLN-Zero块产生的FID低于corss-attention和in-content，同时计算效率最高。adaLNZero将每个 DiT 块初始化为恒等函数，显着优于普通的 adaLN。 对于文章的其余部分，所有模型都将使用 adaLN-Zero DiT 模块。
   

2. 缩放模型大小和patch大小（模型越大、patch size越小生成图像质量越好）

训练了12个 DiT 模型，涵盖模型配置（S、B、L、XL）和patch大小（8、4、2）。与其他配置相比，DiT-L 和 DiT-XL 在相对 Gflop 方面明显更接近。缩放模型尺寸，通过增加 Transformer 深度 / 宽度或增加输入 token 数量，具有较高 Gflops 的 DiT 始终具有较低的 FID

1. 图 2（左）概述了每个模型的 Gflops 及其在 400K 训练迭代时的 FID。 在所有情况下，增加模型大小和patch大小可以显着改进扩散模型。
2. 图 6（上）展示了 FID 如何随着模型尺寸的增加而变化并且patch尺寸保持不变。在所有四种配置中，通过使Transformer更深更宽，FID 在训练的所有阶段都获得了显着改进。
3. 图 6（下）显示了随着patch尺寸减小而模型尺寸保持不变的 FID。通过简单地扩展 DiT 处理的token数量并保持参数大致固定，整个训练过程中 FID 有了相当大的改进。
   
   

3. DiT Gflops 对于提高性能至关重要（计算开销与模型效果关系：Gflops越大模型效果越好）

图 6 的结果表明参数计数并不能唯一地确定 DiT 模型的质量。当模型大小保持不变并且patch大小减小时，Transformer的总参数实际上没有改变（实际上，总参数略有减少），并且仅 Gflops 增加。这些结果表明，缩放模型 Gflops 实际上是提高性能的关键。 为了进一步研究这一点，在图 8 中针对模型 Gflops 绘制了 400K 训练步骤下的 FID-50K。结果表明，当总 Gflops 相似时，不同的 DiT 配置会获得相似的 FID 值（例如，DiT-S/2 和 DiT- B/4）。发现模型 Gflops 和 FID-50K 之间存在很强的负相关性，这表明额外的模型计算是改进 DiT 模型的关键因素。 在图 12（附录）中，发现这种趋势也适用于其他指标，例如 Inception Score。

4. 较大的 DiT 模型的计算效率更高（计算开销与模型效果关系：模型越大计算约高效）

在图 9 中，将 FID 绘制为所有 DiT 模型的总训练计算的函数。将训练计算估计为模型 Gflops · 批量大小 · 训练步骤 · 3，其中因子 3 大致近似为后向传递的计算量是前向传递的两倍。发现相对于训练步骤较少的大型 DiT 模型而言，小型 DiT 模型即使训练时间较长，最终也会变得计算效率低下。同样，发现除了patch大小之外相同的模型即使在控制训练 Gflop 时也具有不同的性能配置文件。 例如，在大约 1010 Gflops 后，XL/4 的性能优于 XL/2。

5. 可视化
   在图 7 中可视化缩放对样本质量的影响。在 400K 训练步骤中，使用相同的起始噪声 xtmax 、采样噪声和类标签从 12 个 DiT 模型中的每一个中采样图像。可以直观地解释缩放如何影响 DiT 样本质量。 事实上，缩放模型大小和token数量可以显着提高视觉质量。
   

6. SOTA Diffusion Models

256×256 ImageNet

在缩放分析之后，在表 2 中展示最终结果。当使用无分类器引导时，DiT-XL/2 优于所有先前的扩散模型，将 LDM 实现的先前最佳 FID-50K 3.60 降低到 2.27。图 2（右）显示 DiT-XL/2 (118.6 Gflops) 相对于 LDM-4 (103.6 Gflops) 等潜在空间 U-Net 模型具有计算效率，并且比 ADM 等像素空间 U-Net 模型更高效 (1120 Gflops) 或 ADM-U (742 Gflops)。文章方法实现了所有现有生成模型中最低的 FID，包括之前最先进的 StyleGANXL 。 最后，还观察到，与 LDM-4 和 LDM-8 相比，DiT-XL/2 在所有测试的无分类器指导量表中实现了更高的召回值。当仅训练 2.35M 步骤（类似于 ADM）时，XL/2 的 FID 为 2.55，仍然优于所有先前的扩散模型。

512x512 ImageNet

在 ImageNet 上以 512 × 512 分辨率训练新的 DiT-XL/2 模型，进行 3M 次迭代，其超参数与 256 × 256 模型相同。 patch大小为 2 时，该 XL/2 模型在patch 64 × 64 × 4 输入潜在变量 (524.6 Gflops) 后总共处理 1024 个token。 表 3 显示了与最先进方法的比较。 在此分辨率下，XL/2 再次优于所有先前的扩散模型，将 ADM 实现的先前最佳 FID 3.85 提高到 3.04。 即使token数量增加，XL/2 仍然保持计算效率。 例如，ADM使用1983 Gflops，ADM-U使用2813 Gflops； XL/2 使用 524.6 Gflops。 在图 1 和附录中展示了来自高分辨率 XL/2 模型的样本。

7. 缩放模型与采样计算（扩大采样计算不能弥补模型计算的不足）

扩散模型的独特之处在于，它们可以通过在生成图像时增加采样步骤数来在训练后使用额外的计算。 考虑到模型 Gflops 对样本质量的影响，研究较小模型计算 DiT 是否可以通过使用更多采样计算来优于较大模型？？？ 在 400K 训练步骤后计算所有 12 个 DiT 模型的 FID，每个图像使用 [16, 32, 64, 128, 256, 1000] 个采样步骤。主要结果如图 10 所示。考虑使用 1000 个采样步骤的 DiT-L/2 与使用 128 个采样步骤的 DiT-XL/2。 在这种情况下，L/2 使用 80.7 Tflops 来采样每个图像； XL/2 使用减少 5 倍的计算量（15.2 Tflops）来对每个进行采样图像。 尽管如此，XL/2 的 FID-10K 更好（23.7 vs 25.9）。 一般来说，扩大采样计算不能弥补模型计算的不足。

结论

文章主要引入了 Diffusion Transformers (DiTs)，一种基于 Transformer 的简单扩散模型主干，其性能优于之前的 U-Net 模型，并继承了 Transformer 模型类出色的缩放特性。在Gflops与Stable Diffusion相当的DiTs-XL/2的结构上，把ImageNet 256×256数据集上的FID指标优化到了2.27，达到了SOTA的水平。未来的工作应该继续将 DiT 扩展到更大的模型和token数量。 DiT 还可以作为 DALL·E 2 和稳定扩散等文本到图像模型的直接骨干进行探索。

1. DiT的本质思路就是用 Transformer 替换掉 UNet，并验证了不同scaling 能力和条件注入的结构。
2. 在 Diffusion 成功之后，以 DiT 为代表的，有很多用 Transformer 替换 UNet 的工作，但都没有受到很大的关注。究其原因，应该是UNet 已经做的足够好，计算开销也低，使用 Transformer scaling 上去，虽然生成质量有所提升，但也没那么显著。
3. 在最近 OpenAI 的 Soray以及Stability.AI 的SD3 炸裂登场之后，借助 Transformer 实现建模的方式再次步入人们视线。

关于DIT作者进阶的论文SIT《SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers 》介绍，下一篇文章解析！！（知乎：二度；公众号：二度简并）

参考

论文地址：https://arxiv.org/pdf/2212.09748.pdf
GitHub地址：https://github.com/facebookresearch/DiT
项目地址：https://www.wpeebles.com/DiT.html