)
目录前言1. 引言2. 模型设计2.1 条件注入消融实验2.2 AdaLN-Zero3. DiT Block 代码实现结语参考前言学习 UP 主 霹雳吧啦Wz 的 Scalable Diffusion Model with Transformers (DiTs) 视频了解一下 DiT 的相关工作记录下个人学习笔记和大家一起分享交流videoScalable Diffusion Model with Transformers (DiTs)1. 引言本期内容我们简单聊一聊 DiT本次内容分为三个部分引言模型设计DiT Block 代码实现在代码实现部分我们主要聚焦 DiT Block 的具体实现。在之前的内容当中我们简单介绍过 LDM 的相关内容为了减少扩散模型的单次推理计算量我们会先将图像通过 VAE 的 encoder 编码到更小的隐空间—编码时把原始图像的高宽各下采样f ff倍隐空间特征图的像素数量面积就只有原来的1 / 64 1/641/64f 8 f8f8时即为各下采样 8 倍。推理时我们直接在隐空间完成去噪再通过 VAE 的 decoder 解码还原回图像空间。现在我们已经可以通过 LDM 降低去噪迭代的成本了接下来要考虑的就是如何进一步提升生成效果。如果默认还是采用 DDPM 的理论那么最容易想到的就是提升去噪模型本身的能力而本期内容要讲的 DiT聚焦的正是这一点。这篇论文的标题是 Scalable Diffusion Models with Transformers下载链接是 https://arxiv.org/abs/2212.09748论文最早发表于 2022 年底。回想当时大语言模型的发展非常火热。2020 年OpenAI 提出了 Scaling Laws在大语言模型领域用 Transformer 结合 Scaling Laws通过不断增加模型参数、扩大训练数据量、投入更多训练资源就能得到更好的效果所以当时是一个 “大力出奇迹” 的时代。既然大语言模型那边发展得这么好作者就想到能不能把这些能力迁移到图像领域—于是就有了这篇工作。这篇论文把 LDM 中采用 U-Net 架构的扩散模型换成了纯 Transformer 的模型。首先来看图 8横坐标是 Transformer 模型的 Gflops即推理时的理论计算量纵坐标是 FID可以认为是生成效果的指标数值越低生成效果就越好。作者在论文中设计了不同的模型大小和输入分辨率并在相同训练步数下绘制了这幅图。可以看到模型的计算量越大FID 越低两者呈强相关这也体现了 Transformer 强大的扩展能力。接下来再看图 9横坐标是训练计算量纵坐标同样是 FID。先看论文当中框住的这个局部区域如果在这里做一条垂线就代表在训练算力相同的情况下比较不同模型的效果。从图中可以发现模型越大效果越好—所以在训练算力预算充足时更大的 DiT 模型算力利用率更高。但如果看这个区域左侧的部分反而是小模型的效果更优也就是说算力预算比较低的时候用小模型更划算。当然作者想表达的结论是算力足够时使用更大的 DiT 模型效率更高。接着再看这幅图横坐标向右代表增大 Transformer 模型的参数量纵坐标向下对应减小 patch size后者也就是刚才图中的 8、4、2。patch size 越大后续 patchify 时对隐空间高宽的压缩倍率就越大DiT Block 的运算量也就越小。通过上一幅图可以看到模型的参数量越大、patch size 越小或者说输入分辨率越高生成效果就越好。实际上目前的一些主流方案默认都采用 patch size 为 2很少会用到 4 甚至 8。2. 模型设计接着来看 DiT 的架构下图就是论文当中展示的整体架构。首先这里的带噪 latentNoised Latent就是我们通过 VAE encoder 编码得到的隐空间变量。比如训练使用 256×256 分辨率时下采样倍率f 8 f8f8隐空间对应的高宽就是 32×32。接下来会先把带噪 latent 做一次 patchify具体实现如下所示它和我们之前讲 ViT 时的 patchify 逻辑是一样的。以这里的带噪 latent 为例I II对应 32会在此基础上进一步压缩压缩用的 patch 大小是p × p p \times pp×p其中p pp可以取 2、4 或 8。具体实现时可以直接用一个卷积来做当 patch size 为 2 时用一个 kernel 大小为2 × 2 2 \times 22×2、stride 为 2 的卷积进行下采样下采样之后再将所有 token 展平就得到长度为T TT的 token 序列。接着再来看输入的条件。对于类别条件场景比如基于 ImageNet 训练时这里的 labely yy就是我们要生成的类别的索引信息除此之外还需要当前去噪的时间步t tt。labely yy和时间步t tt共同组成了要注入的条件把它注入模型之后再送入 N 个 DiT Block。关于 DiT Block 的具体实现可以看上面架构图里的 DiT Block with adaLN-Zero 部分。先不看右边的部分左边的 Layer Norm、Multi-Head Self-Attention、Pointwise Feedforward 都是非常常见的结构唯一没见过的是右边的 MLP 以及 Scale/Shift。这部分其实是这篇工作中额外提出的 adaLN-Zero 模块细节我们后面再讲。通过 N 个 DiT Block 之后再经过一个 Layer Norm、一个 Linear 以及 reshape就能得到模型预测的最终输出。可以看到在 DiT 这篇工作中模型的输出其实有两部分一个是预测的噪声另一个是协方差Σ \SigmaΣ。在 DDPM 当中协方差是固定的所以只需要预测噪声而 DiT 继承了 ADM 这篇工作让模型去预测协方差推理时就可以用预测出的协方差进行运算。但实际上大部分开源实现还是默认只预测噪声也就是回退到原始的 DDPM 做法。下面来看论文中关于条件注入的消融实验。2.1 条件注入消融实验在这个实验中作者一共做了四组实验多出的一组是 adaLN不带 Zero但下图只绘制了其中三种最右边采用的是 In-Context 条件注入方式它把条件t tt和c cc分别编码成两个 token再与图像 token 在序列维度拼接然后经过标准的 Layer Norm、Multi-Head Self-Attention、Layer Norm 以及 Pointwise Feedforward 这些模块。这种注入方式额外引入的理论计算量比较小因为它在原始 token 序列上只增加了两个 token。除此之外作者还尝试了 Cross-Attention 的条件注入方式中间这幅图它会额外引入大概 15% 的计算开销。但在 DiT 这篇工作中作者最终采用的是 adaLN-Zero 这种条件注入形式它的效果在这几种方式中最好额外引入的 Gflops 也最小—它通过条件生成自适应归一化的参数γ \gammaγ、β \betaβ以及门控参数α \alphaα以此完成条件注入。关于这三种方式的计算量作者在论文当中也有统计以 DiT-XL/2 这个模型为例采用 In-Context 形式时 Gflops 约为 119.4Cross-Attention 形式是 137.6adaLN 是 118.6adaLN-Zero 同样是 118.6。这些方法的效果在论文图 5 中有展示。从这幅图可以看到采用 adaLN-Zero 的效果最好。不过作者这个结论其实是基于类别标签这类任务的UP 提到对于 “在序列维度拼接效果不太理想” 这一点他持怀疑态度。因为目前主流的条件注入方式基本都是这篇工作的 adaLN-Zero 与序列维度拼接条件的结合。至于为什么在这篇工作中序列维度拼接效果不理想他的猜测是时间步t tt和条件c cc各自只对应一个 token这两个 token 在整个序列中占比实在太低做 attention 时注意力很容易被其他 token 分散掉。所以他认为并不是序列维度拼接本身效果不好而是在当前这个任务中条件信息对应的 token 占比太低。2.2 AdaLN-Zero大家可能困惑这里的 adaLN-Zero 到底是怎么来的作者是怎么想到这样一种条件注入形式的呢。在讲这个问题之前我们先来回顾一下 Layer Norm 的计算公式。y x − E [ x ] Var [ x ] ϵ ∗ γ β \boldsymbol{y} \frac{\boldsymbol{x} - \text{E}[\boldsymbol{x}]}{\sqrt{\text{Var}[\boldsymbol{x}] \epsilon}} * \boldsymbol{\gamma} \boldsymbol{\beta}yVar[x]ϵx−E[x]∗γβ对于每一个 token我们会减去均值、除以标准差然后再乘上γ \gammaγ加上β \betaβ。其中γ \gammaγ和β \betaβ是可学习的仿射参数这两个参数一旦训练完成就固定下来也就意味着无论当前的时间步是多少、输入的条件c cc是什么这里的仿射参数都是固定不变的。如果这两个参数能够根据不同的时间步t tt以及条件c cc动态地调整特征分布就好了。所以在 adaLN 中我们把γ \gammaγ和β \betaβ改成了由条件信息实时生成的形式从而实现条件的注入。使用 adaLN 时需要把 Layer Norm 的仿射参数关掉具体怎么关呢在 Layer Norm 的创建过程当中有一个elementwise_affine参数如果将它设置成 False就不会有这两个可学习参数。然后我们再额外引入一个 MLP 结构让这个 MLP 根据输入的条件去生成γ \gammaγ和β \betaβ就可以了。但在 DiT 这篇工作当中我们使用的其实还不是 adaLN而是 adaLN-Zero—两者的区别在于adaLN-Zero 额外引入了一个门控参数α \alphaα或者说 gate也就是图中的α 1 \alpha_1α1、α 2 \alpha_2α2。有了这两个门控参数我们就能让 DiT Block 在初始化时变成一个恒等映射只要把 MLP 的权重weight和偏置bias都初始化为零它输出的γ \gammaγ、β \betaβ以及α \alphaα就全部为零于是 Multi-Head Self-Attention 和 Pointwise Feedforward 的输出都会被屏蔽掉DiT Block 的输出就等于 input tokens随着训练进行这些参数才会逐步发挥作用。这种初始化方式目前也非常常用它能让我们的训练更加稳定、收敛更快也算是训练工程上的一些小 trick 了。在本期内容的最后我们再简单看一下源码里是如何实现这个 DiT Block 的。3. DiT Block 代码实现相关源码可以查看https://github.com/facebookresearch/dit################################################################################## Core DiT Model ##################################################################################classDiTBlock(nn.Module): A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning. def__init__(self,hidden_size,num_heads,mlp_ratio4.0,**block_kwargs):super().__init__()self.norm1nn.LayerNorm(hidden_size,elementwise_affineFalse,eps1e-6)self.attnAttention(hidden_size,num_headsnum_heads,qkv_biasTrue,**block_kwargs)self.norm2nn.LayerNorm(hidden_size,elementwise_affineFalse,eps1e-6)mlp_hidden_dimint(hidden_size*mlp_ratio)approx_gelulambda:nn.GELU(approximatetanh)self.mlpMlp(in_featureshidden_size,hidden_featuresmlp_hidden_dim,act_layerapprox_gelu,drop0)self.adaLN_modulationnn.Sequential(nn.SiLU(),nn.Linear(hidden_size,6*hidden_size,biasTrue))defforward(self,x,c):shift_msa,scale_msa,gate_msa,shift_mlp,scale_mlp,gate_mlpself.adaLN_modulation(c).chunk(6,dim1)xxgate_msa.unsqueeze(1)*self.attn(modulate(self.norm1(x),shift_msa,scale_msa))xxgate_mlp.unsqueeze(1)*self.mlp(modulate(self.norm2(x),shift_mlp,scale_mlp))returnx首先定义了一个名为DiTBlock的类在初始化过程中创建了一个 Layer Normnorm1它对应的就是图中左边的 Layer Norm它的elementwise_affine设置的是 False。attn对应图中的 Attention 模块norm2对应另一个 Layer Normmlp对应 Pointwise Feedforward最后这里的adaLN_modulation就是我们刚刚讲的 adaLN-Zero可以看到它内部核心就是一个 Linear 层输入维度是 hidden size输出是 6 倍的 hidden size。再看前向传播forward这里输入的x就是 input tokensc就是 conditioning。我们先把条件c传入定义的adaLN_modulation模块接着把它分成六份分别对应shift_msa、scale_msa、gate_msa以及shift_mlp、scale_mlp、gate_mlp。其中shift_msa对应β 1 \beta_1β1scale_msa对应γ 1 \gamma_1γ1gate_msa对应图中的α 1 \alpha_1α1后面这三个参数也是类似的就不去讲了。接着我们将输入的x先过一遍 Layer Norm再经过一个 modulate 函数它的定义如下defmodulate(x,shift,scale):returnx*(1scale.unsqueeze(1))shift.unsqueeze(1)可以看到这个函数输入有三个参数x、shift以及scale它的计算公式是x先乘以 (1 scale)再加上shift。注意这里有个小细节并不是直接乘以scale。为什么写成这种形式呢因为刚刚提到过在 adaLN-Zero 当中初始化时我们希望γ \gammaγ、β \betaβ以及α \alphaα全部为零但这些参数是根据条件动态生成的所以只能通过初始化 MLP 的参数来控制γ \gammaγ、β \betaβ以及α \alphaα的初始数值。实际初始化 adaLN-Zero 参数时对应的就是下面这段代码。# Zero-out adaLN modulation layers in DiT blocks:forblockinself.blocks:nn.init.constant_(block.adaLN_modulation[-1].weight,0)nn.init.constant_(block.adaLN_modulation[-1].bias,0)可以看到初始化就是把 adaLN modulation 的最后一个 layer也就是这里的 Linear的 weight 以及 bias 全部置零。置零之后条件c经过这个模块输出的六个参数就全部为零也就意味着整个 DiT Block 的输入和输出完全一样等价于恒等映射scale 和 shift 也不会起作用后续在实际训练中这几个参数才会慢慢起作用。再回到 forward 函数当中。刚刚讲到的modulate处理完之后先通过 attention 模块再乘上 adaLN-Zero 输出的 gate 参数加上残差连接x接着经过 Layer Norm、modulate、Mlp再乘上对应的 gate最后加上残差连接x就得到整个 DiT Block 的输出。讲到这儿相信大家已经对整个 DiT 的结构有了非常细致的了解。到此本期内容就全部结束了希望对大家有帮助。结语本篇文章我们跟随 UP 主梳理了 DiT 的提出动机、模型设计、条件注入的消融实验以及 adaLN-Zero 的具体实现最后过了一遍 DiT Block 的源码。DiT 的出发点很直接LDM 已经通过隐空间压缩把去噪迭代的成本降了下来那么下一步自然就是提升去噪模型本身的能力。当时大语言模型领域正凭借 Transformer Scaling Laws 一路高歌猛进作者便顺势把 LDM 中的 U-Net 换成了纯 Transformer 架构。论文中的图 8 和图 9 也验证了这一点模型的计算量与 FID 呈强相关并且在训练算力预算充足的前提下更大的 DiT 模型算力利用率更高。模型结构上DiT 沿用了 ViT 的思路—先把 VAE 编码得到的带噪 latent 通过 patchify 切成 token 序列这一步可以直接用一个 kernel 大小为2 × 2 2 \times 22×2、stride 为 2 的卷积来实现随后送入 N 个 DiT Block最终输出预测的噪声以及继承自 ADM 的协方差Σ \SigmaΣ实际开源实现中大部分还是只预测噪声。至于 patch size论文显示参数量越大、patch size 越小则效果越好目前主流方案默认取 2。条件注入是这篇工作最核心的部分。作者对比了 In-Context、Cross-Attention 和 adaLN-Zero 三种方式adaLN-Zero 的效果最好额外引入的 Gflops 也最小—它的思路是让条件信息通过一个 MLP 实时生成 Layer Norm 的仿射参数γ \gammaγ、β \betaβ以及门控参数α \alphaα使归一化后的缩放与偏移能够随时间步t tt与条件c cc动态调整而 Zero 的含义则是把 MLP 的 weight 和 bias 初始化为零让 DiT Block 在初始化时等价于一个恒等映射训练因此更加稳定、收敛也更快这也成了后续生成模型当中非常常用的一个初始化 trick。这里也提到 UP 主个人的不同看法作者关于 “序列维度拼接效果不理想” 的结论未必是拼接方式本身的问题更可能是因为时间步和类别各自只对应一个 token在整个序列中占比实在太低做 attention 时注意力容易被稀释掉。而目前的主流方案恰恰是 adaLN-Zero 与序列维度拼接条件的结合这也从侧面印证了这一点。总的来说DiT 把扩散模型的骨架从 U-Net 换成了 Transformer证明了 Transformer 在图像生成领域同样具备良好的可扩展性这也为后续 SD3、Flux 等工作的架构选择打下了基础。感兴趣的同学可以多看看 UP 主的视频非常感谢 UP 主的精彩讲解 。参考https://github.com/facebookresearch/ditScalable Diffusion Models with TransformersScalable Diffusion Model with Transformers (DiTs)
