Sora核心之一：当Diffusion遇到Transformer，DiT的前世今生_transformer diffusion

Diffusion与Tranformer结合的代表性文章（DiT的前世今生）

一、DiT

题目: Scalable Diffusion Models with Transformers
机构：UC Berkeley，纽约大学
论文: https://arxiv.org/pdf/2212.09748.pdf
代码：https://github.com/facebookresearch/DiT
任务: 图像生成
特点: 将Diffusion中常常用到的UNet替换为transformer结构，并分析模型复杂度以及生成质量的关系，相比于UNet在效率上的优势，以及模型架构的可扩展性
前置相关工作：DDPM，Improved DDPM，ADM，LDM，Classifier-free guidance，Parti
同期类似工作：UViT
后续衍生以及改进工作：HDiT，SiT，FiT，PIXART-a
王炸级旗舰产品：Sora，SD3

围绕DiT的前世今生，重要的相关工作可以梳理如下，截止2024.03.02

[1]Denoising Diffusion Probabilistic Models Ho et al, UC Berkeley 2020.12
[2]Improved Denoising Diffusion Probabilistic Models Nichol et al, OpenAI 2021.02
[3]Diffusion Models Beat GANs on Image Synthesis Nichol et al, OpenAI 2021.06
[4]High-Resolution Image Synthesis with Latent Diffusion Models Rombach et al, CompVis Runway 2021.12 [CVPR2022]
[5]CLASSIFIER-FREE DIFFUSION GUIDANCE Ho et al, Google 2022.07
[6]Scaling Autoregressive Models for Content-Rich Text-to-Image Generation Yu et al, Google 2022.06
[b]All are Worth Words: A ViT Backbone for Diffusion Models Bao et al, Tsinghua University 2022.09 [CVPR2023]
[7]Scalable Diffusion Models with Transformers Peebles et al, UC Berkeley, NYU V1 2022.12, V2 2023.03 [ECCV2023]
[8]PIXART-α: FAST TRAINING OF DIFFUSION TRANSFORMER FOR PHOTOREALISTIC TEXT-TO-IMAGE SYNTHESIS Chen et al, Huawei, 2023.12
[9]Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers Crowson et al, Stability AI 2024.01
[10]SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers Ma et al, NYU, 2024.01
[11]https://openai.com/research/video-generation-models-as-world-simulators Brooks et al, OpenAI 2024.02
[12]https://stability.ai/news/stable-diffusion-3 Stability AI 2024.02
[13]FiT: Flexible Vision Transformer for Diffusion Model Lu et al, Shanghai Artificial Intelligence Laboratory, 2024.02

1、动机

Transformer结构在自然语言处理以及视觉等领域取得了广泛应用，但是在pixel-level的生成任务上，往往和自回归方法相结合，因此本文将Transformer架构与Diffsuion生成方法相结合，提出基于transformer的扩散模型（Diffusion Transformers，简称DiTs），并分析这些模型在不同复杂度下的可扩展性和性能。具体来说，论文探讨了以下几个关键问题：

1. Transformer rather than UNet：论文研究了将Transformer架构作为扩散模型（特别是去噪扩散概率模型，DDPMs）的骨干网络（backbone）的可行性，以替代传统的基于U-Net的架构。
2. 模型复杂度与生成质量的关系：通过构建和评估DiT模型的设计空间，论文分析了模型复杂度（以GFLOPS为度量）与生成样本质量（以Fréchet Inception Distance，FID为度量）之间的关系。
3. DiT模型的可扩展性：论文展示了通过增加Transformer的深度、宽度或输入tokens数量（即增加GFLOPS）可以显著降低FID，从而提高生成图像的质量。
4. DiT模型在特定任务上的性能：论文在ImageNet数据集上训练DiT模型，并在512×512和256×256分辨率的条件下，与现有的扩散模型进行性能比较，展示了DiT模型在这些任务上的优越性。
5. DiT模型的计算效率：论文还探讨了DiT模型相对于其他基于U-Net的扩散模型（如ADM和LDM）在计算效率方面的表现。

2、相关工作

A: 这篇论文提到了多个与扩散模型（Diffusion Models）、Transformers、以及图像生成相关的研究工作。以下是一些关键的相关研究：

3、论文方法

3.1、DiT VS LDM的核心不同

其中上面的condition包含时间步t和类别信息y，对于时间步而言，使用的是Transformer中介绍的不可学习的绝对位置编码，类别信息而言，单独训练了一个科学系的embedding层，然后c = t + y

3.2、DiT方法详解

整体提出了一种新的扩散模型架构，称为Diffusion Transformers（DiTs）。这种架构基于Vision Transformers（ViTs），将图像的潜在表示（latent patches）作为输入，并通过一系列Transformer块进行处理。

3.2.1、Patchify

• 比如输入$224\ast 224$的图像，经过VAE编码得到$32\ast 32\ast 4$的隐藏层特征，然后输入后面的transformer结构
• 第一步就是类似ViT一样的patchify操作，将上面的空间输入转化为长度为T的序列tokens，其中tokens的位置信息，用不可学习的绝对位置编码（sin/cos），这儿主要探讨了通过调整patch size，实现对序列长度的控制，进而影响整个DiT的复杂度（GFLOPS），我们可以知道的是，当patch size减半时，比如从4 -> 2，我们的tokens序列长度是变为了原来的4倍，如果对ViT比较熟悉的话，可以推导整个网络的GFLOPS（关于MSA，FFN等模块的复杂度与序列长度s，hidden size h的关系），至少应该GFLOPS应该会是原来的4倍，值得注意的是，这儿tokens的数目增加，对参数量是没有影响的（更准确的说法是影响微乎其微，可以忽略不计），一个侧面也反映了本文选择用GLOPS而不是参数量来当作模型复杂度的测度标准的原因。

3.2.2、条件信息注入

条件化策略：为了处理额外的条件信息（如噪声时间步、类别标签等），论文探索了不同的Transformer块设计，包括in-context conditioning、cross-attention、adaptive layer norm（adaLN）和adaLN-Zero。这些设计允许模型在不改变标准ViT块的情况下，有效地利用条件信息。其实对于条件信息的注入，这儿的四种方法设计，DiT本身的创新性不大，唯一相对又一点创新性的就是adaLN-zero，对于adaLN而言，它是利用上文提到的条件信息c，去学习预测layer norm的缩放和平移参数$\gamma$和$\beta$，而adaLN-zero而言，是额外预测了门控参数$\alpha$，并且初始化门控为0，实现类似Resnet或者ControlNet那样的zero初始化，使训练更稳定以及效果更优。

如果对adaLN不太了解的，可以去看看adaIN这篇经典的工作：Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization，这篇论文的主要目标是实现实时的、任意风格的风格迁移（style transfer），主要方法就是自适应实例标准化（Adaptive Instance Normalization，AdaIN），将内容图像（content image）特征的均值和方差对齐到风格图像（style image）的均值和方差。


就adaLN而言，其实就是把instance normalization换成layer normalization，在传统在BN，IN，CIN中，网络会学习仿射变换参数$\gamma$ 和$\beta$ ，AdaIN则无需学习这两个参数，直接用style image的特征的均值和标准差代替这两个参数，而AdaLN，是利用条件信息c来预测回归这两个参数，具体的代码实现如下：

上述就是adaLN-Zero的具体实验，可以发现，它是把LayerNorm层的element_affine的参数给置为了false，如果设为False，则LayerNorm层不含有任何可学习参数，如果设为True（默认是True）则会包含可学习参数weight和bias，用于仿射变换，即对输入数据归一化到均值0方差1后，乘以weight，再加上bias，所以上面就是在用c来进行回归预测，具体可以参见forward实现，最终实现，输入c，然后输出shift，scale，gate的预测，值得注意的是分别预测了msa以及mlp两个层的adaLN，并且不同DiT block之间是不共享的对于adaLN modulation的参数初始化，可以发现，其实现了zero initialization。

2. 模型复杂度与性能分析：论文通过构建DiT模型的设计空间，包括不同的模型大小（如DiT-S、DiT-B、DiT-L和DiT-XL）和输入tokens的数量（通过改变patchify层的patch size），来研究模型复杂度（以GFLOPS为度量）与样本质量（以FID为度量）之间的关系。
   
   就实现而言，有一个比较有意思的点，就是关于cfg的实现：
   

4、论文实验

A: 论文进行了一系列的实验来探索和验证DiT模型的性能和可扩展性。以下是主要的实验内容：

1. DiT块设计实验：不同的块设计：in-context conditioning、cross-attention、adaptive layer norm (adaLN) 和 adaLN-Zero，其中最佳是adaLN-Zero

2. 模型大小和patch size的缩放实验：论文训练了12个DiT模型，覆盖了不同的模型配置（S、B、L、XL）和patch sizes（8、4、2）。这些实验旨在研究模型大小（通过增加Transformer层数和隐藏层大小）和输入令牌数量（通过减小patch size）对FID的影响。

3. DiT模型与现有扩散模型的比较：在ImageNet数据集上，论文将DiT模型与现有的扩散模型（如ADM、LDM等）进行了性能比较。这包括在256×256和512×512分辨率下的FID、sFID、Inception Score、Precision和Recall等指标。

4. 计算效率分析：
   论文分析了DiT模型在训练过程中的计算效率，包括训练损失曲线和训练迭代次数。此外，还比较了DiT模型与基于U-Net的模型在采样计算量和模型计算量方面的差异。

5. VAE解码器的消融研究：为了评估VAE解码器对DiT模型性能的影响，论文在不同的预训练VAE解码器（如ft-MSE和ft-EMA）之间进行了比较。

6. DiT模型的可视化：论文展示了DiT模型生成的图像样本，以及在不同分类器自由引导（classifier-free guidance）尺度下的样本质量。

7. DiT模型的长期训练：对于DiT-XL/2模型，论文进行了长达7百万步的训练，并观察了FID随训练步骤的变化。

这些实验不仅验证了DiT模型在图像生成任务中的有效性，还揭示了模型复杂度与生成图像质量之间的强相关性，以及DiT模型相对于传统U-Net模型在计算效率方面的优势。

5、Future work

尽管论文已经取得了显著的成果，但仍有许多潜在的研究方向可以进一步探索：
更大的模型规模： 论文展示了DiT模型在增加计算量（GFLOPS）时性能的提升。未来的研究可以继续扩大模型规模，探索更大的模型是否能够进一步提高图像生成的质量。
不同的条件化策略：虽然论文已经尝试了几种条件化策略，但可能还有其他有效的方法可以进一步改善DiT模型的性能。例如，可以尝试结合多种策略或开发新的条件化机制。
跨领域应用： DiT模型在图像生成领域表现出色，但其架构的通用性意味着它可能适用于其他领域，如文本到图像生成、视频生成等。研究DiT在这些领域的应用是一个有趣的方向。
训练效率和稳定性： 尽管DiT模型在性能上取得了突破，但训练大型模型仍然是一个挑战。研究如何提高训练效率和稳定性，例如通过改进优化算法或引入正则化技术，是值得探索的。
模型解释性和可解释性： 理解DiT模型的内部工作机制，以及如何解释生成的图像，对于提高模型的透明度和可信度至关重要。研究模型解释性可以帮助用户更好地理解和信任生成的图像。
数据集和任务的多样性： 论文主要在ImageNet数据集上进行了实验。将DiT模型应用于更多样化的数据集和任务，如医学成像、艺术风格迁移等，可以进一步验证其泛化能力。
模型压缩和部署： 为了在资源受限的环境中部署DiT模型，研究模型压缩和加速技术，如知识蒸馏、量化和剪枝，是必要的。
与人类创造力的结合： 探索如何将DiT模型与人类的创造力相结合，例如通过交互式生成或辅助设计工具，可以开辟新的应用场景。

这些方向不仅有助于推动DiT模型本身的发展，还可能为整个生成模型领域带来新的突破。

6、总结

背景与动机：论文指出，尽管Transformer在自然语言处理和视觉领域取得了巨大成功，但在图像生成模型中的应用相对较少。特别是，扩散模型（如Denoising Diffusion Probabilistic Models, DDPMs）通常采用基于U-Net的架构。作者旨在探索将Transformer架构引入扩散模型的潜力。

DiT架构：论文提出了DiT架构，它遵循Vision Transformers（ViTs）的最佳实践，将图像的潜在表示（latent patches）作为输入，并通过一系列Transformer块进行处理。DiT架构包括patchify层、DiT块（包含不同的条件化策略）、模型层和Transformer解码器。

实验设计：作者在ImageNet数据集上训练了不同配置的DiT模型，并与现有的扩散模型进行了比较。实验包括分析模型复杂度（GFLOPS）与样本质量（FID）之间的关系，以及DiT模型在不同分辨率下的性能。

主要发现：

1. DiT模型在增加计算量（GFLOPS）时，样本质量（FID）显著提高。
2. 在256×256分辨率的ImageNet数据集上，DiT-XL/2模型达到了2.27的FID，这是当时的最佳性能。
3. DiT模型在计算效率上优于基于U-Net的模型，尤其是在训练和生成过程中。
4. 可扩展性分析：论文展示了DiT模型具有良好的可扩展性，通过增加模型大小或输入令牌数量，可以进一步提高生成图像的质量。

二、PIXART-α

题目： PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis
机构：华为
论文：https://arxiv.org/abs/2310.00426
代码：https://pixart-alpha.github.io/

这篇论文介绍了PIXART-α，这是一个基于Transformer的文本到图像（T2I）扩散模型，旨在解决现有T2I模型训练成本高和环境影响大的问题。以下是论文的主要内容总结：

1、问题背景：

当前先进的T2I模型，如Imagen、Stable Diffusion等，需要巨大的计算资源进行训练，这不仅增加了成本，还导致了显著的CO2排放。

2、PIXART-α的核心设计：

训练策略分解：

1. 将复杂的文本到图像生成任务分解为三个子任务：学习自然图像的像素分布、学习文本与图像的对齐以及提高图像的审美质量。
2. 在第一阶段，使用低成本的类别条件模型初始化T2I模型，显著降低学习成本。
3. 在第二阶段，通过在信息密度丰富的文本-图像对数据上进行预训练，然后对具有更高审美质量的数据进行微调，提高训练效率。

高效的T2I Transformer： 在Diffusion Transformer（DiT）的基础上，引入交叉注意力模块以注入文本条件，并简化计算密集的类别条件分支。提出了一种重参数化技术，允许调整后的文本到图像模型直接加载原始类别条件模型的参数，从而利用从ImageNet学到的关于自然图像分布的先验知识，为T2I Transformer提供合理的初始化并加速其训练。

高信息量数据：利用大型视觉-语言模型自动标记密集的伪标题，以提高文本-图像对齐学习。使用LLaVA（Liu et al., 2023）在SAM（Kirillov et al., 2023）数据集上生成描述，SAM数据集由于其丰富多样的对象集合，是创建高信息密度文本-图像对的理想资源。

通过这些设计，PIXART-α在保持与现有最先进图像生成器竞争力的同时，显著减少了训练成本和CO2排放。例如，PIXART-α的训练时间仅为Stable Diffusion v1.5的12%，训练成本仅为其1%，同时减少了90%的CO2排放。此外，与更大的SOTA模型RAPHAEL相比，PIXART-α的训练成本仅为其1%。这些成果表明，PIXART-α在图像质量、艺术性和语义控制方面表现出色，为AIGC社区和初创企业提供了新的见解，帮助他们以更低的成本从头开始构建自己的高质量生成模型。

3、实验结果：

PIXART-α在保持与现有最先进模型竞争力的同时，显著降低了训练成本和CO2排放。

在图像质量、艺术性和语义控制方面表现出色，用户研究显示其在质量和对齐方面优于现有SOTA模型。

对模型的关键修改进行了消融研究，验证了结构修改和重参数化设计的有效性。

4、 应用扩展：

展示了PIXART-α在不同应用场景下的生成能力，如与DreamBooth和ControlNet结合进行定制化图像生成。

5、 结论：

PIXART-α为AIGC社区和初创企业提供了新的见解，帮助他们以更低的成本构建高质量的T2I模型。

6、未来工作：

提出了进一步探索的方向，如模型扩展性、风格和内容控制、交互式生成、多模态输入、数据集和训练策略优化、环境影响优化、模型可解释性、跨领域应用以及模型压缩和加速。

整体而言，这篇论文通过提出PIXART-α模型，展示了在文本到图像生成领域实现高效、低成本和环境友好的训练方法的可能性。