小yo 博客

Back

《Scalable Diffusion Models with Transformers》论文精读

核心前置理论

隐空间扩散模型(LDM)

原始的 DDPMDDPM是直接在像素级别上加噪和去噪的。对于高分辨率图像,这种方式的计算成本非常高。隐空间扩散模型(LDM)通过引入两阶段的方法解决了这个问题。

  • 压缩(Encode):首先训练一个预训练的自编码器(通常是 VAE),利用其编码器 EE 将高维的像素图像压缩到空间维度更小、但通道数更密集的隐空间表示 z=E(x)z = E(x) 。
  • 扩散:接下来的所有扩散过程(加噪、去噪)都在这个隐空间 zz 上进行,而不再是处理原始图片。这样做能在保留极高生成质量的同时,将 Gflops 降低到原来的几分之一。
  • 解码(Decode):在生成时,扩散模型先在隐空间去噪出一个特征变量 zz ,最后通过 VAE 的解码器 x=D(z)x = D(z) 将其还原为像素图像。

Gflops: Gflops 指的是“十亿次浮点运算”,在论文中被用来衡量模型前向传播的计算复杂度。 传统上,人们喜欢用“参数量”来衡量模型的大小,但作者指出这在图像生成领域具有误导性,因为参数量无法体现“图像分辨率”带来的巨大计算压力。

DiT 本质上就是一个工作在隐空间的扩散模型,论文中提到,一张 256×256×3256 \times 256 \times 3 的图像,会先被 VAE 压缩成 32×32×432 \times 32 \times 4 的隐变量。DiT 的网络结构接收的就是这个 32×32×432 \times 32 \times 4 的张量。

无分类器引导(CFG)

条件扩散模型需要接收额外的信息(如类别标签 cc )来控制生成结果。CFG 是一种 在推理阶段大幅提高条件生成图像质量和条件匹配度的核心采样技术。

  • 训练阶段:模型需要同时学习有条件预测 ϵθ(xt,c)\epsilon_\theta(x_t, c) 和无条件预测。具体做法是,在训练时有一定概率(例如随机 dropout 掉条件 cc )用一个可学习的“空嵌入”(null embedding ∅\emptyset)来代替真正的标签。
  • 推理阶段:模型在每一个去噪布,都会分别计算一次有条件输出和一次无条件输出,然后按以下公式进行外推: ϵ^θ(xt,c)=ϵθ(xt,∅)+s⋅(ϵθ(xt,c)−ϵθ(xt,∅))\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \emptyset) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset)) 这里的 s>1s > 1 就是引导尺度。 (ϵθ(xt,c)−ϵθ(xt,∅))(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset)) 代表了“加入分类条件后,噪声预测发生改变的方向”,我们把这个方向放大 ss 倍,就能让生成的图像更强烈地贴合我们给定的标签。 DiT 能够取得惊艳的生成效果,极大地依赖于这种引导技术。

学习方差

在原始的 DDPM 中,作者在实验后发现,把逆向过程的方差 Σθ\Sigma_\theta 固定为常数工程简化:固定方差 ,模型依然可以生成很好的图像。因此最原始的 DDPM 只需要一个预测噪声的网络 ϵθ\epsilon_\theta 。 但是在后续的“improved DDPM”工作中,研究人员发现预测方差可以让模型在推理阶段以更少的步数达到同样的生成质量,所以后来的高级扩散模型普遍会把方差作为可学习的参数。

预测方差

在扩散模型中,逆向过程的真实方差在理论上有一个固定的上限 βt\beta_t 和一个下限 β~t\tilde{\beta}_t(即 1−αˉt−11−αˉtβt\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t} \beta_t)工程简化:固定方差。既然方差被夹在这两个值之间,我们便可以让神经网络去预测一个插值系数。 神经网络的方差预测头实际上输出的是一个与图像维度相同的张量 vv。为了数值计算的稳定性,插值是在对数空间(Log Space)中进行的。最终模型预测的协方差公式为: log⁡Σθ=vlog⁡βt+(1−v)log⁡β~t\log \Sigma_\theta = v \log \beta_t + (1-v) \log \tilde{\beta}_t

损失函数

为了让模型学习到逆向过程的协方差 Σθ\Sigma_\theta ,只用简单的 MSE 损失是不够的,必须去优化包含完整 KL 散度的损失项。 具体的组合和训练策略如下:

  • 噪声预测使用简单损失:网络预测的噪声 ϵθ\epsilon_\theta 依然使用简单的均方误差进行训练: Lsimple(θ)=∥ϵθ(xt)−ϵt∥22\mathcal{L}_{simple}(\theta) = \|\epsilon_\theta(x_t) - \epsilon_t\|_2^2
  • 方差预测使用完整损失:网络的方差预测部分 Σθ\Sigma_\theta 则是使用完整的 L\mathcal{L} 推导标准答案进行训练。
  • 组合与“停止梯度”技巧:实际上,总的损失函数是这两个损失的加权和(通常表示为 Lhybrid=Lsimple+λLvlb,λ=0.001\mathcal{L}_{hybrid} = \mathcal{L}_{simple} + \lambda \mathcal{L}_{vlb},\lambda = 0.001 )。但是这里有一个极其关键的工程技巧:在计算完整的 VLB 损失时,会对预测噪声的均值部分施加“停止梯度”操作。
  • 结论:Lsimple\mathcal{L}_{simple} 只负责更新网络中用于预测噪声的权重;而完整的 Lvlb\mathcal{L}_{vlb} 损失传导回来的梯度,只负责更新网络中用于预测方差的权重。

DiT 核心架构与创新

图像的 Token 化

既然要用纯 Transformer ,第一步就是要把原本具有空间维度的特征图变成一维的序列。DiT 完全借鉴了 ViT 的做法。

  • 切分与展平:DiT 接收来自 VAE 的空间特征张量 zz(例如 256×256×3256 \times 256 \times 3 的图像,此时 zz 的形状是 32×32×432 \times 32 \times 4) 。第一层被称为 “Patchify”,它通过一个线性层将输入转换为长度为 TT、隐藏层维度为 dd 的 Token 序列 。紧接着,应用标准的基于频率的位置编码位置编码 。
  • 计算复杂度:序列长度 TT 是由 Patch 大小 pp 决定的,公式为 T=(I/p)2T=(I/p)^2 。这是 DiT 调控计算量(Gflops)最关键的旋钮。如果将 pp 减半(例如从 4 变成 2),序列长度 TT 将变为原来的 4 倍,从而导致 Transformer 的总 Gflops 至少增加 4 倍 。

条件信息的特征对齐

无论采用哪种条件注入方案,必须要把时间步 tt 和类别标签 cc 对齐到相同的维度(通常就是 Transformer 的隐藏层维度 dd )。

  • 关于 tt :论文附录中提到,模型先用一个 256 维的频率编码来处理 tt,然后再让它穿过一个带有 SiLU 激活函数的两层 MLP,最终输出的维度刚好等于 Transformer 的隐藏层维度。

SiLU激活函数: 论文中明确提到了使用 SiLU (Sigmoid Linear Unit,也常被称为 Swish) 作为多层感知机的激活函数。

  • SiLU 的公式: f(x)=x⋅σ(x)=x1+e−xf(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}}
  • 与 ReLU的区别: ReLU 的公式是 f(x)=max⁡(0,x)f(x) = \max(0, x)。ReLU 的问题在于它在 x=0x=0 处不可导,并且对于所有负数输入,梯度直接变成了死寂的 0(Dead ReLU 问题)。 相比之下,SiLU 是一个平滑的、非单调的函数。当 xx 是一个较小的负数时,它允许少量的负面信息流过,并且处处可导 。
  • 为什么选择 SiLU: 在现代扩散模型中,SiLU 已经成为了事实上的标配激活函数。在极深的网络中,SiLU 的平滑特性和非零负梯度能让反向传播时的梯度流动更加稳定,极大地帮助了生成模型的收敛。
  • 关于 cc :模型通常会使用一个标准的 Embedding 层来提取类别标签,并将其映射到同样的隐藏层维度。

Embedding 层: 这里的 Embedding 层其实就是一个标准的、可学习的查表操作。对于这篇在 ImageNet 上训练的论文,类别总数是 1000。但为了实现无条件引导,模型需要一个额外的“空标签”。所以,这个 Embedding 层实际上包含 1001 个词条,每一个词条对应一个维度为 dd(Transformer 隐藏层维度)的向量。

条件注入机制

作者极其严谨地测试了四种微调标准 ViT Block 的方案。

上下文条件

直接把 tt 和 cc 的 Embedding 向量当作两个额外的 Token,追加到图像 Token 序列的末尾,就像 ViT 里的 cls token 一样 。这种方法几乎不增加 Gflops 。

交叉注意力

交叉注意力机制
  • 自注意力: 相当于“自己看自己”。特征序列自己生成 Query (Q)、Key (K) 和 Value (V)。图像的每个 Patch 都在寻找与其他 Patch 的关联。
  • 交叉注意力: 相当于“带着问题去找答案”。它的核心在于 Q 和 K、V 来源于不同的输入。 在条件生成的场景下(比如给定类别生成图片):
  • Query (Q):来自于正在生成的图像特征。
  • Key (K) 和 Value (V): 来自于条件信息(比如这篇论文里的时间步 tt 和类别标签 cc)。 图像特征(Q)会去跟条件特征(K)计算相似度打分,然后把对应的条件信息(V)提取过来,融合到自己的特征中。
DiT 中的使用

作者把 tt 和 cc 拼成一个长度为 2 的独立序列,然后在 Transformer 的自注意力层之后,额外加一层多头交叉注意力层 。这种方法增加了约 15% 的 Gflops 额外开销,是所有方案中最昂贵的。

自适应层归一化(adaLN)

自适应归一化

要想理解“自适应”,我们得先回顾一下普通的归一化(比如 Layer Norm)。 普通的 Layer Norm: 在标准的 Transformer 中,Layer Norm 会把输入特征 hh 处理成均值为 0、方差为 1 的状态,然后再乘以一个缩放参数 γ\gamma ,加上一个平移参数 β\beta 。 公式如下: Output=γ⋅h−μσ+βOutput = \gamma \cdot \frac{h - \mu}{\sigma} + \beta

在普通的网络中,这里的 γ\gamma 和 β\beta 是模型的可学习参数。也就是说,无论输入什么图片,这两个参数的值在推理阶段都是固定的。 自适应归一化 (Adaptive Norm / adaLN): 自适应归一化的巧妙之处在于:γ\gamma 和 β\beta 不再是固定的参数了,而是根据输入的“条件”实时计算出来的。

DiT 中的使用

在 DiT 论文中,这里的“条件”就是时间步 tt 和类别标签 cc 的嵌入向量之和。 因为前面一步我们已经把 tt 和 cc 映射成了两个长度同为 dd 的一维向量,所以可以直接把它们加在一起:Condition=Vectort+VectorcCondition = Vector_t + Vector_c 。 得到相加后的条件向量后,adaLN 层会将这个向量送入一个 SiLU 激活函数,紧接着经过一个线性层 。这个线性层会直接输出当前 Transformer 块所需要的缩放参数 γ\gamma 和平移参数 β\beta。 这种方法计算效率最高,增加的 Gflops 最少 。

adaLN-Zero

这是 DiT 的灵魂设计。除了 γ\gamma 和 β\beta,它还回归出一个维度级别的缩放参数 α\alpha,在任何残差连接之前应用 。同时,它将输出这些参数的 MLP 初始化为输出零向量。这导致在训练最开始时,整个 DiT Block 实际上是一个“恒等映射” 。这种初始化极大地提升了大规模网络的训练稳定性,以极小的 Gflops 代价换取了最好的生成效果 。

缩放参数 α\alpha

在标准的 Transformer 块中,残差连接的公式通常是:Output=Input+SubBlock(Input)Output = Input + SubBlock(Input) (这里的 SubBlock 可能是自注意力层或前馈神经网络层)。 论文提到,α\alpha 是一个维度级别的缩放参数,它被应用在DiT 块内任何残差连接之前 。 结合 adaLN-Zero,完整的公式如下:

  1. 先生成参数:通过前面的 MLP 回归出 γ\gamma、β\beta 和 α\alpha。
  2. 应用自适应归一化: h=γ⋅LayerNorm(x)+βh = \gamma \cdot \text{LayerNorm}(x) + \beta
  3. 应用 α\alpha 进行缩放并加入残差: xnew=x+α⋅SubBlock(h)x_{new} = x + \alpha \cdot \text{SubBlock}(h) 这里的 α\alpha 控制了网络当前这个子模块(SubBlock)的输出能在多大程度上加回到原始输入 xx 上。
初始化为零

这是深度学习中一个极其经典且优雅的调参技巧。 论文中明确写道:“我们将 MLP 初始化为对所有的 α\alpha 输出零向量;这会将整个 DiT 块初始化为恒等映射(identity function)” 。 原理解析: 结合我们上面的公式: xnew=x+α⋅SubBlock(h)x_{new} = x + \alpha \cdot \text{SubBlock}(h) 如果在训练的最开始,我们将生成 α\alpha 的神经网络权重初始化为 0,那么初始状态下 α\alpha 就会是一个零向量。 此时公式变成了: xnew=x+0⋅SubBlock(h)=xx_{new} = x + 0 \cdot \text{SubBlock}(h) = x 这意味着,在训练的第 0 步,无论这个 Transformer 块内部有多复杂,它对数据的处理效果就相当于什么都没做,输入什么就输出什么。 为什么这很有用? 论文提到,这是借鉴了 ResNet 的经验 。 在极深的网络中,如果一开始权重是随机的,信号经过几十层复杂的变换很容易彻底丢失(梯度消失)或爆炸。 如果一开始所有的层都是“恒等映射”,信号就能完美、无损地直接从网络第一层传递到最后一层。这极大地保证了网络在训练初期的稳定性,使得模型可以从一个极其稳定的状态开始,一点点“缓慢而安全”地学习如何改变特征 。

Transformer 解码器与输出

经过 NN 层 DiT Block 处理后,我们需要将这些一维的 Token 变回特征图,并输出预测的噪声和方差。 DiT 使用一个标准的线性解码器,对每个 Token 应用最终的归一化,并将其线性解码为一个形状为 p×p×2Cp \times p \times 2C 的张量 。 这里的 2C2C 个通道:一半通道用来输出预测的噪声 ϵθ\epsilon_\theta,另一半用来输出对数方差的插值系数 vv 预测方差。 最后,将这些解码后的 Token 重新排列回原始的空间布局 。通过在一开始注入的位置编码,Transformer 记住了每个 Token 的绝对二维坐标,因此最后可以直接按序重组回原始的网格形状而不乱套。


实验解析与 Scaling Law

实验设置

训练设置

  • 数据集与分辨率:模型在 ImageNet 数据集上训练,分别测试了 256×256256 \times 256 和 512×512512 \times 512 两种高分辨率图像生成。
  • 优化器与超参数:使用了深度学习中最常用的 AdamW 优化器。学习率固定在 1×10−41 \times 10^{-4} ,并且没有使用权重衰退,批量大小设定为 256 。数据增强仅仅使用水平翻转。
  • 摒弃 Transformer 的“传统依赖”:以往在训练 ViT 时,通常需要精心设计学习率预热和各种正则化手段来防止模型崩溃。但是作者发现,DiT 的训练极度稳定,完全不需要这些技巧,也没有出现常见的损失值剧烈波动。
  • EMA 模型评估:这也是生成模型训练的标准操作EMA。作者在训练过程中维护了模型权重的指数移动平均(EMA),衰减率设定为 0.9999。论文中汇报的所有结果,都是基于这个更加平滑、稳定的 EMA 模型得到的。

扩散与潜空间设置

  • 直接用 VAE:为了节省算力并专注于扩散部分的研究,作者没有自己训练 VAE,而是直接使用了 Stable Diffusion 中现成的、经过预训练的 VAE 模型。这个 VAE 的下采样率是 8,完美地将 256×256×3256 \times 256 \times 3 的图像压缩成了 32×32×432 \times 32 \times 4 的隐变量空间。
  • 继承前人的参数:在扩散过程的超参数上,作者几乎完全照搬了 ADM 模型的设置,比如 tmax=1000t_{max} = 1000 的线性方差调度(从 1×10−41 \times 10^{-4} 到 2×10−22 \times 10^{-2})前向过程参数调度,以及 ADM 处理协方差参数化的方法预测方差 。这种严格的控制变量做法,能说明性能的跃升完全来自于 DiT 架构本身,而不是调参的功劳。

评估指标

在生成模型的评测上,作者选用了多种指标。

FID

这是评估图像生成质量的黄金标准。也是该论文中的核心指标。简单来说,它通过预训练的图像分类网络提取真实图像和生成图像的特征,然后计算这两个特征分布之间的距离。FID 数值越小,说明生成的图像越逼真、分布越接近真实世界。FID 论文中主要汇报的是 FID-50K,意思是模型生成了 50000 张图像来和验证集进行对比,采样步数设定为标准 250 步 DDPM 采样。

为什么采样步数是 250 步而不是 1000 步?

  • 固定方差的“短视”: 原始的 DDPM 必须使用 1000 步,因为它把方差固定为了常数。如果在采样时强行减少步数,这个固定的方差假设就会奔溃,导致生成质量断崖式下跌。 原始 DDPM 假设走 1000 步,每一步的跨度极小。在数学推导中,单步逆向过程 q(xt−1∣xt,x0)q(x_{t-1}|x_t, x_0) 的真实方差 β~t\tilde{\beta}_t 是极其微小的。模型把方差固定为这个极小值是没问题的,因为 1000 步的密集采样足以覆盖整个分布。 但是,如果想加速采样,比如从 1000 步直接跨到 990 步(一次跨了 10 步的距离),此时真实的“跨步方差”在数学上会急剧变大。如果依然使用原始 DDPM 那个基于“单步微小位移”推导出的固定极小方差,模型就会因为方差过小、缺乏足够的随机性(不确定性范围)而无法覆盖跨步带来的误差。
  • 学习方差的“自适应”能力: 当模型通过预测插值系数 vv 来学习方差 Σθ\Sigma_\theta 时,它实际上是在学习当前去噪步的不确定性。在 Improved DDPM 和 DiT 中,网络具备了输出较大方差的能力。当强行跳步采样时,网络可以通过预测更大的 vv 值来输出一个更宽广的方差分布。这个变大的方差可以很好地吸收减少步数带来的离散化截断误差,从而赋予了模型“跨大步”依然能平稳落地的能力。
  • 行业标准: 也是为了在算力开销和评测准确度之间取得平衡,250 步已经成为了使用学习方差的扩散模型计算 FID 跑分的标准操作 。

特别需要注意的是,在进行基础的Scaling Law 分析时,为了反映模型最真实的性能,作者在这里汇报的 FID 跑分均没有使用 CFG 引导技术,除非在后续对比最强模型时特别声明。

辅助指标

在评估生成模型的时候,FID 不是万能的,所以作者补充了一些次级指标。

Inception Score (IS)

IS 是早期 GAN 时代非常流行的指标,它看重生成图像的清晰度和多样性。IS 越高越好。IS

sFID

常规 FID 在计算时需要把图片扔进 InceptionV3 分类网络中,提取最后一层全局池化后的特征向量。但这会丢失大量的空间分布和纹理细节。 sFID 取的是池化之前的特征图,将其空间维度展平,作为局部特征集合来计算 Fréchet 距离。它的核心公式与 FID 完全一致,只是统计的对象变了:

sFID=∥μr−μg∥22+Tr(Σr+Σg−2(ΣrΣg)1/2)sFID = \|\mu_r - \mu_g\|_2^2 + Tr(\Sigma_r + \Sigma_g - 2(\Sigma_r \Sigma_g)^{1/2})

  • μr,Σr\mu_r, \Sigma_r:真实图像空间特征的均值向量和协方差矩阵。
  • μg,Σg\mu_g, \Sigma_g:生成图像空间特征的均值向量和协方差矩阵。
    它对图像的底层空间结构、边缘和纹理更加敏感。
Precision & Recall

这一对指标将生成模型的评估拆分成了两个独立的维度:

  • Precision: 衡量的是“模型生成的图片中,有多少落在了真实数据分布的范围内”。Precision 越高,说明模型生成的废图越少,肉眼看着越真实。
  • Recall:衡量的是“真实世界的数据分布中,有多少被模型生成的图片覆盖到了”。Recall 越高,说明模型没有发生“模式崩溃”,能生成各种各样不同的猫、狗、风景,而不是只会画同一种脸。论文特别指出 DiT-XL/2 在这个指标上超越了之前的 LDM 模型 。 这两个指标不假设数据是高斯分布,而是利用k-近邻 (k-NN) 在高维特征空间中构建真实图像和生成图像的“流形边界”。 首先,对于特征空间中的任何一个真实样本 xrx_r,我们找到它的第 kk 个最近邻的真实样本,它们之间的距离设为半径 。所有真实样本周围这个半径内的空间,就构成了真实图像的流形空间 Φr\Phi_r。同理可以构建生成图像的流形空间 Φg\Phi_g。
  • Precision 公式: Precision=1Ng∑j=1NgI(xg(j)∈Φr)Precision = \frac{1}{N_g} \sum_{j=1}^{N_g} \mathbb{I}(x_g^{(j)} \in \Phi_r)
    • NgN_g 是生成图像的总数。
    • I\mathbb{I} 是指示函数(落入区间为 1,否则为 0)。
    • 计算生成的图片中,有多少张(比例)落在了真实图片的流形空间内。只要落进去了,就算作一张好图。它不关心是否生成了所有种类的图,只关心生成的图是不是看起来像真的。
  • Recall 公式: Recall=1Nr∑i=1NrI(xr(i)∈Φg)Recall = \frac{1}{N_r} \sum_{i=1}^{N_r} \mathbb{I}(x_r^{(i)} \in \Phi_g)
    • NrN_r 是真实图像的总数。
    • 计算真实的图片中,有多少张(比例)被生成的图片的流形包围了。如果模型发生了模式崩溃,那么真实世界里的图片就会落在生成流形之外,Recall 就会变得很低。

条件注入消融实验

作者在这里固定了模型大小(使用了最大的 DiT-XL/2),分别用四种不同的条件注入机制条件注入机制训练了四个模型,并比较了它们的 FID 和 Gflops。

  • FID:实验结果显示,adaLN-Zero 在训练的各个阶段,FID 得分都远低于交叉注意力和上下文条件。
  • Gflops:adaLN 和 adaLN-Zero 增加的算力最少(均为 118.6 Gflops),In-context 略高(119.4 Gflops),而 Cross-Attention(交叉注意力)的算力开销最大(高达 137.6 Gflops,增加了约 15% 的额外计算)。
  • 零初始化:最令人惊讶的是 adaLN 和 adaLN-Zero 的对比。两者的网络结构和算力开销完全一模一样,唯一的区别就在于 adaLN-Zero 把输出缩放参数的 MLP 权重初始化为了 0。结果证明,仅仅是这一个初始化的改变,极大地加速了模型的收敛,最终的 FID 显著降低。 这个实验最终确定:条件注入机制不仅影响算力,更致命地影响模型生成质量;而良好的初始化是训练大模型的关键。

控制变量实验

控制 Patch Size 不变

假设固定 pp为 8/4/2 其中一个。此时随着模型从 S 变到 B、L、XL(也就是 Transformer 的层数更深、隐藏层更宽),算力逐渐增加。 结果:FID 曲线完美地呈现出阶梯状下降。模型越大,生成质量越好。

控制模型大小不变

假设固定 架构大小不变(例如 DiT-B)。此时将 Patch size 从 8 缩小到 4,再缩小到 2。结合公式 T=(I/p)2T = (I/p)^2,Patch 减半意味着 Token 数量直接翻 4 倍,计算量也随之暴增图像的 Token 化。 结果:同样地,pp 越小,FID 曲线越低,图像质量越好。

核心结论:Gflops 才是硬通货

过去大家习惯用“参数量”来衡量生成模型强不强。但是 DiT 证明了这是错的。比如,当模型大小固定时,把 Patch size 从 8 缩小到 2,模型的总参数量其实根本没有增加(甚至因为位置编码变长而让线性映射层的参数微调,整体几乎不变),但是 Gflops却翻了十几倍。 在图 8 中,作者把这 12 个模型放在一张图里,横坐标是前向传播的 Gflops,纵坐标是最终的 FID-50K。结果发现它们极其完美地拟合在了一条直线上(皮尔逊相关系数高达 -0.93) 核心结论: 不管通过“加深/加宽网络”来增加算力,还是通过“切分更小的 Patch”来增加算力,只要最终的 Gflops 是一样的,它们的 FID 得分就极其相似(比如图中的 DiT-S/2 和 DiT-B/4)。这就是 DiT 揭示的 Scaling Law 核心!

大模型 vs 小模型

算力效率之争

既然算力那么关键,那如果小模型训练时间足够长,消耗同样的算力,能不能追上大模型?

总训练算力

作者为了比较不同模型的“算力性价比”,提出了估算总训练算力 (Total Training Compute) 的公式 : 训练算力=模型前向Gflops×批次大小(BatchSize)×训练步数(TrainingSteps)×3训练算力 = 模型前向 Gflops \times 批次大小 (Batch Size) \times 训练步数 (Training Steps) \times 3

  • 模型前向 Gflops:这是模型处理单张图片进行一次前向传播(也就是从输入到输出算出预测噪声)所需要的浮点运算次数 。
  • 批次大小 ×\times 训练步数:这两项乘起来,代表了模型在整个训练过程中,一共处理了多少张图片 。比如 Batch Size 是 256,训练了 400K 步,就意味着模型一共处理了 256×400,000256 \times 400,000 张图。
  • 乘数因子 3: 神经网络的训练包含前向传播(算误差)和反向传播(算梯度并更新权重)。在深度学习框架的底层实现中,根据经验估算,计算一次反向传播所需的算力大约是前向传播的两倍 。所以,前向 (1) + 反向 (2) = 总消耗是前向传播的 3 倍 。 这个公式将时间、硬件规模和模型规模统筹在了一个统一的维度下,科学地量化了算力的消耗量。
结论

作者发现,大模型(比如 XL/2)在相同的训练算力消耗下,FID 远低于小模型。即“一个小模型训练很久”不仅效果打不过“一个大模型只训练一小会儿”,而且在算力利用率上也是极其低效的。这告诉我们,如果手头有足够的算力,直接上大模型是最经济的选择。

采样算力能否弥补模型算力

扩散模型可以在推理时增加去噪步数来提高图片质量,那么如果用一个小模型在推理时采样 1000 步,能否打败大模型只跑 128 步? 在图 10 中,作者做了极其详尽的对比。比如拿 DiT-L/2 跑 1000 步,和 DiT-XL/2 跑 128 步对比。尽管 L/2 跑 1000 步消耗的“总采样算力(80.7 Tflops)”是 XL/2 跑 128 步(15.2 Tflops)的 5 倍多,但 XL/2 的 FID 依然吊打 L/2。 结论:模型算力是硬伤。如果模型本身没学到足够好的数据分布,再多的采样算力也没用。采样算力的增加,永远无法弥补模型算力的先天不足。

最终成绩:SOTA

在完成 Scaling Law 的验证后,作者将算力最大的模型 DiT-XL/2 继续训练到了 700 万步,并开启了 CFG。

  • 吊打 U-Net 架构:在 ImageNet 256×256256 \times 256 的条件生成任务上,DiT-XL/2 取得了 2.27 的 FID 得分,远低于之前基于 U-Net 架构的最强模型 LDM(FID 3.60)和 ADM。
  • 算力性价比极高:相比于在像素空间计算、极其昂贵的 ADM(1120 Gflops),DiT-XL/2 仅用了 118.6 Gflops 就取得了更好的成绩。
  • 历史意义:这个结果无可辩驳地证明了:在扩散模型领域,U-Net 的归纳偏置并非不可替代。纯粹的 Transformer 架构不仅能做扩散模型,而且能在规模化后做得比 U-Net 更好。这也为后来 Sora 等视频生成大模型全面拥抱 DiT 架构奠定了基础。

DiT 的历史地位与局限性

DiT 的历史贡献

  • 大一统的曙光:DiT 证明了在生成模型领域,Transformer 完全可以取代 U-Net。这让计算机视觉和自然语言处理在底层架构上进一步走向了统一。
  • 极强的可预测性:它的 Scaling Law 极其优美。只要你有钱、有算力,你就可以闭着眼睛扩大模型规模(增加 Gflops),且你可以非常确定地知道生成质量一定会提升。这种“大力出奇迹”的稳定性,是工业界(比如后来的 Sora 模型全面采用 DiT)最看重的。

DiT 的局限性

  • 平方级复杂度:Transformer 的 Self-Attention 计算复杂度是关于序列长度 TT 成平方关系 O(T2)O(T^2) 的。即使 DiT 在隐空间工作,但如果要生成 1024×10241024 \times 1024 甚至更高分辨率的图像,或者生成长视频,Token 数量会呈指数爆炸,导致算力开销变得极其恐怖。这也是为什么很多后续工作都在研究 Linear Attention(线性注意力)或 Mamba 等架构来替代标准 Transformer。
  • 并非“纯粹”的 Transformer:DiT 依然重度依赖基于 CNN 训练的 VAE 来进行特征压缩和解码。它并没有完全摆脱卷积的影子。如何训练一个真正端到端(End-to-End)的纯 Transformer 像素级扩散模型,依然是一个挑战。
  • 缺乏空间归纳偏置:相比于 U-Net(卷积网络天生懂得“相邻的像素大概率是相关的”),Transformer 一开始是一张白纸,它必须通过消耗海量的数据和算力,才能自己慢慢“悟”出图像的二维空间结构。所以它在小数据集上极其容易过拟合,是一个真正的“算力巨兽”。
《Scalable Diffusion Models with Transformers》论文精读
https://smallyo.top/blog/papers-DiT
Author小yo
Published atMarch 1, 2026