扩散变压器：开创图像生成新纪元

 在深度学习领域，变换器（Transformer）架构已经成为自然语言处理、视觉识别等多个领域的核心技术。最近，一项新的研究探索了基于变换器的一类新型扩散模型——扩散变压器（Diffusion Transformers，简称DiTs），这些模型在图像生成任务中展现出了前所未有的性能。

Diffusion Transformers（扩散变压器）

在深度学习中，变换器架构因其在处理序列数据方面的强大能力而闻名，特别是在自然语言处理（NLP）领域取得了巨大成功。然而，在图像生成领域，尤其是扩散模型中，U-Net架构一直是主要使用的骨干网络。DiTs的出现打破了这一传统，引入了变换器作为核心组件，以期利用其在视觉任务中表现出的卓越性能。

扩散模型，作为一种生成模型，其工作原理是通过一个由噪声逐渐增加的过程来生成数据，然后训练一个模型来逆转这一过程，从而生成高质量的图像。在传统的扩散模型中，U-Net作为一种卷积神经网络（CNN），负责逐步去除噪声并恢复图像。U-Net的结构通常包含多个卷积层和残差连接，这些设计有助于捕捉图像的局部特征并增强模型的表征能力。

DiTs的核心思想是将变换器架构应用于这一过程中，特别是在图像的潜在空间中进行操作。潜在空间是指通过某种编码方式（如自编码器）将原始图像数据映射到的低维表示空间。在DiTs中，原始图像首先被编码为潜在表示，然后这些潜在表示被分割成小块，即“patches”，这些小块随后被用作变换器的输入。

变换器架构的一个关键优势是其能够处理长距离依赖关系，这在图像中可以对应于全局上下文信息。与传统的卷积层相比，变换器可以更有效地捕捉图像的全局结构和复杂的模式。在DiTs中，通过自注意力机制，每个图像块都能与输入中的其他所有块进行交互，这使得模型能够更好地理解和生成图像的全局结构。

DiTs的设计还考虑了计算效率和模型的扩展性。通过在潜在空间而不是原始像素空间中进行操作，DiTs能够显著降低计算成本，同时保持或提高生成图像的质量。这一点在大规模图像生成任务中尤为重要，因为它允许模型在保持高效性的同时生成高分辨率的图像。

DiTs的设计哲学深受视觉变换器（Vision Transformers，ViTs）的影响，ViTs在处理图像识别任务时，表现出了超越传统卷积网络的能力。ViTs之所以能够取得这样的成就，很大程度上归功于其独特的处理图像的方式：它将图像划分为一系列的小块（patches），并将这些小块视为输入序列，通过自注意力机制来捕捉图像中的全局上下文信息。这种方法使得ViTs能够有效地处理图像中的长距离依赖关系，这对于理解复杂的视觉场景至关重要。

在DiTs中，这种处理图像的方式被进一步发扬光大。DiTs采用了与ViTs相似的策略，将图像编码到潜在空间中，然后在这个空间内进行操作。具体来说，DiTs利用变分自编码器（VAE）来将高维的图像数据压缩成一个低维的潜在表示，这个潜在表示通常具有更小的空间尺寸，但仍然保留了图像的重要信息。通过这种方式，DiTs能够在一个更紧凑的表示空间内进行图像生成，这大大减少了模型的计算负担。

在潜在扩散模型（Latent Diffusion Models，LDMs）的框架下，DiTs的训练过程被进一步优化。LDMs是一种两阶段的方法，首先通过学习一个自编码器来压缩图像，然后在得到的潜在表示上训练扩散模型。这种方法的优势在于，它允许模型在较低维度的空间中进行操作，从而减少了模型的参数数量和计算复杂度。同时，由于潜在空间的连续性，扩散过程可以更加平滑地进行，这有助于生成更高质量的图像。

在DiTs的训练过程中，模型首先需要学习如何将输入的噪声数据逐步转换为潜在表示中的图像。这个过程中，模型需要预测每一步中噪声的分布，以便在逆过程中准确地重建图像。随着训练的进行，DiTs逐渐学会了如何有效地在潜在空间中导航，从而生成高质量的图像。

为了模型的扩展性通过增加变换器的深度和宽度，或者通过增加输入token的数量，DiTs能够灵活地扩展其容量，以适应不同复杂度的任务。这种设计使得DiTs不仅在小规模的数据集上表现出色，也能够在大规模的数据集上发挥出强大的性能。

图3展示了扩散变压器（DiT）的架构设计，这一部分是理解DiT如何运作的核心。图的左侧描述了条件潜在DiT模型的训练过程，而右侧则详细展示了DiT块的设计细节。

条件潜在DiT模型的训练

在图3的左侧，我们可以看到条件潜在DiT模型的训练涉及将输入的潜在表示（latent）分解为多个小块（patches），这些小块随后被送入多个DiT块中进行处理。这个过程是DiT生成高质量图像的关键步骤。

• 输入潜在表示的分解：在潜在扩散模型（LDMs）框架下，原始图像首先被编码为一个潜在表示，这个表示通常是高维空间中的一个点。在DiT中，这个连续的潜在表示被进一步分解为一系列的小块，这样做的目的是为了降低处理的复杂度，同时保留图像的重要特征信息。
• DiT块的处理：每个小块随后被送入一系列DiT块中。每个DiT块都是一个变换器模块，负责进一步处理这些小块，以生成最终的图像数据。

DiT块的设计细节

图3的右侧深入介绍了DiT块的设计，特别是它们如何通过不同的条件化技术来增强模型的性能。

• 自适应层归一化（adaLN）：在实验中，研究者们尝试了多种标准变换器块的变体，其中包括通过自适应层归一化来实现条件化的方法。adaLN是一种特殊的层归一化技术，它可以根据条件信息动态调整归一化的参数。在DiT中，adaLN通过将条件信息（如类别标签或时间步长）的嵌入向量与小块的特征相结合，来动态调整每个小块的处理。
• 交叉注意力：除了adaLN，研究者们还尝试了交叉注意力机制，这是一种允许模型在处理小块时考虑到其他小块信息的技术。这种方法可以帮助模型更好地理解图像的不同部分之间的关系，从而生成更加连贯和真实的图像。
• 额外输入token：在某些变体中，DiT块还加入了额外的输入token，这些token可以携带额外的条件信息，以指导图像的生成过程。

DiTs的架构精妙地融合了视觉变换器（ViTs）的高效实践和潜在扩散模型（LDMs）的优化框架，通过将变换器架构应用于潜在空间，显著提升了图像生成的效率和质量。这种创新的设计策略，在减少模型计算负担的同时，增强了生成图像的精细度，为图像生成技术开辟了崭新的发展方向。

在DiTs的设计空间内，变换器块的设计、模型规模的大小以及潜在图像块的尺寸是几个关键维度，它们相互协作，共同塑造了DiT在图像生成任务中的性能表现。通过细致调整这些参数，研究者们能够针对不同的应用需求，定制出性能最优化的模型配置。这种灵活性和可调性，使得DiT不仅在理论上具有创新性，而且在实际应用中也展现出了强大的适应性和实用性。

在变换器块的设计方面，DiT采用了几种不同的方法来处理条件输入，如自适应层归一化（adaLN）和交叉注意力机制。自适应层归一化是一种特殊类型的层归一化，它通过学习数据的统计特性来调整每个隐藏层的缩放和平移参数。这种方法使得模型能够更加灵活地适应不同的输入特征，从而提高性能。在DiT中，adaLN通过将条件信息（如时间步长和类别标签）的嵌入向量相加并输入到激活函数和线性层来工作，生成对每个隐藏层进行调整的参数。

交叉注意力机制则是另一种处理条件输入的方法，它允许模型在处理图像块时考虑到额外的条件信息。在DiT中，交叉注意力通过将条件信息的嵌入向量与图像块序列分开处理，并在变换器块中加入额外的多头交叉注意力层来实现。这种方法使得模型能够在生成图像时，更好地利用条件信息来指导图像的生成过程。

模型大小是影响DiT性能的另一个重要因素。模型大小通常由其层数、隐藏层的维度以及注意力头的数量来决定。较大的模型通常能够捕捉更复杂的特征和模式，但同时也需要更多的计算资源。在实验中，研究者们探索了不同大小的DiT模型，从较小的模型到较大的模型，并评估了它们在图像生成任务上的表现。

潜在图像块的大小也是DiT设计中的关键考虑因素。在将图像编码到潜在空间后，图像被划分为较小的块，这些块的大小直接影响到模型的计算效率和生成图像的质量。较小的块可以提供更细粒度的图像信息，但会增加序列的长度，从而增加计算复杂度。较大的块可以减少计算量，但可能会丢失一些细节信息。因此，选择合适的块大小对于平衡模型性能和计算效率至关重要。

通过系统地探索这些设计空间，研究者们能够确定最有效的DiT配置。例如，他们发现adaLN-Zero块在所有训练阶段都优于其他块设计，并且在计算效率上也更高。此外，通过减少潜在图像块的大小并增加模型大小，可以显著提高FID分数，这表明增加模型的计算量是提高性能的关键。这些发现为DiT的设计和优化提供了宝贵的见解，并为未来的研究指明了方向。

实验

实验部分，研究者们在ImageNet数据集上训练了不同配置的DiT模型，分辨率分别为256×256和512×512。所有模型均使用AdamW优化器进行训练，学习率为1×10^-4，没有权重衰减，批量大小为256。为了评估模型性能，使用了Fréchet Inception Distance（FID）作为主要指标，这是一种评估图像生成模型的标准度量。

实验结果显示，随着模型复杂性的增加（以Gflops衡量），DiT模型的FID分数显著下降，表明模型的生成质量随着计算量的增加而提高。特别是，DiT-XL/2模型在256×256分辨率的ImageNet基准测试中取得了2.27的FID分数，这是当时的最佳结果。

研究者们还比较了不同DiT块设计的性能，发现adaLN-Zero块在所有训练阶段均优于其他块设计，并且在计算效率上也更高。此外，通过减少潜在图像块的大小并增加模型大小，可以显著提高FID分数，这表明增加模型的计算量是提高性能的关键。

在更大规模的模型和更高分辨率的实验中，DiT-XL/2模型同样展现出了卓越的性能，其在512×512分辨率的ImageNet基准测试中取得了3.04的FID分数，超越了之前的最佳模型。

图6在论文中展示了扩散变压器（DiT）模型规模扩展对Fréchet Inception Distance（FID）指标的影响，FID是评估生成模型性能的一个重要指标，其值越低表示生成的图像质量越高。该图通过跟踪12种不同配置的DiT模型在50,000次迭代过程中的FID-50K（即在50,000个样本上的FID）变化，揭示了模型规模扩展对性能的积极作用。

图中的上半部分固定了潜在图像块的大小，仅改变模型的大小（即DiT块的数量和模型层的深度与宽度），展示了不同模型规模对FID的影响。而下半部分则固定了模型的大小，仅改变潜在图像块的大小，来展示块的大小对FID的影响。实验结果表明，在所有模型尺寸和块尺寸的组合中，扩展变换器骨干网络的规模能够显著提高生成模型的性能，即生成更高质量的图像。

研究者们在所有12种不同配置的DiT模型上进行了采样，这些模型经过了400,000次训练迭代。为了公平比较，所有模型使用相同的输入潜在噪声和类别标签进行图像生成。随着模型中Gflops的增加，即通过增加变换器的深度和宽度或增加输入token的数量，生成图像的视觉保真度得到了显著提升。这意味着更大的模型计算量能够带来更高质量的图像生成结果。这种提升在放大查看时尤为明显，因为它允许观察到图像中的细节和清晰度的改善。

扩散变压器（DiTs）的研究为图像生成领域带来了新的突破，证明了变换器架构在这一任务中的有效性。随着计算资源的不断增长，未来的工作可以继续扩展DiT模型的规模和潜在空间的容量，进一步推动图像生成技术的发展。DiTs也有可能被用作文本到图像模型的骨干网络，如DALL·E 2和Stable Diffusion等，为这些模型带来性能上的飞跃。

论文链接：https://arxiv.org/pdf/2212.09748