DiT（Diffusion Transformer）详解——AIGC期间的新宠儿
https://i-blog.csdnimg.cn/direct/bc2cb0b989b84d85a6c84437c5796b72.png扩散模型相关知识点参考：小白也能读懂的AIGC扩散（Diffusion）模型系列解说
推荐阅读：
[*]一文带你搞懂DiT（Diffusion
Transformer）
[*]Sora 基础作品之 DiT：Scalable Diffusion Models with
Transformer
[*]此文DiT部分：视频生成Sora的全面剖析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等
部分摘录于：https://zhuanlan.zhihu.com/p/684068402
论文
Scalable Diffusion Models with Transformers
定义
DiT是基于Transformer架构的扩散模型。用于各种图像（SD3、FLUX等）和视频（Sora等）视觉生成任务。
DiT证实了Transformer思想与扩散模型结合的有效性，并且还验证了Transformer架构在扩散模型上具备较强的Scaling能力，在稳步增大DiT模型参数目与加强数据质量时，DiT的生成性能稳步提升。
此中最大的DiT-XL/2模型在ImageNet 256x256的种别条件生成上达到了当时的SOTA【最先进的（State Of The Art）】（FID为2.27）性能。同时在SD3和FLUX.1中也阐明了较强的Scaling能力。
架构
DiT架构如下所示：
https://i-blog.csdnimg.cn/direct/affe4797c41e4e779ba05241ca8c1064.png
图3.扩散Transformer（DiT）架构。左：我们训练条件潜在DiT模型。输入的潜在被分解成补丁和处置惩罚的几个DiT块。右图：DiT区块的详细信息。我们用标准Transformer块的变体进行了实验，这些块通过自适应层归一化、交叉留意和**额外输入的令牌（上下文环境）**来进行调节。自适应层规范结果最好。
[*] 左侧主要架构图：训练条件潜在DiT模型(conditional latent DiT models)， 潜在输入和条件被分解成patch并结合条件信息通过几个DiT blocks处置惩罚。本质就是噪声图片减掉预测的噪声以实现渐渐复原。
[*]DiT blocks前： 比如当输入是一张256x256x3的图片，得到32x32x4的Noised Latent，之后进行Patch和位置编码，结合当前的Timestep t、Label y作为输入。
[*]DiT blocks后： 经过N个Dit Block(基于transformer)通过MLP进行输出，在 DiT 模型 的最后一个 Transformer 块（DiT block）之后，需要将生成的图像 token 序列解码为以下两项输出：噪声“Noise预测”以及对应的协方差矩阵，最后经过T个step采样，得到32x32x4的降噪后的latent。
[*] 右侧DiT Block实现方式：DiT blocks的细节，作者试验了标准transformer块的变体，这些变体通过自适应层归一化、交叉留意和额外输入token来加入条件(incorporate conditioning via adaptive layer norm, cross-attention and extra input tokens，这个coditioning相当于就是带条件的去噪)，此中自适应层归一化结果最好。
下文将按照这个架构进行论述，从左到右。
与传统(U-Net)扩散模型区别
架构
[*]DiT将扩散模型中经典的U-Net架构完全替换成了Transformer架构。能够高效地捕捉数据中的依赖关系并生成高质量的结果。
噪声调理战略
[*]DiT扩散过程的采取简单的Linear scheduler（timesteps=1000，beta_start=0.0001，beta_end=0.02）。在传统的U-Net扩散模型（SD）中，所采取的noise scheduler通常是Scaled Linear scheduler。
TODO： 【也有说在传统的U-Net扩散模型（SD）中，所采取的noise scheduler是带调优参数后的线性调理器（Linear Scheduler）。】
与传统扩散的相同
[*]DiT的团体框架并没有采取通例的Pixel Diffusion（像素扩散）架构，而是使用和Stable Diffusion相同的Latent Diffusion（潜变量扩散）架构，使用了和SD一样的VAE模型将像素级图像压缩到低维Latent特征。这极大地降低了扩散模型的计算复杂度（减少Transformer的token的数目）。
输入图像的Patch化（Patchify）和位置编码
在图像领域使用Transformer，起首想到的模型就是ViT（参考：万字长文解读深度学习——ViT、ViLT），和ViT一样，DiT也需要经过Patch和位置编码，如下图红框。
https://i-blog.csdnimg.cn/direct/9c8ea2c0090f4a068b36b51a64e1096a.png
Patch化
DiT和ViT一样，起首采取一个Patch Embedding来将输入图像Patch化，主要作用是将VAE编码后的二维特征转化为一维序列，从而得到一系列的图像tokens，ViT具体如下图所示：
https://i-blog.csdnimg.cn/direct/2e76918932c04cfda9145413c608e929.png
DiT在这个图像Patch化的过程中，设计了patch size超参数，它直接决定了图像tokens的大小和数目，从而影响DiT模型的团体计算量。DiT论文中共设置了三种patch size，分别是 2, 4, 8。patch size 为 2*2 是最理想的。（结论来自：视频生成Sora的全面剖析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等）
Latent Diffusion Transformer结构中，输入的图像在经过VAE编码器处置惩罚后，生成一个Latent特征，Patchify的目的是将Latent特征转换成一系列 T 个 token（将Latent特征进行Patch化），每个 token 的维度为 d。Patchify 创建的 token 数目 T 由补丁大小超参数 p 决定。如下图所示，将 p 减半会使 T 增长四倍，因此至少使整个 transformer Gflops 增长四倍。具体流程如下图所示：
https://i-blog.csdnimg.cn/direct/0240b3cb396c44cabe3f5e6da817c824.png
图4. DiT 的输入规格 。给定patch size是 p × p p \times p p×p，空间表现（来自 VAE 的加噪潜变量），其外形为 I × I × C I \times I \times C I×I×C，会被“划分成补丁”（patchified）为一个长度为 T = ( I / p ) 2 T = (I/p)^2 T=(I/p)2 的序列，埋伏维度为 d d d。较小的补丁大小 p p p 会导致序列长度更长，因此需要更多的计算量（以 Gflops 表现）。
位置编码
在执行 patchify 之后，我们对全部输入 token 应用标准的 ViT 基于频率的位置嵌入（正弦-余弦版本）。图像tokens后，还要加上Positional Embeddings进行位置标记，DiT中采取经典的非学习sin&cosine位置编码技术。
ViT（vision transformer）采取的是2D Frequency Embeddings（两个1D Frequency Embeddings进行concat操作），详情请参考：深度学习——3种常见的Transformer位置编码【sin/cos、基于频率的二维位置编码（2D Frequency Embeddings）、RoPE】
DiT Block模块详细信息
DiT在完成输入图像的预处置惩罚后，就要将Latent特征输入到DiT Block中进行特征的提取了，与ViT不同的是，DiT作为扩散模型还需要在Backbone（主干）网络中嵌入额外的条件信息（不同模态的条件信息等），这里的条件信息就包罗了Timesteps以及种别标签（文本信息）。
DiT中的Backbone网络进行了两个主要工作：
[*]通例的特征提取抽象；
[*]对图像和特征额外的多模态条件特征进行融合。
额外信息都可以采取一个Embedding来进行编码，从而注入DiT中。DiT论文中为了加强特征融合的性能，一共设计了四种 【三种】方案来实现两个额外Embeddings的嵌入(说白了，就是怎么加入conditioning)。实现方式如下图模型架构的后半部分（红框）：
【有的文章表现设计了四种，其实是将AdaLN和AdaLN-Zero分为两种，这里按照论文中图进行解释，分为三种】
https://i-blog.csdnimg.cn/direct/e7059feb714d43868e62fa5f02c2be21.png
Diffusion Transformer模型架构图中由右到左的次序分别是：
[*] 上下文条件（In-context conditioning）
[*] 交叉留意力块（Cross-Attention）
[*] 自适应层归一化块（Adaptive Layer Normalization, AdaLN）
下面将按次序详细先容。
上下文条件化
实现机制：
[*]在上下文条件化中，条件信息 (t)（时间步嵌入） 和 (c)（其他条件，如种别或文本嵌入） 被表现为两个独立的嵌入向量。
[*]这些向量被附加到输入图像 token 序列的开头，形成一个扩展后的输入序列。
[*]Transformer 模块对 (t) 和 (c) 的嵌入与图像 tokens 一视同仁，这些条件化 tokens 通过多头自留意力机制与图像 token 一起参与信息交换。
这与 ViT 中的 cls tokens 类似，它允许我们无需修改就使用标准的 ViT 模块。在最后一个模块之后，我们从序列中移除条件化 tokens。这种方法对模型的新 Gflops 增长可以忽略不计。
交叉留意力模块
实现机制：
[*]条件信息 (t)（时间步嵌入） 和 (c)（其他条件，如种别或文本嵌入） 被拼接（concat）为一个长度为 2 的序列，与图像 token 序列分开。
[*]Transformer 模块被修改为在多头自留意力模块后添加一个 多头交叉留意力层，专门用于让图像 token 与条件 token 进行交互，从而将条件信息显式注入到图像特征中。
[*]图像特征作为Cross Attention机制的查询（Query）。
[*]条件信息的Embeddings作为Cross Attention机制的键（Key）和值（Value）。
这种方式是Stable Diffusion等文生图大模型常用的方式，交叉留意力对模型的 Gflops 增长最多，约莫增长了 15% 的开销。
adaLN-Zero 模块
起首需要了解什么是Adaptive Layer Normalization（AdaLN），而AdaLN之前又要先知道LN，下面将一步步优化解说：
Layer Normalization（LN）
起首在理解AdaLN之前，我们先简单回首一下Layer Normalization。
其他归一化方法参考：深度学习——优化算法、激活函数、归一化、正则化
[*]层归一化（Layer Normalization, LN） 是 Transformer 中的一个关键组件，其作用是对输入的每个特征维度归一化，从而稳定训练和加速收敛。
Layer Normalization的处置惩罚步骤主要分成下面的三步：
[*] 计算输入权重的均值和标准差：计算模型每一层输入权重的均值和标准差。
[*] 对输入权重进行标准化：使用计算得到的均值和标准差将输入权重标准化，使其均值为0，标准差为1。
[*] 对输入权重进行缩放和偏移：使用可学习的缩放参数和偏移参数，对标准化后的输入权重进行线性变换，使模型能够拟合任意的分布。
LN 的公式：
LN ( x ) = γ ⋅ x − μ σ + β \text{LN}(x) = \gamma \cdot \frac{x - \mu}{\sigma} + \beta LN(x)=γ⋅σx−μ+β
此中：
[*] x x x：输入特征。
[*] μ , σ \mu, \sigma μ,σ：输入的均值和标准差（按特征维度计算）。
[*] γ , β \gamma, \beta γ,β：可学习的缩放和偏移参数，用于调解归一化后的分布。
在条件生成任务（如扩散模型）中，需要让条件信息（如时间步 (t) 或种别 (c)）对生成过程产生影响。为此，传统的 LN 被改进为 自适应层归一化（adaLN），可以动态调解归一化的参数以包含条件信息。
Adaptive Layer Normalization（AdaLN）
在 GANs 和具有 UNet 骨干的扩散模型中广泛使用自适应归一化层之后，探索用自适应层归一化（adaLN）替换 transformer 模块中的标准归一化层。adaLN 并不是直接学习维度规模的缩放和偏移参数 γ 和 β，而是从 t 和 c 的嵌入向量之和中回归得到它们。
AdaLN的核心思想是根据输入的不同条件信息，自适应地调解Layer Normalization的缩放参数 γ \gamma γ 和偏移参数 β \beta β，增长的 Gflops 非常少，得当大规模任务。
adaLN 相比LN的改进：
[*] 不再直接使用固定的可学习参数 γ , β \gamma, \beta γ,β。
[*] 相反，AdaLN 仅通过动态生成缩放参数 γ ada \gamma_{\text{ada}} γada 和偏移参数 β ada \beta_{\text{ada}} βada，以条件信息 c c c 为输入，用于调解 Layer Normalization 的行为：
AdaLN ( x , c ) = γ ada ⋅ x − μ σ + ϵ + β ada \text{AdaLN}(x, c) = \gamma_{\text{ada}} \cdot \frac{x - \mu}{\sigma + \epsilon} + \beta_{\text{ada}} AdaLN(x,c)=γada⋅σ+ϵx−μ+βada
[*] γ ada = f γ ( c ) \gamma_{\text{ada}} = f_\gamma(c) γada=fγ(c)：由条件信息 c c c 通过神经网络（如 MLP）生成的缩放参数。
[*] β ada = f β ( c ) \beta_{\text{ada}} = f_\beta(c) βada=fβ(c)：由条件信息 c c c 通过神经网络生成的偏移参数。
阐明：此中 f ( c ) f(c) f(c) 是一个小型神经网络（通常是一个多层感知机，MLP），以条件信息（如时间步 t t t 和种别 c c c 的嵌入向量）为输入，输出对应的 γ , β \gamma, \beta γ,β。
AdaLN的核心步骤
AdaLN的核心步骤包罗以下三步【详细的步骤会在下一节 adaLN-Zero的核心步骤总结】：
[*] 提取条件信息
从输入的条件（如Text Embeddings、标签等）中提取信息，一般来说会专门使用一个神经网络模块（比如全连接层等）来处置惩罚输入条件，并生成与输入数据相对应的缩放参数 γ \gamma γ 和偏移参数 β \beta β。
[*] 生成自适应的缩放和偏移参数
利用提取的条件信息，生成自适应的缩放和偏移参数。假设输入条件为 c c c，经过一个神经网络模块（比如全连接层等）生成缩放参数和偏移参数如下：
γ ada = f γ ( c ) , β ada = f β ( c ) \gamma_{\text{ada}} = f_{\gamma}(c), \quad \beta_{\text{ada}} = f_{\beta}(c) γada=fγ(c),βada=fβ(c)
[*] 使用自适应参数
使用这些自适应参数对输入权重进行 Layer Normalization 处置惩罚：
AdaLN ( x , c ) = γ ada ⋅ x − μ σ + ϵ + β ada \text{AdaLN}(x, c) = \gamma_{\text{ada}} \cdot \frac{x - \mu}{\sigma + \epsilon} + \beta_{\text{ada}} AdaLN(x,c)=γada⋅σ+ϵx−μ+βada
在我们探索的三种模块设计中，adaLN 增长的 Gflops 最少，因此是最计算高效的。它也是唯一一个限制对全部 tokens 应用相同函数的条件化机制。
adaLN-Zero
AdaLN-Zero 在 AdaLN 的基础上新增了 残差缩放参数 α \alpha α，用于动态控制残差路径的影响。通过将 α \alpha α 初始化为零，模型的初始状态被设置为恒等函数，从而确保输出在训练初期的稳定性。这种设计显著提升了模型的训练稳定性和收敛速度。
之前的 ResNets 工作发现，将每个残差块初始化为恒等函数是有益的。例如，在监督学习环境中，将每个块中最后的批量归一化缩放因子 γ 零初始化可以加速大规模训练。
扩散 U-Net 模型使用了类似的初始化战略，在任何残差连接之前零初始化每个块中的终极卷积层。
我们探索了对 adaLN DiT 模块的修改，它做了同样的事变。除了回归 γ 和 β，我们还回归了在 DiT 模块内的任何残差连接之前作用的 dimension-wise 的缩放参数 α。初始化 MLP 以输出全部 α 为零向量；这将完整的 DiT 模块初始化为恒等函数。与标准的 adaLN 模块一样，adaLNZero 对模型的 Gflops 增长可以忽略不计。
adaLN-Zero的核心步骤
下面将根据下图来论述：
https://i-blog.csdnimg.cn/direct/72bcfb2b88e149e2b2f872c091bdccdd.png
[*]AdaLN 有 4 个参数： γ 1 , β 1 , γ 2 , β 2 \gamma_1, \beta_1, \gamma_2, \beta_2 γ1,β1,γ2,β2，分别用于自留意力和 MLP 模块的归一化操作，没有残差缩放参数。
[*]AdaLN-Zero 增长了 2 个参数： α 1 , α 2 \alpha_1, \alpha_2 α1,α2，用于控制残差路径的输出，显著提升了训练稳定性和适应性，因此总共 6 个参数【如上图】。
[*]如果任务需要更强的稳定性（如深层 Transformer 模型或大规模扩散模型训练），AdaLN-Zero 是更优的选择。
adaLN-Zero的核心步骤包罗以下三步，和adaLN的步骤相似，只不外需要在过程中加入维度缩放参数 α \alpha α。
需要做的额外工作如下：
[*]在第一步中提取回归缩放参数 α \alpha α
[*]在第二步中生成自适应的缩放参数 α \alpha α
[*]在第三步中使用 α \alpha α残差路径进行控制
[*] 提取条件信息、缩放参数 α \alpha α：从输入的条件（如Text Embeddings、标签等）中提取信息，一般来说会专门使用一个神经网络模块（比如全连接层等）来处置惩罚输入条件，并生成与输入数据相对应的缩放和偏移参数。
在DiT的官方实现中，使用了一个全连接层+SiLU激活函数来实现这样一个输入条件的特征提取网络：
# 输入条件的特征提取网络
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
nn.Linear(hidden_size, 6 * hidden_size, bias=True)
)
# c代表输入的条件信息
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
同时，DiT 在每个残差模块之后还使用了一个回归缩放参数 α \alpha α 来对权重进行缩放调解，这个 α \alpha α 参数也是由上述条件特征提取网络提取的。
上面的代码示例中和上图（DiT Block with adaLN-Zero）我们可以看到，adaLN_modulation 计算了 6 个变量 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp，这 6 个变量分别对应了多头自留意力机制(MSA)的 AdaLN 的归一化参数与缩放参数（下图中的 β 1 \beta_1 β1, γ 1 \gamma_1 γ1, α 1 \alpha_1 α1）以及 MLP 模块的 AdaLN 的归一化参数与缩放参数（下图中的 β 2 \beta_2 β2, γ 2 \gamma_2 γ2, α 2 \alpha_2 α2）。
【在 DiT Block 中，MSA（多头自留意力模块） 和 MLP（多层感知机模块） 都需要分别进行一次 adaLN-Zero 归一化处置惩罚。每个模块的 Layer Normalization（LN） 都会被替换为 adaLN-Zero，并且两者的 归一化参数 ( γ , β (\gamma, \beta (γ,β） 和 残差路径缩放参数 ( α ) (\alpha) (α) 是独立的（需要分别提取），具体如下：】
Transformer 模块由 MSA 和 MLP 两部分构成，而它们在功能上的分工导致必须对每部分单独设计对应的条件化机制。这是因为：
[*]MSA 的核心在于捕捉全局依赖关系，因此其动态参数 ( β 1 , γ 1 , α 1 ) (\beta_1, \gamma_1, \alpha_1) (β1,γ1,α1)主要控制 全局特征的动态调解。
[*]MLP 的核心在于非线性特征提取，加强局部特征表达，因此其动态参数 ( β 2 , γ 2 , α 2 ) (\beta_2, \gamma_2, \alpha_2) (β2,γ2,α2)主要用于控制 局部特征的动态变换。
[*] 生成自适应的缩放和偏移参数、缩放参数 α \alpha α：
利用提取的条件信息，生成自适应的缩放和偏移参数。假设输入条件为 c c c，经过一个神经网络模块（比如全连接层等）生成缩放参数和偏移参数如下：
γ ada = f γ ( c ) , β ada = f β ( c ) \gamma_{\text{ada}} = f_{\gamma}(c), \quad \beta_{\text{ada}} = f_{\beta}(c) γada=fγ(c),βada=fβ(c)
为残差路径增长一个新的维度缩放参数 α \alpha α，由条件信息动态生成：
α ada = f α ( c ) \alpha_{\text{ada}} = f_{\alpha}(c) αada=fα(c)
初始化为零：在训练开始时， α ada = 0 \alpha_{\text{ada}} = 0 αada=0，使得模块输出仅为主路径输出，实现恒等初始化。
[*] 使用自适应参数、缩放参数 α \alpha α：
[*]在 AdaLN 的基础上，加入 α ada \alpha_{\text{ada}} αada 对残差路径进行缩放控制：
AdaLN-Zero ( x , c ) = α ada ⋅ ( γ ada ⋅ x − μ σ + ϵ + β ada ) \text{AdaLN-Zero}(x, c) = \alpha_{\text{ada}} \cdot \left(\gamma_{\text{ada}} \cdot \frac{x - \mu}{\sigma + \epsilon} + \beta_{\text{ada}}\right) AdaLN-Zero(x,c)=αada⋅(γada⋅σ+ϵx−μ+βada)
[*]残差路径的输出被动态调节：通过 α ada \alpha_{\text{ada}} αada 的渐渐增长，残差路径的影响逐渐加强。
[*]当 α ada = 0 \alpha_{\text{ada}} = 0 αada=0 时，整个模块行为等效于恒等函数。
初始化 MLP 以输出全部 α 为零向量；这将完整的 DiT 模块初始化为恒等函数。adaLNZero 对模型的 Gflops 增长可以忽略不计，与标准的 adaLN 模块一样。
阐明
[*]公式 AdaLN-Zero ( x , c ) \text{AdaLN-Zero}(x, c) AdaLN-Zero(x,c)描述的是残差路径的动态调解过程，输出为 Residual Path Output。
[*]完整的模块输出是路径输出与残差路径输出的加权和（AdaLN完整的模块输出换成对应公式即可）：
Output = Main Path Output + AdaLN-Zero ( x , c ) \text{Output} = \text{Main Path Output} + \text{AdaLN-Zero}(x, c) Output=Main Path Output+AdaLN-Zero(x,c)
[*]这种联系确保了主路径与残差路径的协同作用，结合条件化调解和归一化机制，使模型更加稳定高效地处置惩罚生成任务。
DiT中具体的初始化
设置如下所示：
[*]对DiT Block中的AdaLN和Linear层均采取参数0初始化。
[*]对于其它网络层参数，使用正态分布初始化和xavier初始化。
https://i-blog.csdnimg.cn/direct/0c867a317e344f2ebfa06582fb98f5ce.png
图5.比较不同的条件反射战略。adaLN-Zero在训练的各个阶段都优于交叉留意和情境条件反射。
DiT论文中对四种方案进行了对比试验，发现采取AdaLN-Zero结果是最好的，所以DiT默认都采取这种方式来嵌入条件Embeddings。与此同时，AdaLN-Zero也成为了基于DiT架构的AI绘画大模型的必备战略。
U-ViT（U-Net Vision Transformer）
参考：此文U-ViT部分：视频生成Sora的全面剖析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等
DiT 和 U-ViT 的对比
特性DiTU-ViT模型设计灵感基于 ViT 的纯 Transformer 架构结合 U-Net 和 ViT 的混淆架构网络结构标准 Transformer 堆叠Encoder-Transformer-Decoder 框架局部特征建模依赖 Patch Embedding 和 MLP，局部建模较弱使用 U-Net 的卷积模块，局部特征建模强全局特征建模完全由 Transformer 捕捉全局上下文信息通过嵌入 ViT 加强全局建模能力跳跃连接（Skip）无跳跃连接具有跳跃连接，保存细粒度信息输入表现Patch Embedding 序列化输入原始图像直接输入实用任务高分辨率潜在空间生成任务低分辨率生成任务计算复杂度随序列长度增长计算复杂度显著提升U-Net 局部操作高效，团体复杂度较低 https://i-blog.csdnimg.cn/direct/5bd7297416e5422f88ca905fe288abfb.png
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。
页:
[1]