(ICCV-2023)使用 Transformer 的可扩展扩散模型

使用 Transformer 的可扩展扩散模型

Paper Title:Scalable Diffusion Models with Transformers

Paper是UC Berkeley发表在ICCV 2023的工作

paper地址

code地址

图 1：具有 transformer 主干的扩散模型实现了最先进的图像质量。我们展示了在 ImageNet 上以 512$\times $512 和 256$ \times$256 分辨率训练的两个类条件 DiT-XL/2 模型中的选定样本。

Abstract

我们探索了一种基于transformer架构的新类扩散模型。我们训练图像的潜在扩散模型，将常用的 $U$ -Net主干替换为在潜在patch上操作的transformer。我们通过Gflops衡量的前向传递复杂性分析了我们的扩散transformer（DiTs）的可扩展性。我们发现，通过增加transformer的深度/宽度或增加输入token的数量来提高Gflops的DiTs，其FID（Fréchet Inception Distance）值始终较低。除了具有良好的可扩展性外，我们最大的DiT-XL/2模型在类条件ImageNet $512 \times 512$ 和 $256 \times 256$ 基准上优于所有之前的扩散模型，在后者上实现了2.27的最先进FID。

1. Introduction

机器学习正在经历由 transformer 推动的复兴。在过去五年中，自然语言处理 [42, 8]、视觉 [10] 和其他几个领域的神经架构已被 transformer [60] 所取代。然而，许多类别的图像级生成模型仍然没有赶上这一趋势——虽然 transformer 在自回归模型中得到了广泛的应用 [43, 3, 6, 47]，但在其他生成建模框架中却很少采用。例如，扩散模型一直处于图像生成领域最新进展的前沿 [9, 46]；然而，它们都采用卷积 U-Net 架构作为事实上的主干选择。

Ho 等人 [19] 的开创性工作首次为扩散模型引入了 U-Net 主干。U-Net 最初在像素级自回归模型和条件 GAN [23] 中取得成功，之后从 PixelCNN++ [52, 58] 继承而来，并进行了一些更改。该模型是卷积的，主要由 ResNet [15] 块组成。与标准 U-Net [49] 相比，额外的空间自注意力块（Transformer 中必不可少的组件）以较低的分辨率散布。Dhariwal 和 Nichol [9] 放弃了 UNet 的几种架构选择，例如使用自适应归一化层 [40] 为卷积层注入条件信息和通道计数。但是，Ho 等人的 UNet 的高级设计基本保持不变。

通过这项工作，我们旨在揭开扩散模型中架构选择的重要性，并为未来的生成模型研究提供经验基线。我们表明，U-Net 归纳偏差对扩散模型的性能并不重要，它们可以很容易地用标准设计（如transformer）取代。因此，扩散模型可以从最近的架构统一趋势中受益 - 例如，通过继承其他领域的最佳实践和训练方法，以及保留可扩展性、稳健性和效率等有利特性。标准化架构还将为跨领域研究开辟新的可能性。

在本文中，我们重点介绍一类基于 Transformer 的新型扩散模型。我们将其称为扩散 Transformer，简称 DiT。DiT 遵循 Vision Transformer (ViT) [10] 的最佳实践，事实证明，ViT 在视觉识别方面的扩展效果比传统卷积网络（例如 ResNet [15]）更有效。

更具体地说，我们研究了 Transformer 相对于网络复杂度与样本质量的扩展行为。我们表明，通过在潜在扩散模型 (LDM) [48] 框架下构建和基准测试 DiT 设计空间（其中扩散模型在 VAE 的潜在空间内训练），我们可以成功地用 Transformer 替换 U-Net 主干。我们进一步表明，DiT 是扩散模型的可扩展架构：网络复杂度（以 Gflops 衡量）与样本质量（以 FID 衡量）之间存在很强的相关性。通过简单地扩大 DiT 的规模并使用高容量主干（118.6 Gflops）训练 LDM，我们就能够在类条件 256$\times$256 ImageNet 生成基准上实现 2.27 FID 的最佳结果。

2. Related Work

Transformer。Transformer [60] 已取代语言、视觉 [10]、强化学习 [5, 25] 和元学习 [39] 等领域的特定领域架构。它们在增加模型大小、训练计算和语言领域数据 [26] 的情况下表现出了出色的扩展特性，可用作通用自回归模型 [17] 和 ViT [63]。除了语言之外，Transformer 还被训练用于自回归预测像素 [38, 7, 6]。它们还在离散码本 [59] 上被训练为自回归模型 [11, 47] 和掩码生成模型 [4, 14]；前者在多达 20B 个参数的情况下表现出出色的扩展行为 [62]。最后，人们在 DDPM 中探索了使用 Transformer 来合成非空间数据；例如，在 DALL·E2[46, 41] 中生成 CLIP 图像嵌入。在本文中，我们研究了变压器作为图像扩散模型的骨干时的缩放特性。

去噪扩散概率模型 (DDPM)。
扩散 [54, 19] 和基于分数的生成模型 [22, 56] 作为图像生成模型 [35, 46, 50, 48] 尤其成功，在许多情况下甚至优于之前最先进的生成对抗网络 (GAN) [12]。过去两年来，DDPM 的改进主要得益于改进的采样技术 [19, 55, 27]，最显著的是无分类器指导 [21]、重新制定扩散模型以预测噪声而不是像素 [19] 以及使用级联管道，其中低分辨率基础扩散模型与上采样器并行训练 [20, 9]。对于上面列出的所有扩散模型，卷积 U-Nets [49] 是事实上的主干架构选择。并行工作 [24] 为 DDPM 引入了一种基于注意力机制的新颖、高效的架构；我们探索了纯 Transformer。

架构复杂性。在评估图像生成文献中的架构复杂性时，通常使用参数计数。一般来说，参数计数不能很好地代表图像模型的复杂性，因为它们没有考虑图像分辨率等因素，而图像分辨率会显著影响性能 [44, 45]。相反，本文的大部分分析都是通过计算的视角进行的。
这使我们与架构设计文献保持一致，其中广泛使用flops来衡量复杂性。在实践中，黄金指标将取决于特定的应用场景。Nichol 和 Dhariwal 改进扩散模型的开创性工作 [36, 9] 与我们最相关——在那里，他们分析了 U-Net 架构类的可扩展性。在本文中，我们重点关注 Transformer 类。

3. Diffusion Transformers

3.1. Preliminaries

扩散公式。在介绍我们的架构之前，我们简要回顾一些理解扩散模型（DDPMs）[54, 19]所需的基本概念。高斯扩散模型假设一个前向加噪过程，逐渐将噪声应用于真实数据 $x_0: q\left(x_t \mid x_0\right)=\mathcal{N}\left(x_t ; \sqrt{\bar{\alpha}_t} x_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right)$ ，其中常数 $\bar{\alpha}_t$ 是超参数。通过应用重新参数化技巧，我们可以采样 $x_t=\sqrt{\bar{\alpha}_t} x_0+\sqrt{1-\bar{\alpha}_t} \epsilon_t$ ，其中 $\epsilon_t \sim \mathcal{N}(0, \mathbf{I})$ 。

扩散模型被训练以学习反向过程，反转前向过程中的扰动： $p_\theta\left(x_{t-1} \mid x_t\right)=\mathcal{N}\left(\mu_\theta\left(x_t\right), \Sigma_\theta\left(x_t\right)\right)$ ，其中神经网络用于预测 $p_\theta$ 的统计信息。反向过程模型通过 $x_0$ 的对数似然的变分下界[30]进行训练，简化为 $\mathcal{L}(\theta)=-p\left(x_0 \mid x_1\right)+\sum_t \mathcal{D}_{K L}\left(q^*\left(x_{t-1} \mid x_t, x_0\right) \| p_\theta\left(x_{t-1} \mid x_t\right)\right)$ ，不包括训练无关的附加项。通过将 $\mu_\theta$ 重参数化为噪声预测网络 $\epsilon_\theta$ ，可以通过预测噪声 $\epsilon_\theta\left(x_t\right)$ 与采样高斯噪声 $\epsilon_t$ 之间的简单均方误差来训练模型： $\mathcal{L}_{\text {simple }}(\theta)=\left\|\epsilon_\theta\left(x_t\right)-\epsilon_t\right\|_2^2$ 。但为了训练具有学习的反向过程协方差 $\Sigma_\theta$ 的扩散模型，需要优化完整的 $\mathcal{D}_{K L}$ 项。我们遵循Nichol和Dhariwal的方法[36]：用 $\mathcal{L}_{\text {simple }}$ 训练 $\epsilon_\theta$ ，并用完整的 $\mathcal{L}$ 训练 $\Sigma_\theta$ 。一旦 $p_\theta$ 训练完成，可以通过初始化 $x_{t_{\max }} \sim \mathcal{N}(0, \mathbf{I})$ 并通过重参数化技巧采样 $x_{t-1} \sim p_\theta\left(x_{t-1} \mid x_t\right)$ 来采样新图像。

无分类器指导。条件扩散模型将额外的信息作为输入，例如类别标签 $c$ 。在这种情况下，反向过程变为 $p_\theta\left(x_{t-1} \mid x_t, c\right)$ ，其中 $\epsilon_\theta$ 和 $\Sigma_\theta$ 以 $c$ 为条件。在这种情况下，无分类器指导可用于鼓励采样过程找到使 $\log p(c \mid x)$ 较高的 $x$ [21]。根据贝叶斯定理， $\log p(c \mid x) \propto \log p(x \mid c)-\log p(x)$ ，因此 $\nabla_x \log p(c \mid x) \propto \nabla_x \log p(x \mid c)-\nabla_x \log p(x)$ 。通过将扩散模型的输出解释为得分函数，可以通过以下方式指导DDPM采样过程采样具有高 $\mid c)$ 的 $x$ ： $\hat{\epsilon}_\theta\left(x_t, c\right)=\epsilon_\theta\left(x_t, \emptyset\right)+s \cdot \left(\epsilon_\theta\left(x_t, c\right)-\epsilon_\theta\left(x_t, \emptyset\right)\right)$ ，其中 $s > 1$ 表示指导的规模（注意 $s = 1$ 恢复标准采样）。通过在训练期间随机丢弃 $c$ 并用学习的“空”嵌入 $\emptyset$ 替换它来实现对 $c=\emptyset$ 的扩散模型的评估。无分类器指导广为人知，可显著改善样本质量，相较于通用采样技术[21,35,46]，这一趋势在我们的DiT模型中也得以保持。

潜在扩散模型。直接在高分辨率像素空间中训练扩散模型可能会在计算上令人望而却步。潜在扩散模型（LDMs）[48]通过两阶段方法解决了这个问题：（1）学习一个自动编码器，将图像压缩成更小的空间表示，使用学习到的编码器 $E$ ；（2）训练表示 $z = E (x)$ 的扩散模型，而不是图像 $x$ 的扩散模型（ $E$ 是冻结的）。然后可以通过从扩散模型中采样表示 $z$ ，并随后用学习到的解码器将其解码为图像 $x = D (z)$ 来生成新图像。

如图2所示，LDMs在使用比像素空间扩散模型如ADM少得多的Gflops的情况下，实现了良好的性能。由于我们关注计算效率，这使它们成为架构探索的一个有吸引力的起点。在本文中，我们将DiTs应用于潜在空间，尽管它们也可以无需修改地应用于像素空间。这使我们的图像生成管道成为一种基于混合的方法；我们使用现成的卷积VAE和基于transformer的DDPMs。

图 2：使用 Diffusion Transformers (DiTs) 生成 ImageNet。气泡区域表示扩散模型的flops。左图：我们的 DiT 模型在 400K 次训练迭代中的 FID-50K（越低越好）。随着模型失败次数的增加，FID 的性能稳步提高。右图：我们最好的模型 DiT-XL/2 具有计算效率，并且优于所有之前基于 U-Net 的扩散模型，例如 ADM 和 LDM。

3.2. Diffusion Transformer Design Space

我们引入了扩散变换器 (DiT)，这是一种新的扩散模型架构。我们的目标是尽可能忠实于标准变换器架构，以保留其缩放属性。由于我们的重点是训练图像的 DDPM（具体来说，是图像的空间表示），因此 DiT 基于对patch序列进行操作的视觉变换器 (ViT) 架构 [10]。DiT 保留了 ViT 的许多最佳实践。图 3 显示了完整 DiT 架构的概览。在本节中，我们将描述 DiT 的前向传递，以及 DiT 类设计空间的组成部分。

图 3：扩散变换器 (DiT) 架构。左图：我们训练条件潜在 DiT 模型。输入潜在被分解为块并由多个 DiT 块处理。右图：我们的 DiT 块的细节。我们尝试了标准变换器块的变体，这些变体通过自适应层规范、交叉注意和额外的输入标记结合了条件。自适应层规范效果最好。

切块（Patchify）。DiT的输入是一个空间表示 $z$ （对于 $256 \times 256 \times 3$ 的图像， $z$ 的形状为 $32 \times 32 \times 4$ ）。DiT的第一层是“切块”，它通过线性嵌入输入中的每个patch，将空间输入转换为一个 $T$ 个token的序列，每个token的维度为 $d$ 。在切块之后，我们对所有输入token应用标准的ViT频率基位置嵌入（正弦-余弦版本）。切块生成的token数量 $T$ 由patch大小的超参数 $p$ 决定。如图4所示，将 $p$ 减半会使 $T$ 增加四倍，因此至少会使transformer的总Gflops增加四倍。尽管它对Gflops有显著影响，但注意到改变 $p$ 对下游参数数量没有有意义的影响。我们将 $p = 2, 4, 8$ 添加到DiT设计空间。

图4：DiT的输入规格。给定patch大小 $\times p$ ，形状为 $\times I \times C$ 的空间表示（来自VAE的噪声潜变量）被“切块”成一个长度为 $T=(I / p)^2$ 的序列，隐藏维度为 $d$ 。较小的patch大小 $p$ 会导致序列长度更长，从而增加Gflops。

DiT 块设计。在 patchify 之后，输入标记由一系列转换器块处理。除了噪声图像输入外，扩散模型有时还会处理其他条件信息，例如噪声时间步长 t、类标签 c、自然语言等。我们探索了四种以不同方式处理条件输入的转换器块变体。这些设计对标准 ViT 块设计进行了细微但重要的修改。所有块的设计如图 3 所示。

上下文条件。我们只需将 t 和 c 的向量嵌入作为两个附加标记附加到输入序列中，将它们与图像标记处理在一起。这类似于 ViT 中的 cls 标记，它允许我们使用标准 ViT 块而无需修改。在最后一个块之后，我们从序列中删除条件标记。这种方法为模型引入了可忽略不计的新 Gflops。
交叉注意力模块。我们将 t 和 c 的嵌入连接成一个长度为 2 的序列，与图像标记序列分开。Transformer 模块经过修改，在多头自注意力模块之后包含一个额外的多头交叉注意力层，类似于 Vaswani 等人 [60] 的原始设计，也类似于 LDM 用于条件化类标签的设计。交叉注意力为模型增加了最多的 Gflops，大约 15% 的开销。
自适应层归一化 (adaLN) 块。继自适应归一化层 [40] 在 GANs [2,28] 和具有 UNet 主干的扩散模型 [9] 中的广泛应用之后，我们探索用自适应层归一化 (adaLN) 替换 transformer 块中的标准层归一化层。我们不是直接学习维度缩放和平移参数 $\gamma$ 和 $\beta$ ，而是从 $t$ 和 $c$ 的嵌入向量之和回归这些参数。在我们探索的三种块设计中，adaLN 增加的 Gflops 最少，因此是最计算高效的。它也是唯一一种限制对所有 token 应用相同函数的条件机制。

adaLN-Zero 块。以前关于 ResNets 的研究发现，将每个残差块初始化为恒等函数是有益的。例如，Goyal 等人发现，在监督学习设置中零初始化每个块中最终批归一化缩放因子 $\gamma$ 加速了大规模训练 [13]。扩散 U-Net 模型使用类似的初始化策略，在任何残差连接之前零初始化每个块中的最终卷积层。我们探索了 adaLN DiT 块的一种修改，它也执行相同的操作。除了回归 $\gamma$ 和 $\beta$ ，我们还回归维度缩放参数 $\alpha$ ，该参数在 DiT 块内的任何残差连接之前立即应用。我们将MLP初始化为对所有 $\alpha$ 输出零向量；这将整个DiT块初始化为恒等函数。与普通adaLN块一样，adaLNZero为模型增加的Gflops可以忽略不计。

我们在DiT设计空间中包括上下文内、交叉注意力、自适应层归一化和adaLN-Zero块。

模型大小。我们应用一系列 $N$ 个DiT块，每个块都在隐藏维度大小 $d$ 上运行。按照ViT的方式，我们使用共同扩展 $N$ 、 $d$ 和注意力头的标准transformer配置[10,63]。具体来说，我们使用四种配置：DiT-S、DiT-B、DiT-L和DiT-XL。它们涵盖了从0.3到118.6 Gflops的广泛模型大小和浮点运算分配范围，使我们能够评估扩展性能。表1给出了配置的详细信息。我们在DiT设计空间中添加B、S、L和XL配置。

Transformer解码器。在最终的DiT块之后，我们需要将图像token序列解码为输出噪声预测和输出对角协方差预测。两个输出的形状都与原始空间输入相等。我们使用标准线性解码器来完成这一操作；我们应用最后的层归一化（如果使用adaLN则为自适应的）并将每个token线性解码为 $\times p \times 2C$ 张量，其中 $C$ 是DiT的空间输入的通道数。最后，我们将解码后的token重新排列为其原始的空间布局，以获得预测的噪声和协方差。