在这里插入图片描述

1. 论文试图解决什么问题?

打破 U-Net 的架构垄断:在本文发表之前,基于扩散模型的图像生成(如 DDPM、LDM)几乎全部默认使用带有卷积层的 U-Net 作为主干网络 。

验证标准架构的通用性与可扩展性:论文试图证明 U-Net 所带来的空间归纳偏置并不是生成高质量图像所必需的,标准的 Transformer 架构可以被用来替代 U-Net 。

探索 Scaling Law:作者希望验证 Transformer 优秀的“可扩展性”在扩散模型中是否同样适用,即通过增加模型复杂度的计算量(Gflops)能否持续带来图像质量的提升 。

2. 模型输入输出是什么?

这篇论文的 DiT(Diffusion Transformers)是在潜扩散模型(Latent Diffusion Models)框架下工作的,它不直接处理像素,而是处理经过预训练 VAE 压缩后的潜变量空间 。

输入(添加了噪声的潜变量):一个空间特征张量,例如对于 256x256x3 的图像,经过下采样率为 8 的 VAE 编码后,输入尺寸为 32x32x4 。

输入(条件信息):包含了扩散过程的当前时间步(Timestep)以及图像的类别标签(Class label) 。

输出(预测的噪声和协方差):模型最终输出预测的噪声以及对角协方差,输出的张量形状与输入的潜变量完全一致,随后可通过 VAE 的解码器还原成像素图像 。

3. 模型的网络结构是什么样子的?

DiT 的整体架构极大程度地保留了标准 Vision Transformer (ViT) 的设计 。

Patchify(块分割与展平):第一步是将输入的空间潜变量划分为多个小块(Patch)并线性映射为 Token 序列,同时加上标准的 ViT 正弦-余弦位置编码 。较小的 Patch 会产生更长的序列,从而大幅增加总体的计算量(Gflops) 。

DiT Block(条件注入机制):为了注入时间步和类别标签,作者探索了四种 Transformer 块变体:In-context(将条件作为额外 Token 拼接)、Cross-attention(引入交叉注意力层)、Adaptive layer norm(使用自适应层归一化)以及 adaLN-Zero 。

adaLN-Zero(最终核心设计):这是论文最终采用的结构,它通过条件信息回归层归一化的参数以及残差连接前的维度缩放参数 。最关键的是,它将这些缩放参数初始化为零,使得整个 DiT Block 在初始状态下等同于一个恒等映射函数 。

Transformer Decoder:经过一系列 DiT 块后,模型采用一个标准线性解码器将 Token 还原为原始的空间布局结构 。

4. 使用了什么数据集?

  • 模型是在 ImageNet 数据集上进行训练的,主要针对类别条件图像生成任务 。

  • 论文核心评估了 256x256 和 512x512 两种高挑战性的图像分辨率 。

5. 核心 Contribution 是什么?

提出了 DiT 全新架构:成功将纯 Transformer 应用于图像扩散模型,证明了没有特定卷积设计也能做到甚至超越目前的表现 。

验证了扩散模型的“可扩展性”:系统构建了 DiT 的设计空间,并证明增加网络复杂度(Gflops),不管是通过加深/加宽网络,还是增加输入 Token 数量,都能一致地降低 FID 分数(提升图像质量) 。

大幅刷新 SOTA 成绩:其最大规模的模型 DiT-XL/2 超越了之前所有的 U-Net 扩散模型,在类别条件 ImageNet 256x256 生成基准上实现了 2.27 的最先进 FID ;在 512x512 分辨率上达到了 3.04 的 FID 。


Adaln-zero

LN 的参数是写死的(静态的),而 adaLN-Zero 的参数是根据条件动态生成的(动态的),并且 adaLN-Zero 在训练初始阶段会巧妙地将整个网络层变成“透明”的(恒等映射)。

以下是它们在数学直觉和架构上的核心差异对比:

1. LN (Standard Layer Normalization)

标准层归一化是我们在 Transformer 中最常见的组件。

  • 计算逻辑:它将输入特征的均值归零,方差归一,然后乘以一个缩放参数 γ\gammaγ,再加上一个平移参数 β\betaβ
  • 参数特性(静态)γ\gammaγβ\betaβ作为模型权重直接存在网络里的。它们在训练过程中被梯度下降更新,一旦训练完成,无论你输入什么图片或文本,这两个参数都是固定不变的。

2. adaLN (Adaptive Layer Normalization)

“Adaptive(自适应)”是为了解决条件注入的问题(比如在扩散模型中,我们需要告诉模型当前是第几个时间步 ttt,或者生成什么类别 ccc)。

  • 计算逻辑:归一化过程与 LN 一样,但 γ\gammaγβ\betaβ 的来源变了。
  • 参数特性(动态)γ\gammaγβ\betaβ 不再是固定的权重,而是由一个额外的 MLP(多层感知机)根据条件(tttccc)实时计算出来的。这意味着:输入不同的类别或时间步,层归一化的缩放和平移方式就会动态改变。这就把外部条件“注入”到了主干网络中。

3. adaLN-Zero (DiT 的核心创新)

这是在 adaLN 基础上的究极进化版,也是 DiT (Diffusion Transformers) 能够成功训练超大模型的关键魔法。它在 adaLN 的基础上做了两项致命修改:

  • 增加残差缩放参数 α\alphaα:那个用来处理条件信息的 MLP,除了实时计算 γ\gammaγβ\betaβ 之外,还会额外计算出一个维度缩放参数 α\alphaα。这个 α\alphaα 会被乘在 Transformer 块(自注意力机制或 FFN)的输出上,然后再加回残差连接。

  • 公式大致为:Output=Input+α⋅TransformerBlock(adaLN(Input))\text{Output} = \text{Input} + \alpha \cdot \text{TransformerBlock}(\text{adaLN}(\text{Input}))Output=Input+αTransformerBlock(adaLN(Input))

  • “Zero” 初始化(零初始化):负责输出 γ\gammaγβ\betaβα\alphaα 的 MLP 最后一层,被强制初始化为 0

  • 奇妙的化学反应:在训练的第 0 步,由于 MLP 输出全是 0,导致 α=0\alpha = 0α=0。带入上面的公式你会发现:Output=Input+0⋅TransformerBlock(...)\text{Output} = \text{Input} + 0 \cdot \text{TransformerBlock}(...)Output=Input+0TransformerBlock(...)

  • 结果:整个极其复杂的 Transformer 块在初始阶段直接变成了恒等映射(Identity function)。输入什么,就原封不动地输出什么。


总结对比表

特性 LN (Layer Norm) adaLN adaLN-Zero
γ\gammaγβ\betaβ 的来源 固定模型权重(静态) MLP 根据条件动态计算 MLP 根据条件动态计算
残差缩放 α\alphaα 有(MLP 动态计算)
能否处理条件输入 否(需借助额外结构)
初始训练状态 随机/标准分布 随机/标准分布 恒等映射 (Identity function)

为什么 adaLN-Zero 这么重要?

在深度学习中,网络越深,训练越容易崩溃。对于 DiT 这样极深、计算量极大的 Transformer 模型,如果初始状态下每一层的信号都被随机的注意力和前馈网络剧烈扭曲,模型很难收敛。

adaLN-Zero 的绝妙之处在于:它让一个拥有几十层、上百亿参数的深层网络,在训练开始的第一步,表现得就像一个没有深度的“空管道”。随着训练的推进,模型会慢慢学习让 MLP 输出非零的 α,γ,β\alpha, \gamma, \betaα,γ,β,极其平滑且稳定地“唤醒” Transformer 的每一层,这极大提升了超大扩散模型的训练稳定性和收敛质量。

虽然DiT发现adaLN-Zero效果是最好的,但是这种方式只适合这种只有类别信息的简单条件嵌入,因为只需要引入一个class embedding;但是对于文生图来说,其条件往往是序列的text embeddings,采用cross-attention方案可能是更合适的。 由于对输入进行了token化,所以在网络的最后还需要一个decoder来恢复输入的原始维度,DiT采用一个简单的linear层来实现,直接将每个token映射为p×p×2Cp×p×2Cp×p×2C的tensor,然后再进行reshape来得到和原始输入空间维度一样的输出,但是特征维度大小是原来的2倍,分别用来预测噪音和方差。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐