论文发布时间：2023.3.2

发布会议/期刊：ICCV

Abstract

• 利用transformer架构替代了U-Net，名为Diffusion Transformers (DiTs)。
• 通过增加transformers的深度和宽度或者输入tokends，有更高Gflops的DiTs却始终有较低的FID。
• DiT-XL/2在class-conditional ImageNet 512×512 and 256×256 benchmarks上取得SOTA。

Introduction

目标

将普遍使用于扩散模型的骨干网络U-Net替换为Transformer，并探讨其scaling等性能

动机

1. Transformers已经得到了广泛的应用，而图像生成模型却仍然抵制这一趋势，在许多diffusion模型中还在使用U-Net作为backbone。
2. U-Net的inductive bias对于diffusion的表现并不重要，可以被标准化的架构替代，并且保留scalability, robustness and efficiency，并且标准化架构更有利于跨领域研究

解决方案

1. 利用了四种不同的方案训练diffusion transformer，探讨不同方案的性能。

Diffusion Transformers

Preliminaries

Diffusion formulation

扩散模型（Diffusion Models），特别是去噪扩散概率模型（Denoising Diffusion Probabilistic Models, 简称DDPMs），是一类生成模型，通过逐步添加噪声到真实数据并学习逆过程来生成新数据。

1. 前向扩散过程

扩散模型假设存在一个前向扩散过程，该过程逐步将噪声应用于真实数据。具体来说，前向过程在每个时间步将数据添加少量高斯噪声，生成：

其中，是预先设定的超参数，是单位矩阵。

通过重参数化技巧，我们可以表示为：

其中，是从标准正态分布中采样的噪声。

2. 逆向扩散过程

扩散模型的核心在于学习一个逆向过程，即从噪声数据逐步去噪，恢复出真实数据。逆向过程定义为：

这里，和是由神经网络参数化的均值和协方差，用于预测下一步的均值和协方差。

3. 模型训练

扩散模型通过最大化数据的变分下界（variational lower bound）来训练，其损失函数主要由以下两部分组成：

其中，表示KL散度，用于衡量真实后验和模型后验之间的差异。由于和均为高斯分布，KL散度可以通过均值和协方差直接计算。

为了简化训练过程，通常采用以下策略：

1. 噪声预测网络：通过重参数化，将表示为一个噪声预测网络，模型的训练目标转化为最小化预测噪声与真实噪声之间的均方误差：

2. 协方差矩阵的学习：为了进一步提高模型性能，需要学习逆向过程的协方差矩阵。首先用训练噪声预测网络，然后固定，并优化整个损失函数来训练协方差矩阵。

4. 图像生成

训练完成后，生成新图像的过程如下：

1. 初始化：从标准正态分布中采样一个初始噪声。
2. 迭代去噪：通过逆向扩散过程，逐步采样直到生成最终的图像。

Classifier-free guidance

无分类器引导（Classifier-free Guidance）

条件扩散模型在生成过程中会接收额外的信息作为输入，例如类别标签。在这种情况下，反向过程表示为，其中噪声预测网络和协方差矩阵都是基于条件进行调整的。在此框架下，无分类器引导方法被用于引导采样过程，使得生成的样本满足较高。

根据贝叶斯定理，我们有：

因此，其梯度可以表示为：

在扩散模型中，输出可以被解释为得分函数（score function），因此可以通过以下方式引导DDPM采样过程，使得生成的样本具有较高的条件概率：

其中，表示引导的尺度（注意，当时，恢复为标准采样）。具体来说：

1. 条件与无条件预测：模型同时训练有条件（带标签）和无条件（不带标签，使用一个学习到的“空”嵌入）的噪声预测和。

2. 引导采样：在采样过程中，通过结合有条件预测与无条件预测的差异，并乘以引导尺度，调整噪声预测，从而增强生成样本对条件的符合度。

3. 随机丢弃条件信息：在训练阶段，通过随机丢弃条件并用“空”嵌入替代，使模型能够同时处理有条件和无条件的情况。

无分类器引导通过在模型内部同时处理有条件和无条件的信息，并在采样过程中动态调整生成过程，从而无需额外的分类器即可有效引导生成样本，更加简便且效果显著。

Latent diffusion models

潜在扩散模型（Latent Diffusion Models）

直接在高分辨率像素空间中训练扩散模型，往往会面临计算资源消耗巨大的问题。Latent Diffusion Model(LDMs)针对这一挑战提出了一种两阶段的方法：

1. 学习自动编码器（Autoencoder）：首先，训练一个自动编码器，将高分辨率的图像压缩为较小的空间表示。自动编码器包括一个编码器，用于将图像转换为潜在表示，以及一个解码器，用于将潜在表示还原为图像。在这一过程中，编码器和解码器的参数通过训练学习得到，确保可以有效地压缩和重构图像。

2. 训练潜在空间中的扩散模型：在第二阶段，不再直接在原始像素空间中的图像上训练扩散模型，而是在自动编码器生成的潜在表示上进行训练。具体来说，冻结编码器，仅训练扩散模型来生成潜在表示。

通过这种两阶段的方法，LDMs 能够显著减少计算需求，因为在潜在空间中的表示通常比原始高分辨率图像小得多，从而大幅降低了计算量。

如图所示，LDMs 在保持良好生成性能的同时，大幅降低了计算复杂度。这使得LDMs成为进行架构探索的理想起点，特别是在关注计算效率的情况下。在本文中，作者将DiTs（Diffusion Transformers） 应用于潜在空间。这种方法使得本文的图像生成pipeline成为一种混合方法：先使用现成的卷积变分自编码器（Convolutional VAEs）进行潜在空间的编码与解码，然后结合基于Transformer的去噪扩散概率模型。

Diffusion Transformer Design Space

Patchify

在DiT（Diffusion Transformer）模型中，Patchify 是其第一层关键组件，用于将输入的空间表示转换为适合Transformer处理的序列形式。

1. 输入空间表示

   DiT模型的输入是一个潜在空间表示。分辨率为256×256×3的图像经过编码器处理后得到的形状为32×32×4。图像将被压缩到一个较低维度的潜在空间中，每个空间位置有4个特征通道。

2. Patchify

   Patchify就是将空间表示转换为一系列的tokens，以便于Transformer处理。具体步骤如下：

   • 分块（Patching）：将输入的空间表示按照预定的块大小进行分块。例如，当块大小时，输入的32×32×4的会被分成个块。
   • 线性嵌入（Linear Embedding）：每个patch被线性投影到一个固定维度的向量空间，生成一个维度为的token。
     因此，整个输入在经过Patchify后会被转换为长度为的token序列，每个token的维度为，其中。

3. 位置编码（Positional Embeddings）

   在Patchify之后，对所有生成的token应用标准的ViT中的基于频率的位置编码（sine-cosine版本）。位置编码用于保留空间位置信息，使得模型能够理解token之间的相对和绝对位置关系。

4. 块大小与计算量的关系

   
   如图，patch大小是一个超参数，决定了Patchify过程中每个patch的尺寸及生成的token数量。token数量的增加会显著增加Transformer的计算量（以Gflops计)。因为Transformer的计算复杂度与token数量的平方成正比，减半将导致至少四倍的Gflops增加。然而，的变化并不会显著影响模型的参数数量，因为参数数量主要取决于Transformer的架构和每个令牌的维度。

DiT block design

Patchify之后，输入tokens由一系列transformer blocks处理。除了噪声图像之外，扩散模型有时还会处理附加的条件信息，比如加噪步长t，类标签c，自然语言等。本文探索了4中不同方式处理条件输入的transformer块。

• In-context conditioning：这个方法简单地将t和c地向量嵌入作为输入序列中地两个附加tokens，与图像tokens没有区别，类似于ViT中的[cls]。最后一个块之后，再将附加tokens从序列中拿出。这样的方法引入了可以忽略不计地新Gflops。
• Cross-attention block：这个方法将t和c的embeddings连接成为了一个长度为2的序列，与图像的token序列分开。transformer块会被修改，额外包含一个多头交叉注意力层在多头自注意力块后。多头交叉注意力增加了几乎15%的开销，贡献了最多的Gflops。
• Adaptive layer norm (adaLN) block：用adaptive layer norm(adaLN)替换标准的层归一化。作者不是直接学习缩放和平移参数和，而是通过和的embedding向量和对它们进行回归(通过线性层)。addLN增加了最少的Gflops，因此效率最高。
• adaLN-Zero block：除了回归and之外，作者还回归了一个新的参数，用于在残差连接之前进行缩放：。对于，作者首先初始化所有MLP的输出为零向量; 这使得所有残差连接先变为恒等映射。与普通的adaLN模块一样，adaLN-Zero向模型添加的Gflops可以忽略不计。

最后两个方法在附录有详细讲解

Model size

• N个DiT块，每个以​的隐藏维度运算，并配置注意力头，借鉴的ViT做法。
• 四个版本：DiT-S, DiT-B, DiT-L and DiT-XL。

Transformer decoder

在最终的DiT块之后，需要将图像token解码输出为噪声预测和对角协方差预测。这两个输出的形状都等于原始空间的输入。

使用一个标准线性解码器去实现：作者先使用final layer norm(adaptive if using adaLN)和线性解码器将最终token解码成的tensor，C为输入channels。最终，作者将这些解码过的tokens重排到原始的空间布局去得到预测的噪声和协方差矩阵。

Experimental Setup

命名：根据配置和latent patch size，例如DiT-XL/2表示XLarge配置和。

Training

• 在256 × 256 and 512 × 512 image resolution on the ImageNet dataset中训练。
• 用0或者ViT的标准权重初始化最终线性层。
• 全部用AdamW优化器训练。
• 学习率=，无权重衰减，batch size=256。
• 水平翻转数据增强。
• 没有进行学习率warmup或者正则化。
• 保留exponential moving average (EMA)，weights over training with a decay of 0.9999。
• 所有报告结果使用EMA模型，所有不同尺寸DiT模型用相同超参数，训练超参数几乎保留于ADM。

Diffusion

• 使用Stable Diffusion现成的预训练变分自编码器(VAE)。
• VAE的下采样因子为8，给一个256 × 256 × 3的RGB图像，的shape为32 × 32 × 4。
• 从扩散模型采样新的潜在样本后，再利用VAE将将其解码为像素。
• 保留ADM的超参数，specifically, we use a tmax = 1000 linear variance schedule ranging from 1×10−4 to 2 × 10−2, ADM’s parameterization of the covariance Σθ and their method for embedding input timesteps and labels.

Evaluation metrics

• 通过FID来衡量scaling性能，这是一个评估生成模型图像的标准矩阵。
• 本文报告 FID-50K using 250 DDPM sampling steps。
• 由于FID对小的实现细节很敏感，为保准确，所有报告通过导出样本使用ADM的TensorFlow评估。
• 本文报告的FID没有使用classifier-free guidance，额外报告了Inception Score，sFID，Precision/Recall作为附加评估矩阵。

Compute

• 在JAX上部署了所有模型，利用TPU-v3 pod训练。
• 最费算力的DiT-XL/2，以5.7iterations/second的速度在TPU v3-256 pod训练，全局batch size=256。

Experiments

DiT block design

Scaling model size and patch size

DiT Gflops are critical to improving performance

Larger DiT models are more compute-efficient

Visualizing scaling

State-of-the-Art Diffusion Models

256×256 ImageNet

512×512 ImageNet

Scaling Model vs. Sampling Compute

Conclusion

• 引入了 Diffusion Transformers (DiTs)，一个为diffusion模型设计的基于transformer的backbone，明显优于U-Net作为backbone的模型。
• 取得了SOTA结果，未来可以将DiT扩展到更大的模型和更多的token数量。DiT也可以直接作为 DALL·E 2h或者Stable Diffusion直接骨干。

附录

adaLN模块及其变体adaLN-Zero详解与推导

在Diffusion Transformer (DiT) 模型中，自适应层归一化（adaLN）模块及其变体adaLN-Zero在提升模型性能和计算效率方面发挥了关键作用。

一、标准层归一化（Layer Norm）

在讨论自适应层归一化之前，首先回顾标准的层归一化方法。

标准层归一化公式：

给定输入向量，层归一化（Layer Norm）的计算过程如下：

其中：
-和分别是输入向量的均值和标准差。
-和是可学习的缩放参数和偏移参数。
-是一个小常数，防止数值不稳定。

标准层归一化的特点：

• 固定参数：和是在训练过程中学习得到的固定参数，与输入的条件信息无关。
• 适用于所有样本：同一层的归一化参数在所有输入样本中共享。

二、自适应层归一化（Adaptive Layer Norm, adaLN）

引入自适应层归一化的动机：

在条件生成任务中（如扩散模型），希望生成的样本能够依据条件信息进行调整。标准层归一化中的固定和无法动态适应不同的条件信息，因此引入了自适应层归一化（adaLN）。

自适应层归一化的基本思想：

将归一化的缩放参数和偏移参数从固定参数转变为条件依赖的动态参数，即通过条件信息和的嵌入向量来预测和。

adaLN的具体公式：

1. 条件信息的嵌入：

   首先，将时间步和条件分别通过嵌入层和转换为向量：

2. 合并条件信息：

   将和相加，得到综合条件向量：

3. 回归缩放参数和偏移参数：

   利用一个回归网络（通常是一个小型的全连接网络或多层感知机，MLP来生成和：

4. 应用自适应层归一化：

   将动态生成的和应用于层归一化：

   其中，是标准层归一化的归一化结果。

自适应层归一化的特点：

• 动态调整参数：归一化参数和随条件信息和动态变化，使模型能根据不同条件调整输出。
• 计算效率高：相比其他条件机制（如条件批归一化），adaLN在计算量上更为高效，因为它只需少量参数来回归和。

三、adaLN-Zero 模块

引入adaLN-Zero的动机：

在深度神经网络（如ResNets）的训练中，初始化残差块为恒等函数可以提升训练稳定性和加速收敛。

adaLN-Zero的具体实现和推导：

1. 基础结构：

   在adaLN-Zero模块中，除了回归和外，还回归一个新的参数，用于在残差连接之前进行缩放：

2. 初始化策略：

   将的输出初始化为零向量，即使得在训练开始时：

   这样，在初始阶段：

   即无论残差连接中的内容如何，初始的残差连接对输出不会产生影响，整个adaLN-Zero模块表现为恒等映射。

3. 残差连接中的应用：

   假设一个Transformer块中的输入为，经过adaLN-Zero模块后的输出为：

   在初始阶段，由于，有：

   进一步，由于和也是通过回归得到的，并且在初始阶段如果被适当地初始化，也可以使得和不会对输出产生过大影响。

4. 动态调整：

   在训练过程中，会逐步从零向量变为需要的缩放参数，使得残差连接逐渐起作用，从而实现模型参数的逐步优化。

adaLN-Zero的优势：

• 训练稳定性：通过将残差块初始化为恒等映射，避免了初始阶段的梯度问题，提升了训练的稳定性。
• 快速收敛：有助于加速大规模模型的训练过程，提高训练效率。
• 计算效率高：与标准adaLN相似，adaLN-Zero增加的计算量（Gflops）极少，保持了高效性。