DiT - 基于Transfomer架构的扩散模型

DiT是什么

DiT（Diffusion Transformers）是一种新型的扩散模型，由William Peebles（Sora的研发负责人之一） 与纽约大学助理教授谢赛宁提出，结合了去噪扩散概率模型（DDPMs）和Transformer架构。扩散模型是一种生成模型，通过模拟数据的逐步去噪过程来生成新的样本。DiT的核心思想是使用Transformer作为扩散模型的骨干网络，而不是传统的卷积神经网络（如U-Net），以处理图像的潜在表示。近期伴随OpenAI视频生成模型Sora的大热，DiT被视为Sora背后的技术基础之一而广受关注。

在DiT中，图像首先通过一个自动编码器（如变分自编码器VAE）被压缩成较小的潜在表示，然后在这个潜在空间中训练扩散模型。这样做的好处是可以减少直接在高分辨率像素空间训练扩散模型所需的计算量。DiT模型通过Transformer的自注意力机制来处理这些潜在表示，这使得模型能够捕捉到图像的长距离依赖关系，从而生成高质量的图像。

DiT的官网入口

• 官方项目主页：https://www.wpeebles.com/DiT
• Arixv研究论文：https://arxiv.org/pdf/2212.09748.pdf
• GitHub代码库：https://github.com/facebookresearch/DiT
• Hugging Face空间：https://huggingface.co/spaces/wpeebles/DiT
• Replicate Demo：https://replicate.com/arielreplicate/scalable_diffusion_with_transformers
• Google Colab运行地址：http://colab.research.google.com/github/facebookresearch/DiT/blob/main/run_DiT.ipynb

DiT的技术原理

1. 数据准备：
   • 使用一个预训练的变分自编码器（VAE）将输入图像编码成潜在空间的表示。这个潜在表示通常是图像的低维表示，例如，将256×256×3的RGB图像编码成32×32×4的潜在向量。
   • 这个潜在表示随后被用作DiT模型的输入。
2. 分块化（Patchification）：
   • 输入的潜在表示首先通过一个称为“patchify”的过程，将其分割成一系列的小片段（patches），每个片段对应于Transformer模型的一个输入标记（token）。这个过程类似于将图像分割成小块，以便模型可以逐块处理。
   • 每个片段通过线性嵌入转换为一个固定维度的向量，然后添加位置编码（positional embeddings），以便模型能够理解片段在图像中的位置。
3. Transformer Blocks模块：
   • 输入的标记序列通过一系列的Transformer块进行处理。这些块包括自注意力层、前馈神经网络以及层归一化等组件。
   • 在DiT中，研究者们尝试了不同的Transformer块设计，包括自适应层归一化（adaLN）、交叉注意力（Cross-Attention）和上下文条件（In-Context Conditioning）等，以处理条件信息，如时间步长（timesteps）和类别标签（class labels）。
4. 条件扩散过程：
   • 在训练过程中，DiT模型学习逆向扩散过程，即从噪声数据中恢复出清晰的图像。这个过程涉及到预测噪声的统计特性，如均值（mean）和方差（covariance）。
   • 使用变分下界（Variational Lower Bound, VLB）来训练模型，这涉及到最小化预测噪声和真实噪声之间的均方误差（MSE）。
5. 样本生成：
   • 在训练完成后，可以通过DiT模型生成新的图像。首先，从标准正态分布中采样一个潜在表示，然后通过DiT模型逆向扩散过程，逐步去除噪声，最终解码回像素空间，得到生成的图像。
6. 可扩展性：
   • DiT模型的可扩展性体现在通过增加Transformer的层数、宽度或输入标记的数量来提高模型的计算量（Gflops），从而提高生成图像的质量。这种可扩展性使得DiT模型能够在不同的分辨率和复杂度下生成高质量的图像。

DiT模型通过这种方式，利用Transformer的强大表达能力和扩散模型的生成能力，实现了在图像生成任务中的高效和高质量输出。

DiT的关键特点

• 基于Transformer的架构：DiT采用了Transformer作为其核心架构，这使得模型能够处理图像的序列化表示，并通过自注意力机制捕捉图像中的长距离依赖关系。
• 潜在空间操作：DiT在潜在空间中训练，这通常比直接在像素空间训练更高效。通过使用变分自编码器（VAE）将图像编码到潜在空间，DiT减少了计算复杂度。
• 可扩展性：DiT展示了出色的可扩展性，通过增加模型的计算量（以Gflops衡量），可以显著提高生成图像的质量。这种可扩展性允许DiT在不同的分辨率和复杂度下生成图像。
• 条件生成能力：DiT支持条件生成，能够根据给定的类别标签生成特定类别的图像。这种能力使得DiT在特定领域的图像生成任务中非常有用。
• 自适应层归一化（adaLN）：DiT使用了自适应层归一化技术，这是一种在Transformer块中使用的归一化方法，通过学习来调整层归一化的参数，从而提高模型的表达能力和训练效率。
• 多种Transformer块设计：DiT探索了不同的Transformer块设计，包括自适应层归一化（adaLN）、交叉注意力（Cross-Attention）和上下文条件（In-Context Conditioning），以处理条件信息。
• 高效的训练过程：DiT在训练过程中表现出高度的稳定性，即使在没有使用学习率预热和正则化技术的情况下，也能稳定地训练到高性能。
• 生成图像的多样性和质量：DiT能够生成具有高视觉质量和多样性的图像。通过调整类条件生成的指导强度，可以在生成图像的清晰度和多样性之间进行权衡。
• 高计算效率：在生成图像时，DiT能够在保持高图像质量的同时，实现较高的计算效率。这使得DiT在资源有限的环境中也具有吸引力。
• 应用潜力：DiT在图像生成领域具有广泛的应用潜力，包括艺术创作、游戏开发、虚拟现实、数据增强等，尤其是在需要生成高质量图像的场景中。