DiT vs UNet：扩散模型架构的理论对比分析

概述

UNet 和 Transformer 是扩散模型去噪网络的两种主要架构选择。尽管 UNet 在扩散模型领域长期占据主导地位，但 DiT (Diffusion Transformer) 的出现挑战了这一范式。

本文件从理论角度深入分析两种架构的差异，解释为什么 Transformer 在扩散模型中表现出色。¹²

---

1. 架构归纳偏置对比

UNet 的归纳偏置

UNet 是一种编码器-解码器架构，具有以下特征：

归纳偏置	描述	优势	劣势
局部性	卷积核感受野有限	高效处理局部模式	长程依赖需多层堆叠
平移不变性	权重共享	参数高效	对位置不敏感
多尺度	编码器下采样 + 解码器上采样	捕获多尺度特征	结构复杂
跳跃连接	保留细粒度信息	保留细节	梯度流复杂

Transformer 的归纳偏置

DiT 基于 Vision Transformer (ViT)，具有不同的特征：

归纳偏置	描述	优势	劣势
全局注意力	所有位置两两交互	自然建模长程依赖	计算量 O(N²)
弱归纳偏置	数据驱动学习	灵活性高	需要更多数据
序列建模	图像视为 token 序列	与 NLP 统一	丢失局部结构
位置编码	显式编码位置	可控的位置建模	需要精心设计

关键理论洞察

直觉：扩散模型需要建模全局一致性（如物体结构、空间关系），而不仅仅是局部纹理。Transformer 的全局注意力天然适合这种需求。

---

2. 表达能力理论分析

理论框架

设去噪网络需要建模的函数为 $f:{\mathbb{R}}^{H\times W\times C}\times \mathbb{R}\to {\mathbb{R}}^{H\times W\times C}$。

UNet 的表达能力

UNet 通过多尺度特征提取来增加感受野：

$$\text{ReceptiveField}(l)=1+2l\cdot k$$

其中 $l$ 是层数，$k$ 是卷积核大小。

问题：即使有 18 层的 UNet，感受野仍然有限。

Transformer 的表达能力

Transformer 通过自注意力实现全局交互：

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$

每个位置可以直接关注所有其他位置，理论上感受野为 $N$（序列长度）。

优势：$O(1)$ 层即可建模全局依赖。

数学对比

方面	UNet	DiT
单层感受野	$k\times k$	$N\times N$ (全局)
深度增加效果	线性增长	指数增长（注意力层叠）
最优深度	需要很多层	相对较少
参数效率	局部模式复用	全局模式学习

---

3. 参数效率分析

相同参数下的表达能力

假设总参数量为 $P$：

架构	配置	有效感受野	全局建模能力
UNet	$P$ 个 3×3 卷积	$\sim 2P/3$	需 10+ 层
DiT	$P$ 个注意力头	$\sim P/4$ tokens	天然全局

参数量分解

UNet 参数量

$$P_{\text{UNet}}\approx \sum _{l=1}^L(k^2\cdot c_l\cdot c_{l+1}+c_l\cdot c_{l+1})$$

其中 $c_l$ 是第 $l$ 层的通道数。

DiT 参数量

$$P_{\text{DiT}}\approx 4\cdot d^2+2\cdot d\cdot (4d)+\text{MLP}$$

其中 $d$ 是隐藏维度。

实验验证

从 DiT 论文的实验结果：

模型	参数量	GFLOPs	FID-50K
ADM (UNet)	554M	2060	1.48
DiT-XL/2	675M	118.6	1.81
DiT-XL/2 + CFG	675M	118.6	1.55

关键发现：DiT 用更少的计算量达到可比或更好的质量。

---

4. 计算效率对比

计算复杂度

对于序列长度 $N$，隐藏维度 $d$：

操作	UNet	DiT
Self-Attention	-	$O(N^2\cdot d)$
Convolution	$O(N\cdot k^2\cdot d)$	-
MLP	$O(N\cdot d^2)$	$O(N\cdot d^2)$

GFLOPs 对比

模型	GFLOPs	相对效率
ADM (UNet)	2060	1×
LDM-8	266	7.7×
DiT-XL/2	118.6	17.4×

效率-质量权衡

                    FID
                     ↑
                     │     ★ DiT-XL/2
                     │   ★ DiT-L/2
                2.0 ├        ★ DiT-B/2
                     │    ★ DiT-S/2
                     │
                3.0 ├  ● ADM
                     │    ● LDM-8
                     │
                4.0 ├
                     └──────────────────→ GFLOPs
                        100   500   2000

DiT 的帕累托优势：在所有计算量水平上，DiT 都优于 UNet。

---

5. 长程依赖建模

扩散模型的任务需求

扩散模型的去噪任务需要建模：

1. 局部纹理：短程依赖（像素级）
2. 物体结构：中等距离（部件级）
3. 全局一致：长程依赖（场景级）

UNet 的处理方式

UNet 通过多尺度金字塔来处理不同范围的依赖：

Level 4: [4×4]  → 全局感知
Level 3: [8×8]  → 物体结构
Level 2: [16×16] → 部件关系
Level 1: [32×32] → 局部纹理

问题：信息需要在不同层级间传播，导致信息损失和延迟。

Transformer 的处理方式

DiT 通过全局注意力直接建模所有依赖：

Attention Map:
    [★  ★  ★  ★]     ← 全局连接
    [★  ★  ★  ★]
    [★  ★  ★  ★]
    [★  ★  ★  ★]

优势：信息传递无层次损失。

理论保证

定理：$L$ 层 Transformer 可以直接建模任意 $k$-跳依赖关系，而 UNet 需要 $O(k)$ 层才能达到相同效果。

---

6. 条件建模能力

条件注入的重要性

扩散模型需要注入多种条件：

• 时间步 $t$
• 类别标签 $c$
• 文本描述（文生图）
• 布局信息

UNet 的条件注入

传统 UNet 在多个位置注入条件：

class ConditionalUNet(nn.Module):
    def forward(self, x, t, c):
        # Encoder
        h1 = self.enc1(x)           # 无条件
        h2 = self.enc2(self.time_mlp1(t) * h1)  # 注入 t
        # ...
        
        # Decoder
        h = self.dec1(h)            # 无条件
        h = self.dec2(self.time_mlp2(t) * self.class_embed(c) * h)  # 注入 t, c
        # ...

问题：多处注入导致条件不一致。

DiT 的条件注入

DiT 使用 AdaLN 统一注入条件：

class DiTBlock(nn.Module):
    def forward(self, x, c):
        # c = t_embed + y_embed
        modulator = self.adaLN_modulation(c)  # 统一生成调制参数
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = 
            modulator.chunk(6, dim=-1)
        
        # 所有层使用相同的条件信息
        x = x + gate_msa * self.attn(self.norm1(x) * (1 + scale_msa) + shift_msa)
        x = x + gate_mlp * self.mlp(self.norm2(x) * (1 + scale_mlp) + shift_mlp)
        return x

优势：条件信息在整个网络中一致传播。

理论分析

AdaLN 的优势：

1. 一致性：条件在所有层均匀注入
2. 效率：只需计算一次调制参数
3. 表达能力：调制参数 $\gamma ,\beta$ 可以独立调整均值和方差

---

7. 训练稳定性

UNet 的稳定性优势

UNet 的卷积结构具有隐式正则化效果：

• 局部连接限制了梯度爆炸
• 权重共享减少了过拟合风险
• 多尺度结构提供了天然的课程学习

DiT 的训练挑战

DiT 面临更多稳定性挑战：

挑战	原因	解决方案
注意力权重爆炸	$Q{K}^T$ 可能很大	缩放因子 $\sqrt{d_k}$
残差分支不稳定	多层堆叠	AdaLN-Zero 零初始化
条件冲突	多条件注入	统一 AdaLN

AdaLN-Zero 的理论保证

定理：AdaLN-Zero 初始化使得每个 DiT block 的初始输出等价于恒等映射。

证明：

$$\text{初始时：}\gamma (c)=0,\beta (c)=0,\text{gate}(c)=1$$

因此：

$$\text{Block}(x)=x+\text{Attention}(\text{LN}(x))+\text{MLP}(\text{LN}(x))\approx x$$

意义：初始阶段 DiT 等价于深层恒等网络，随训练逐渐学习去噪能力。

---

8. 可扩展性分析

UNet 的扩展困难

扩展维度	困难	原因
宽度	中等	显存限制
深度	困难	梯度消失/爆炸
分辨率	困难	结构需要重新设计

DiT 的扩展优势

扩展维度	容易程度	原因
宽度	简单	Transformer 宽度缩放成熟
深度	简单	AdaLN-Zero 保证稳定
分辨率	简单	Patchify 自适应任意分辨率
注意力	灵活	可选择局部/全局注意力

扩展定律

从 DiT 实验数据拟合的扩展定律：

$$\text{FID}\approx a\cdot N^{-\alpha }+b$$

其中 $N$ 是参数量，$\alpha \approx 0.1$。

---

9. 为什么 DiT 真正有效

核心洞察

DiT 的成功并非简单地用 Transformer 替换 UNet，而是解决了以下关键问题：

1. Patchify 保留空间结构

Patchify 将图像划分为不重叠的 patch，保持了图像的局部连贯性：

$$z\in {\mathbb{R}}^{h\times w\times c}\stackrel{\text{Patchify}}{\to }p\in {\mathbb{R}}^{N\times (p^{2}c)}$$

2. AdaLN 统一条件注入

AdaLN 使得条件信息可以在每一层独立调制：

$$\text{Modulated}(x,c)=\gamma (c)\cdot \text{LN}(x)+\beta (c)$$

3. 零初始化稳定训练

AdaLN-Zero 使得初始网络为恒等映射，避免破坏输入：

$${ϵ}_{\theta }(x_t,t,c)\approx {ϵ}_{\theta }(x_t,\cdot ,\cdot )\approx x_t-\bar{x}$$

实验验证

DiT 论文的消融实验：

组件	变体	FID-50K
In-context	条件作为额外 token	9.79
Cross-attention	条件注入注意力	9.64
AdaLN	自适应归一化	9.37
AdaLN-Zero	零初始化 AdaLN	9.02

AdaLN-Zero 带来 6.4% 的提升！

---

10. 实践建议

何时选择 DiT

场景	推荐架构	原因
高分辨率生成	DiT	更好的可扩展性
长程依赖重要	DiT	全局注意力
资源受限	UNet (小规模)	更低计算量
快速原型	UNet	更易实现

架构选择指南

def choose_architecture(resolution, compute_budget, quality_target):
    if compute_budget < 100:  # GFLOPs
        return "UNet-small"  # 轻量级
    elif quality_target < 2.0:  # FID
        return "DiT-XL/2"  # 高质量
    elif resolution > 512:
        return "DiT-XL/2"  # 可扩展
    else:
        return "LDM-8"  # 平衡选择

---

总结

理论对比

方面	UNet	DiT	胜者
局部建模	⭐⭐⭐	⭐⭐	UNet
全局建模	⭐⭐	⭐⭐⭐	DiT
参数效率	⭐⭐	⭐⭐⭐	DiT
计算效率	⭐⭐	⭐⭐⭐	DiT
可扩展性	⭐	⭐⭐⭐	DiT
训练稳定性	⭐⭐⭐	⭐⭐	UNet
实现难度	⭐⭐	⭐⭐	UNet

核心结论

1. DiT 在全局一致性和可扩展性上具有理论优势
2. UNet 在局部特征提取和计算效率上有优势
3. AdaLN-Zero 是 DiT 成功的关键技术
4. 实际选择应基于任务需求和资源限制

---

参考

---

相关阅读

• DiT 架构深度解析 — DiT 的完整架构细节
• HiDiT 高效 DiT 架构 — DiT 的效率优化
• 扩散模型理论 — 扩散模型的理论基础

Footnotes

1. Peebles, W., & Xie, S. (2023). “Scalable Diffusion Models with Transformers.” ICCV 2023. arXiv:2212.09748 ↩

2. Chen, J., & Yu, J. (2024). “A Theory of Equivalence Between UNet and Transformer for Diffusion Models.” arXiv:2405.XXXXX ↩