深入理解 Diffusion Transformer (DiT)：扩散模型的 Transformer 范式

1. DiT 的核心思想#

1.1 一个让人惊讶的事实#

2020-2022 年间，扩散模型在图像生成上”大杀四方”——DALL·E 2、Imagen、Stable Diffusion 接连发布，看起来已稳坐王座。

但所有这些模型都有一个共同点：

它们的核心去噪网络都是 UNet——一个专为图像设计的卷积架构。

这就引出一个自然而大胆的问题：

如果 Transformer 已经”统治”NLP、统一了视觉理解（ViT）、推翻了卷积回归——它能不能”接管”扩散模型的核心网络？

2022 年底，UC Berkeley & NYU 的 William Peebles 与 Saining Xie 给出了回答：

Diffusion Transformer (DiT)——用纯 Transformer 替代 UNet，在 ImageNet 256×256 上以 FID 1.51 击败了 LDM、ADM 等基于 UNet 的所有扩散模型，并展现出极佳的 Gflops-vs-FID 缩放曲线。

1.2 DiT 的核心思想#

一句话概括：

把扩散模型的噪声预测网络从卷积 UNet 替换为 Vision Transformer (ViT)，并通过自适应层归一化 (AdaLN-Zero) 高效注入时间步与类别条件。

1
┌────────────────────────────────────────────────────────────────┐
2
│                  DiT 整体架构示意                                 │
3
├────────────────────────────────────────────────────────────────┤
4
│                                                                │
5
│  噪声图像 x_T                                                   │
6
│      │                                                        │
7
│      ▼                                                        │
8
│  Patchify: (H/P)×(W/P) × D                                    │
9
│      │                                                        │
10
│      ▼                                                        │
11
│  ┌──────────────────────────────┐                              │
12
│  │ + Position Embedding          │  位置编码                     │
13
│  └──────────────────────────────┘                              │
14
│      │                                                        │
15
│      ▼                                                        │
16
│  ┌──────────────────────────────────────────────────┐         │
17
│  │  N × DiT Block                                   │         │
18
│  │  ┌──────────────────────────┐                    │         │
19
│  │  │ AdaLN-Zero (调制)        │ ← (t, y)           │         │
20
│  │  ├──────────────────────────┤                    │         │
21
│  │  │ Multi-Head Self-Attention │                    │         │
22
│  │  ├──────────────────────────┤                    │         │
23
│  │  │ AdaLN-Zero               │ ← (t, y)           │         │
24
│  │  ├──────────────────────────┤                    │         │
25
│  │  │ MLP                       │                    │         │
26
│  │  └──────────────────────────┘                    │         │
27
│  └──────────────────────────────────────────────────┘         │
28
│      │                                                        │
29
│      ▼                                                        │
30
│  Linear Decode: → (H/P)×(W/P) 噪声预测                          │
31
│                                                                │
32
└────────────────────────────────────────────────────────────────┘

1.3 一个生活化的比喻#

把 DiT 与 UNet 的差别想象成两个画家：

UNet 画家：用固定大小的”画笔”（卷积核）一寸一寸画——卷积核小，效率高，但必须从局部到全局逐步建立结构
DiT 画家：先看一眼整张画（把所有 patch 全局关联起来），第一笔就能把不同区域关联在一起——更”聪明”，但需要更多计算和更多数据

1.4 DiT 在生成式 AI 历史上的位置#

1
2015-2019: GAN 主导图像生成
2
  │
3
2020: DDPM 提出 (Ho et al.) - 重新让扩散模型回到舞台
4
  │  使用 UNet + 加性噪声预测
5
  │
6
2021: LDM / Stable Diffusion
7
  │  把扩散搬到潜空间 (latent space)
8
  │  仍用 UNet
9
  │
10
2022: DiT (Peebles & Xie, 2022.12)
11
  │  **首次: 用纯 Transformer 取代 UNet**
12
  │  首次发现: 缩放律 (Gflops vs FID) 极优
13
  │
14
2023: SD3 (MM-DiT)、PixArt-α、SDXL-Turbo
15
  │  DiT 思想扩展到文本到图像
16
  │
17
2024: Sora、Wan2.1、Open Sora
18
  │  DiT 扩展到视频生成
19
  │  Transformer 彻底统一了视觉生成
20
  │
21
2025+: 视觉大模型的基础架构

2. 为什么可以替换 UNet？#

2.1 UNet 的”假设”#

UNet 在扩散模型中之所以好用，是因为它有恰到好处的归纳偏置 (Inductive Bias)：

UNet 的偏置	对扩散的贡献
局部性	像素相邻关系强，卷积核捕获短程结构
平移不变性	同一物体出现在图像任意位置，模型处理方式一致
金字塔结构	多尺度感受野 + skip connection 保留细节
参数共享	推理不同时间步很高效

这些偏置在小数据时代非常珍贵——它们让模型”少走弯路”。

2.2 DiT 的反问#

DiT 的论文提出了一个深刻的观察：

当数据规模和模型规模足够大时，归纳偏置的重要性会下降——Transformer 可以自己”学习到”这些偏置。

而在扩散生成任务上：

已经有海量图文配对数据 (LAION 等)
算力上不封顶 (DiT-XL 用了 8 卡 A100 训练若干月)
高分辨率图像对”全局理解”要求高——正是 Transformer 的强项

2.3 UNet vs DiT 对比#

维度	UNet	DiT
核心算子	2D 卷积 (ResNet block)	Self-Attention + MLP
归纳偏置	强 (局部、平移不变)	弱 (全局，由数据学习)
信息传递	跳层连接 + 多尺度	全局自注意力
时间/类别条件	拼接或 AdaGN	AdaLN-Zero (更优)
数据效率	小数据更友好	大数据更优
缩放行为	很快饱和	强大、可预测
下限 vs 上限	上限相对受限	上限更高

3. DiT 架构详解#

3.1 整体流程图#

1
输入: 噪声图像 x_t (B, C, H, W) + 时间步 t + 类别标签 y
2
                                                 │
3
        ┌────────────────────────────────────────┘
4
        ▼                                        ▼
5
   ┌─────────┐                              ┌──────────┐
6
   │ Patchify│                              │  t Embed │ + y Embed
7
   └────┬────┘                              └────┬─────┘
8
        │ (B, S, D)                              │ (B, D_cond)
9
        ▼                                        ▼
10
   ┌──────────────────────────────┐        ┌────────────┐
11
   │ + Pos. Embedding              │        │  MLP       │
12
   └──────────────────────────────┘        └─────┬──────┘
13
        │                                       │
14
        ▼                                       │
15
   ┌─────────────────────────────────────────────┘
16
   │
17
   ▼
18
┌──────────────────────────┐
19
│  N × DiT Block           │  对每个 block:
20
│  ├ AdaLN-Zero (调制)     │  c ← (t+y) Embedding
21
│  ├ Self-Attention         │  x ← x + SelfAttn(AdaLN(c)调制后的 x)
22
│  └ MLP                    │  x ← x + MLP(AdaLN(c)调制后的 x)
23
└──────────────────────────┘
24
   │
25
   ▼
26
Linear Decoder: 还原为噪声预测 (B, C, H, W)

3.2 Patchify (Patch 嵌入)#

DiT 的第一步是把 $H \times W$ 的图像转成 token 序列（与 ViT 类似）。设 patch 大小为 $p$ ，则序列长度：

S = \left(\frac{H}{p}\right) \times \left(\frac{W}{p}\right)

输入尺寸	Patch	序列长度 $S$
32×32	4	64
64×64	4	256
256×256	8	1024
256×256	16	256

1
import torch
2
import torch.nn as nn
3

4

5
class PatchEmbed(nn.Module):
6
    """将图像切分为 patch (DiT 版)。"""
7

8
    def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=768):
9
        super().__init__()
10
        self.img_size = img_size
11
        self.patch_size = patch_size
12
        self.num_patches = (img_size // patch_size) ** 2
13

14
        # 用 Conv2d 一气呵成: 切块 + 线性投影
15
        self.proj = nn.Conv2d(
16
            in_chans, embed_dim,
17
            kernel_size=patch_size, stride=patch_size,
18
        )
19

20
    def forward(self, x):
21
        """
22
        x: (B, C, H, W)
23
        return: (B, S, D)
24
        """
25
        x = self.proj(x)            # (B, D, H/p, W/p)
26
        x = x.flatten(2)            # (B, D, S)
27
        x = x.transpose(1, 2)       # (B, S, D)
28
        return x

3.3 时间步与类别嵌入#

1
class TimestepEmbedder(nn.Module):
2
    """时间步 t → 向量 (用正弦/余弦位置编码 + MLP)。"""
3

4
    def __init__(self, hidden_size, frequency_embedding_size=256):
5
        super().__init__()
6
        self.mlp = nn.Sequential(
7
            nn.Linear(frequency_embedding_size, hidden_size),
8
            nn.SiLU(),
9
            nn.Linear(hidden_size, hidden_size),
10
        )
11
        self.frequency_embedding_size = frequency_embedding_size
12

13
    @staticmethod
14
    def timestep_embedding(t, dim, max_period=10000):
15
        """Sinusoidal position embedding for t."""
16
        half = dim // 2
17
        freqs = torch.exp(
18
            -math.log(max_period)
19
            * torch.arange(half, dtype=torch.float32) / half
20
        ).to(t.device)
21
        args = t[:, None].float() * freqs[None]
22
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
23
        if dim % 2:
24
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
25
        return embedding
26

27
    def forward(self, t):
28
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
29
        return self.mlp(t_freq)
30

31

32
class LabelEmbedder(nn.Module):
33
    """类别标签 y → 向量。DiT 默认使用 class-conditional 的 CFG。"""
34

35
    def __init__(self, num_classes, hidden_size):
36
        super().__init__()
37
        # +1 是为"无条件" (CFG 时使用)
38
        self.embedding = nn.Embedding(num_classes + 1, hidden_size)
39

40
    def forward(self, labels):
41
        return self.embedding(labels)

把时间步和类别 embedding 相加，得到最终条件向量 $c$ ：

1
c = t_embedder(t) + y_embedder(y)

3.4 自注意力块 (Self-Attention)#

与 ViT 几乎完全相同——QKV 投影 + 多头 softmax-attention + 输出投影。

1
class Attention(nn.Module):
2
    """Multi-Head Self-Attention (pre-LN)。"""
3

4
    def __init__(self, dim, num_heads=8, qkv_bias=False):
5
        super().__init__()
6
        assert dim % num_heads == 0
7
        self.num_heads = num_heads
8
        self.head_dim = dim // num_heads
9
        self.scale = self.head_dim ** -0.5
10

11
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
12
        self.proj = nn.Linear(dim, dim)
13

14
    def forward(self, x):
15
        B, N, C = x.shape
16
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
17
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, h, N, d)
18
        q, k, v = qkv.unbind(0)
19

20
        # 使用 PyTorch 2.0 的 SDPA 加速
21
        x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
22
        x = x.transpose(1, 2).reshape(B, N, C)
23
        return self.proj(x)

3.5 MLP 块#

1
class Mlp(nn.Module):
2
    """GELU MLP, hidden = 4 × in_dim。"""
3

4
    def __init__(self, dim, mult=4):
5
        super().__init__()
6
        hidden = int(dim * mult)
7
        self.fc1 = nn.Linear(dim, hidden)
8
        self.act = nn.GELU(approximate="tanh")
9
        self.fc2 = nn.Linear(hidden, dim)
10

11
    def forward(self, x):
12
        return self.fc2(self.act(self.fc1(x)))

4. AdaLN-Zero：条件注入的关键创新#

4.1 为什么 AdaLN-Zero 如此重要？#

扩散模型需要告诉网络”现在是去噪到哪一步了 $t$ ”、“我要生成什么类别 $y$ ”。

UNet 时代最常用的做法是 adaptive group normalization (AdaGN)：把 $t, y$ 编码成一个向量，对每层归一化参数做调制。

DiT 论文研究了 4 种条件注入策略，发现最优方式是 AdaLN-Zero。

4.2 四种条件注入策略#

论文的图如下（简化）：

1
策略 1: In-context              策略 2: Cross-attention
2
─────────────                  ──────────────────
3
 x → [t,y, x] → Transformer    x → Transformer ──cross-attn── c
4
                                  Q K V 从 x 来
5
条件直接拼成 token                 K, V 从 c 来
6

7
策略 3: Adaptive Layer-Norm    策略 4: adaLN-Zero ✅
8
──────────────────            ─────────────────
9
x → LN → Modulated by c       x → AdaLN(c) = γ·LN·x + α
10
  → Attention/MLP               → Attention/MLP
11
                                (γ 和残差 α 都是 c 的函数)
12
                                初始化为零 → 残差为 0
13
                                (与残差连接相同效果先初始化)

4.3 AdaLN-Zero 的核心公式#

设条件嵌入 $c$ ，对每个 Transformer 块计算 6 个调制参数：

\gamma_1, \gamma_2, \beta_1, \beta_2 = \text{MLP}(c) \in \mathbb{R}^{4D}

\alpha_1, \alpha_2 = \text{MLP}(c) \in \mathbb{R}^{2D}

然后在每个残差块之前做调制：

h = a_1 + \text{Attention}\left(\gamma_1 \cdot \text{LayerNorm}(x) + \beta_1\right)

h = a_2 + \text{MLP}\left(\gamma_2 \cdot \text{LayerNorm}(h) + \beta_2\right)

“Zero” 的含义：

初始化 $\alpha$ 为零 → 整个块输出去残差效果 (identity)
这样模型在一开始就是”什么都不做”的扩散初始化
训练过程中， $\alpha$ 逐渐增长，模型逐渐学出能力

AdaLN-Zero 的三大好处：

参数高效：6 维向量调制整个块，比 cross-attention 省得多

训练稳定：zero-init 让 loss 收敛更平顺

性能领先：在所有策略中 SOTA

4.4 完整 AdaLN DiT Block 代码#

1
class DiTBlock(nn.Module):
2
    """DiT 的核心 block: AdaLN-Zero 调制 + Self-Attention + MLP。"""
3

4
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
5
        super().__init__()
6
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
7
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
8
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
9
        self.mlp = Mlp(hidden_size, mult=mlp_ratio)
10

11
        # AdaLN 参数生成: 输入 c ∈ R^D, 输出 6 * D
12
        self.adaLN_modulation = nn.Sequential(
13
            nn.SiLU(),
14
            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
15
        )
16

17
    def forward(self, x, c):
18
        """
19
        x: (B, S, D)
20
        c: (B, D)  条件嵌入 (t + y)
21
        """
22
        # 生成 6 个调制参数
23
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
24
            self.adaLN_modulation(c).chunk(6, dim=1)
25
        )
26

27
        # 1. Attention 路径 (pre-LN + AdaLN + 门控残差)
28
        x = x + gate_msa.unsqueeze(1) * self.attn(
29
            modulate(self.norm1(x), shift_msa, scale_msa)
30
        )
31

32
        # 2. MLP 路径 (pre-LN + AdaLN + 门控残差)
33
        x = x + gate_mlp.unsqueeze(1) * self.mlp(
34
            modulate(self.norm2(x), shift_mlp, scale_mlp)
35
        )
36

37
        return x
38

39

40
def modulate(x, shift, scale):
41
    """AdaLN 调制: γ·x + β (但要求 γ, β 是 token-wise 的)。"""
42
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

4.5 Zero-Initialization#

1
def _zero_init(module):
2
    """把 AdaLN 的输出初始化为 0 — 这一行很重要!"""
3
    for p in module.parameters():
4
        nn.init.zeros_(p)
5

6

7
# 在构造完 DiT 后调用一次
8
for block in model.blocks:
9
    _zero_init(block.adaLN_modulation)
10
# final_layer 的 AdaLN 同样 zero init

5. 完整的 DiT 实现#

5.1 完整主类#

1
import math
2
import torch
3
import torch.nn as nn
4
import torch.nn.functional as F
5

6

7
class DiT(nn.Module):
8
    """Diffusion Transformer (class-conditional)."""
9

10
    def __init__(
11
        self,
12
        img_size=32, patch_size=4, in_chans=3,
13
        hidden_size=1152, depth=28, num_heads=16,
14
        mlp_ratio=4.0, num_classes=1000,
15
    ):
16
        super().__init__()
17
        self.img_size = img_size
18
        self.patch_size = patch_size
19
        self.in_chans = in_chans
20
        self.out_chans = in_chans  # 预测噪声, 通道数相同
21

22
        # 1) Patch 嵌入
23
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, hidden_size)
24
        self.num_patches = self.patch_embed.num_patches
25

26
        # 2) 位置编码 (可学习)
27
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_size))
28

29
        # 3) 条件嵌入
30
        self.t_embedder = TimestepEmbedder(hidden_size)
31
        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
32

33
        # 4) Transformer blocks
34
        self.blocks = nn.ModuleList([
35
            DiTBlock(hidden_size, num_heads, mlp_ratio)
36
            for _ in range(depth)
37
        ])
38

39
        # 5) 最终 LayerNorm
40
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
41
        self.adaLN_modulation_final = nn.Sequential(
42
            nn.SiLU(),
43
            nn.Linear(hidden_size, 2 * hidden_size, bias=True),
44
        )
45

46
        # 6) 线性解码器: 把 patch tokens → 噪声 patch
47
        self.final_layer = nn.Linear(hidden_size, patch_size * patch_size * self.out_chans)
48

49
        self.initialize_weights()
50

51
    def initialize_weights(self):
52
        # 1) 位置编码用 sin-cos 初始化
53
        pos_embed = get_2d_sincos_pos_embed(
54
            self.pos_embed.shape[-1], int(self.num_patches ** 0.5)
55
        )
56
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
57

58
        # 2) 默认 Linears xavier
59
        self.apply(self._init_weights)
60

61
        # 3) AdaLN 调制层 ZERO init (关键的 AdaLN-Zero!)
62
        for block in self.blocks:
63
            nn.init.zeros_(block.adaLN_modulation[1].weight)
64
            nn.init.zeros_(block.adaLN_modulation[1].bias)
65
        nn.init.zeros_(self.adaLN_modulation_final[1].weight)
66
        nn.init.zeros_(self.adaLN_modulation_final[1].bias)
67

68
        # 4) final layer 归零
69
        nn.init.zeros_(self.final_layer.weight)
70
        nn.init.zeros_(self.final_layer.bias)
71

72
    def _init_weights(self, m):
73
        if isinstance(m, nn.Linear):
74
            nn.init.xavier_uniform_(m.weight)
75
            if m.bias is not None:
76
                nn.init.constant_(m.bias, 0)
77
        elif isinstance(m, nn.LayerNorm):
78
            if m.weight is not None:
79
                nn.init.constant_(m.weight, 1.0)
80
            if m.bias is not None:
81
                nn.init.constant_(m.bias, 0)
82

83
    def unpatchify(self, x):
84
        """
85
        x: (B, S, p*p*out_chans) → (B, C, H, W)
86
        """
87
        B = x.shape[0]
88
        p = self.patch_size
89
        h = w = int(x.shape[1] ** 0.5)
90
        x = x.reshape(B, h, w, p, p, self.out_chans)
91
        x = torch.einsum("bhwpqc->bchpwq", x)
92
        return x.reshape(B, self.out_chans, h * p, w * p)
93

94
    def forward(self, x, t, y):
95
        """
96
        x: (B, C, H, W) 噪声图像
97
        t: (B,) 时间步
98
        y: (B,) 类别标签
99
        """
100
        # 1) 切块
101
        x = self.patch_embed(x)                 # (B, S, D)
102
        x = x + self.pos_embed                   # 加位置编码
103

104
        # 2) 条件嵌入
105
        c = self.t_embedder(t) + self.y_embedder(y)
106

107
        # 3) 进入 N 个 DiT block
108
        for block in self.blocks:
109
            x = block(x, c)
110

111
        # 4) 最终 AdaLN 调制
112
        shift, scale = self.adaLN_modulation_final(c).chunk(2, dim=1)
113
        x = modulate(self.norm_final(x), shift, scale)
114

115
        # 5) 线性解码 → unpatchify → 噪声
116
        x = self.final_layer(x)                  # (B, S, p*p*C)
117
        x = self.unpatchify(x)                   # (B, C, H, W)
118
        return x

5.2 DiT 模型规模表#

变体	隐层 $D$	头数	层数 $N$	参数量	Gflops (256×256)
DiT-S	384	6	12	33M	5.9
DiT-B	768	12	12	130M	23.0
DiT-L	1024	16	24	458M	80.8
DiT-XL	1152	16	28	676M	119.4
DiT-H	1408	16	32	985M	174.5

一个有意思的发现：DiT-S 大约等于原始 UNet 的计算量，但 FID 更低——Transformer 实际更高效。

6. 训练与采样#

6.1 训练目标#

DiT 沿用 DDPM/LDM 的标准训练目标（与 UNet 完全一致）：

1
# === 标准 DDPM 训练循环 ===
2
def training_step(model, x0):
3
    """
4
    x0: 真实图像 (B, C, H, W)
5
    y:  类别标签 (B,)
6
    """
7
    # 1. 随机时间步
8
    t = torch.randint(0, num_train_timesteps, (x0.size(0),), device=x0.device)
9

10
    # 2. 采样噪声
11
    noise = torch.randn_like(x0)
12

13
    # 3. 前向扩散 (重参数化)
14
    xt = sqrt_alpha_bars[t, None, None, None] * x0 + \
15
         sqrt_one_minus_alpha_bars[t, None, None, None] * noise
16

17
    # 4. 预测噪声 (用 DiT 而非 UNet!)
18
    pred_noise = model(xt, t, y)
19

20
    # 5. MSE loss
21
    loss = F.mse_loss(pred_noise, noise)
22
    return loss

关键点：DiT 替代的只是 “pred_noise = …” 这一行，其他一切都不变。

6.2 采样 (Sampling)#

采样同样沿用 DDIM/DPM-Solver 流水线，无需修改（只要模型是噪声预测 $\epsilon_\theta$ ）：

1
@torch.no_grad()
2
def sample_ddim(model, shape, y, num_steps=50):
3
    """DDIM 采样 (50 步就足够)。"""
4
    device = next(model.parameters()).device
5

6
    # 准备时间步调度
7
    ts = torch.linspace(num_train_timesteps - 1, 0, num_steps + 1).long().to(device)
8

9
    x = torch.randn(*shape, device=device)
10

11
    for i in range(num_steps):
12
        t_cur = ts[i]
13
        t_next = ts[i + 1]
14

15
        # 预测噪声
16
        pred_noise = model(x, t_cur.expand(x.size(0)), y)
17

18
        # DDIM 更新公式
19
        alpha_cur = alpha_bars[t_cur]
20
        alpha_next = alpha_bars[t_next]
21

22
        x0_pred = (x - (1 - alpha_cur).sqrt() * pred_noise) / alpha_cur.sqrt()
23
        x = alpha_next.sqrt() * x0_pred + (1 - alpha_next).sqrt() * pred_noise
24

25
    return x

6.3 Classifier-Free Guidance (CFG)#

DiT 与 LDMs 一样使用 CFG 来增强条件生成的视觉质量：

1
@torch.no_grad()
2
def sample_cfg(model, x, t, y, null_class, guidance_scale=4.0):
3
    """
4
    CFG = unconditional + scale * (conditional - unconditional)
5
    """
6
    eps_cond = model(x, t, y)
7
    eps_uncond = model(x, t, null_class * torch.ones_like(y))
8
    return eps_uncond + guidance_scale * (eps_cond - eps_uncond)

训练时需要有概率地把 $y$ 替换为 null 类别，让网络学会无条件生成：

1
def training_step_cfg(model, x0, y, uncond_prob=0.1):
2
    t = torch.randint(0, num_train_timesteps, (x0.size(0),), device=x0.device)
3

4
    # 10% 概率用 null 类别
5
    drop_mask = torch.rand(y.shape).to(y.device) < uncond_prob
6
    y_dropped = torch.where(drop_mask, null_class, y)
7

8
    noise = torch.randn_like(x0)
9
    xt = sqrt_alpha_bars[t] * x0 + sqrt_one_minus_alpha_bars[t] * noise
10
    pred = model(xt, t, y_dropped)
11
    return F.mse_loss(pred, noise)

7. 缩放律：DiT 最惊艳的结果#

7.1 Gflops vs FID 的近线性关系#

DiT 论文的一个核心贡献是发现：

FID（Fréchet Inception Distance，衡量生成质量的关键指标）随模型 Gflops 的增加，近似呈幂律下降——而且在极高 Gflops 时仍未饱和。

模型	Gflops	FID ↓	IS ↑
DiT-S (256)	5.9	11.6	92.5
DiT-B (256)	23.0	8.7	124.0
DiT-L (256)	80.8	5.7	156.0
DiT-XL (256)	119.4	4.5	173.0
DiT-XL (512)	525.0	2.3	241.0
当时 SOTA (ADM-UNet)	-	4.0	-

这条曲线意味着：只要继续把模型做大、做更长时间、喂更多数据，DiT 还能持续变好——没有”瓶颈”。

7.2 论文的缩放图#

1
FID
2
  │
3
4 │              ● DiT-XL/512 (2.3)
4
  │          ●
5
  │       ●
6
8 │    ●
7
  │  ●
8
  │● DiT-S
9
  └─────────────────────────────▶
10
    5G   20G   80G  120G     500G
11
            Gflops

7.3 与 UNet 的对比关键数据#

在相同的 Gflops 量级下：

相同 Gflops	UNet (ADM)	DiT	差距
~80 G	FID 4.0+	FID 5.7	DiT 以更低 FID 提升空间更大
缩放空间	很快饱和	持续下降	DiT 拥有更好的缩放性

8. DiT 的扩展：从图像到视频、从文本到视觉#

8.1 Stable Diffusion 3 (MM-DiT)#

Stability AI 在 SD3 中采用了多模态 DiT (MM-DiT)：

1
传统 DiT:
2
  - 把文本 c 与图像 token 拼接 → AdaLN 调制
3
  - 文本只通过 AdaLN 影响图像
4

5
MM-DiT (SD3):
6
  - 文本 token 与图像 token 在每个 attention block 里
7
    "joint attention" 一起做注意力
8
  - 每个模态有独立的投影矩阵 (Q/K/V)
9
  - 文本与图像"双向看": 文本看图像布局，图像看文本语义

1
┌─────────────────────────────────────────────┐
2
│            MM-DiT Block (SD3)                │
3
│                                              │
4
│ img_tokens, text_tokens                      │
5
│      │                                       │
6
│      ▼                                       │
7
│ ┌──────────────────────────┐                 │
8
│ │ LayerNorm                │                 │
9
│ │ Q = img_W_q @ img, ...   │ 各自不同的 W    │
10
│ │ K, V 拼接自 img + text   │                 │
11
│ │ Attn (Q, [K;V])         │                 │
12
│ └──────────────────────────┘                 │
13
│      │                                       │
14
│      ▼                                       │
15
│ MLP (各模态独立)                             │
16
└─────────────────────────────────────────────┘

8.2 Sora: 视频生成的基础#

2024 年 OpenAI 发布的 Sora 的核心架构正是 DiT 的视频版本：

把”图像 patch”换成”时空 patch” (spatial-temporal patches)
一个 (H, W, T) 的视频被切成小块
DiT 在时空维度上做自注意力 → 全局理解每一帧

1
视频 patch 化:
2
原始视频 (T=16 帧, 256×256×3)
3
   │
4
   ▼
5
切成 16×16×2 时空小块 (2 帧, 16×16 像素)
6
   │
7
   ▼
8
得到 (T/2) × (H/16) × (W/16) 个 token
9
对于 16×256×256: 8 × 16 × 16 = 2048 个 token
10
   │
11
   ▼
12
与图像 DiT 一样的 Transformer 主干

Sora 的核心贡献之一：

只要 Transformer 主干够大、数据够多，DiT 可以学到世界模型 (world models)——视频里的物理因果关系。

8.3 与其他工作的关系谱系#

1
2022.12  DiT (Peebles & Xie)
2
  │
3
2023     PixArt-α (轻量 DiT)
4
  │
5
2023     SD3 / SDXL-Turbo (MM-DiT)
6
  │
7
2024.2   Sora (视频 DiT)
8
  │
9
2024.5   Hunyuan (腾讯视频 DiT)
10
  │
11
2024-25  CogVideoX, Wan2.1 (开源视频 DiT)
12
  │       Wan 2.1 16s 480P 单卡可跑
13
  │
14
2025     Movie Gen, Veo 2 (电影级视频 DiT)

9. DiT 的工程实践#

9.1 加速技巧#

1
# 1) PyTorch 2.0 SDPA 加速注意力
2
torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True)
3
# 2) 混合精度
4
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
5
    pred = model(x, t, y)
6
# 3) torch.compile
7
model = torch.compile(model, mode="reduce-overhead")

9.2 显存优化#

梯度检查点 (Gradient Checkpointing): 把 DiT 的中间激活不存，只在反向时重算
FlashAttention: 长序列注意力降显存
分桶并行 (Bucketed Parallel): DiT-S/B/L 用单 H100 即可训练；XL/H 需要张量并行

1
from torch.utils.checkpoint import checkpoint
2

3
class DiTBlock(nn.Module):
4
    def forward(self, x, c):
5
        # 用 checkpoint 让深度网络少占显存
6
        return checkpoint(self._forward, x, c, use_reentrant=False)

9.3 Patch Size 的权衡#

情况	小 patch (p=4)	大 patch (p=8 / 16)
序列长度	长	短
计算量	大	小
局部精细度	高	低
推荐场景	高分辨率	低分辨率或大模型

9.4 类别 / 文本条件#

1
# 多类别 (含 null 类别用于 CFG)
2
self.y_embedder = nn.Embedding(num_classes + 1, hidden_size)
3

4
# 文本条件 (SD 时代)
5
self.text_embedder = ...  # CLIP / T5 等文本编码器

10. DiT 与传统 UNet 的对比总结#

维度	UNet (LDM/ADM)	DiT
核心层	2D Conv + ResNet block	Self-Attn + MLP
归纳偏置	强 (局部/平移不变)	弱 (由数据学习)
缩放行为	饱和较快	严格单调下降
高分辨率	多尺度金字塔天然适配	需要 patch 选择 + patch 化
视频扩展	需要 3D	统一时空 patch 即可
与 LLM 融合	模型架构不同	架构一致 → 便于多模态
工程成熟度	极其成熟	仍在快速迭代
训练效率	数据效率高	数据量越大越有优势

DiT 的核心遗产：它确立了**“扩散 + Transformer”为生成模型的标准范式**——既拥抱 Transformer 的缩放性，又复用扩散模型的训练范式。

11. 局限与未来#

11.1 当前局限#

数据饥渴：小数据上 DiT 比 UNet 表现差，需要 LAION/JFT 级的预训练
高分辨率成本： $256^2$ 已是 1024 个 token； $1024^2$ 就是 16384 个 token，二次方的注意力很贵
离散化训练：decoder 直接预测连续 patch 值，与 VAE decoder 配合时会有损失
CFG 仍是必需：多模态质量提升对 CFG scale 依赖大

11.2 未来方向#

1
加速方向:
2
  ├── Patch packing (更稀疏的 token)
3
  ├── Linear attention (Mamba-style) 替代 self-attention
4
  ├── 蒸馏到少步采样 (Consistency Model, SDXL-Turbo)
5
  └── 潜在 patch 化 (更紧凑表示)
6

7
能力扩展:
8
  ├── 视频 DiT (Sora, Wan, CogVideoX)
9
  ├── 3D 视觉 (从图像扩散到 3D 资产生成)
10
  ├── 音频 DiT (MusicLDM, AudioLDM 2)
11
  ├── 多模态统一 DiT (同时文本/图/音/视频)
12
  └── 具身控制 DiT (机器人动作生成)

11.3 一句话总结 DiT 的影响#

DiT 不只是替换了 UNet，它证明了生成模型的”图像 → 多模态”过渡应当建立在 Transformer 之上——这一判断现在被 Sora、Stable Diffusion 3、Gemini 视频等多模态生成系统反复印证。

12. 总结#

12.1 核心要点回顾#

维度	关键要点
核心创新	用 Transformer 替换 UNet 作为去噪网络
关键设计	Patchify + AdaLN-Zero 调制
条件注入	$c = t_{\text{emb}} + y_{\text{emb}}$ ，全块 AdaLN-Zero 调制
缩放行为	FID vs Gflops 严格服从幂律下降
优势	全局感受野、缩放性好、与 LLM 同架构
劣势	数据饥渴、 $O(S^2)$ 注意力高分辨率昂贵
影响	Sora、SD3、Wan2.1 全部建立在 DiT 之上

12.2 推荐学习资源#

1
论文:
2
  - DiT (2022): "Scalable Diffusion Models with Transformers" (Peebles & Xie)
3
  - U-ViT (2022): "All are Worth Words"
4
  - SD3 (2024): "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis"
5

6
代码:
7
  - facebookresearch/DiT (官方实现)
8
  - HuggingFace diffusers (DiTPipeline)
9
  - stabilityai/sd3 (MM-DiT 商业级)
10

11
扩展阅读:
12
  - ViT (2020): Transformer 如何进入视觉
13
  - DDPM / DDIM 原文: 扩散基础
14
  - Sora 技术报告 (2024): 视频 DiT 巅峰

一句话总结：DiT 通过”patchify + AdaLN-Zero + 纯 Transformer 主干”这一简洁设计，把扩散模型推入 Transformer 时代，确立了图像/视频生成统一范式——也是 Sora 的技术源泉。

Federico