深入理解 MMDiT：多模态扩散 Transformer 与 Stable Diffusion 3 革命

1
┌────────────────────────────────────────────────────────────────┐
2
│                  MMDiT 双流架构示意                               │
3
├────────────────────────────────────────────────────────────────┤
4
│                                                                │
5
│  文本 token  (TxtStream)  ─────┐                                │
6
│                               │                                │
7
│                               ▼                                │
8
│                     ┌──────────────────┐                       │
9
│                     │ Joint Attention  │  ← 两流都喂 Q/K/V      │
10
│                     │   (双流 Transformer Block)                │
11
│                     └────────┬─────────┘                       │
12
│                              │                                 │
13
│  图像 token  (ImgStream)  ───┘                                 │
14
│                                                               │
15
│  重复 N 次                                                    │
16
│  ↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓                                              │
17
│                                                               │
18
│  最终: 图像侧 token → 噪声预测 (用于 Rectified Flow)            │
19
│                                                                │
20
└────────────────────────────────────────────────────────────────┘

1
2020 CLIP: 图文对齐 (对比学习)
2
  │  解决了"图像和文本能不能放在同一空间"的问题
3
  │
4
2022 DiT: Transformer 替换 UNet
5
  │  解决了"扩散能不能用 Transformer 主干"
6
  │  但还只是 class-conditional
7
  │
8
2023.4 PixArt-α: 把 DiT 加上 T5 文本条件 → 早期多模态尝试
9
  │
10
2024.6 SD3 论文: "Multimodal Diffusion Transformer" (MMDiT 正式登场)
11
  │  Stability AI
12
  │  Esser et al.
13
  │  核心创新: 双流联合注意力
14
  │  配套: Rectified Flow (RF) 替代 DDPM
15
  │
16
2024.7 SD3 Medium 发布
17
  │  2B 参数的开源/中等模型
18
  │
19
2024.8 FLUX.1 (black forest labs)
20
  │  把 MMDiT + RF + DiT 推到极致
21
  │  直追 SD3 表现
22
  │
23
2024-2025 FLUX 系列: Pro, Dev, Schnell, 1.1
24
  └─ FLUX.1 一举成名

1
          文本 embeddings (B, T, D)
2
                     │
3
                     ▼
4
  ┌─────────────────────────────┐
5
  │   T5 / CLIP 文本编码器        │  (冻结或微调)
6
  └─────────────────────────────┘
7
                     │
8
                     │ txt_tokens (B, N_t, D_t)
9
                     ▼
10
          Rectified Flow timestep t
11
                     │
12
                     ▼
13
   ┌─────────────────────────────┐
14
   │  Sinusoidal timestep emb     │
15
   │  + MLP                       │  -- c_t
16
   └─────────────────────────────┘
17
                     │
18
                     ▼
19
   噪声图像 z_t (B, C, H, W)  -- Patchify --> img_tokens (B, N_v, D)
20
                     │
21
                     ▼
22
  ┌────────────────────────────────────────────┐
23
  │  N × MM-DiT Block                          │
24
  │  ┌──────────────────────────────┐          │
25
  │  │ 文本侧 LN → Q_txt, K_txt, V_txt│          │
26
  │  │ 图像侧 LN → Q_img, K_img, V_img│          │
27
  │  │                                │          │
28
  │  │ Joint Attn(Q_all=[Q_txt;Q_img],│          │
29
  │  │           K_all=[K_txt;K_img],│          │
30
  │  │           V_all=[V_txt;V_img])│          │
31
  │  │                                │          │
32
  │  │ Resplit → txt_out, img_out    │          │
33
  │  └──────────────────────────────┘          │
34
  │  + Modulation (AdaLN-Zero by c_t)          │
35
  │  + MLP (per-stream)                         │
36
  └────────────────────────────────────────────┘
37
                     │
38
                     ▼
39
   Linear Decode (img stream only)
40
                     │
41
                     ▼
42
   噪声预测 ε_θ (B, C, H, W)

1
def joint_attention(self, img_tokens, txt_tokens):
2
    """Joint Attention: 两流共享注意力矩阵，但有独立投影。"""
3
    # 1) 计算各流的 Q, K, V
4
    q_img, k_img, v_img = self.img_proj_qkv(img_tokens)  # 各自的 Linear
5
    q_txt, k_txt, v_txt = self.txt_proj_qkv(txt_tokens)  # 不同的 Linear
6

7
    # 2) 拼接: (B, N_v+N_t, D)
8
    q = torch.cat([q_img, q_txt], dim=1)
9
    k = torch.cat([k_img, k_txt], dim=1)
10
    v = torch.cat([v_img, v_txt], dim=1)
11

12
    # 3) 标准 Multi-Head Attention
13
    out = F.scaled_dot_product_attention(q, k, v)
14

15
    # 4) 拆分回两流
16
    out_img = out[:, :N_v, :]
17
    out_txt = out[:, N_v:, :]
18

19
    # 5) 各自输出投影
20
    out_img = self.img_out_proj(out_img)
21
    out_txt = self.txt_out_proj(out_txt)
22

23
    return out_img, out_txt

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

5

6
class RMSNorm(nn.Module):
7
    """FLUX 风格: 用 RMSNorm 代替 LayerNorm, 训练更稳。"""
8

9
    def __init__(self, dim, eps=1e-6):
10
        super().__init__()
11
        self.eps = eps
12
        self.weight = nn.Parameter(torch.ones(dim))
13

14
    def forward(self, x):
15
        normed = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
16
        return normed * self.weight
17

18

19
class TimestepEmbedder(nn.Module):
20
    """Sinusoidal timestep embedding."""
21

22
    def __init__(self, hidden_size, frequency_embedding_size=256):
23
        super().__init__()
24
        self.frequency_embedding_size = frequency_embedding_size
25
        self.mlp = nn.Sequential(
26
            nn.Linear(frequency_embedding_size, hidden_size),
27
            nn.SiLU(),
28
            nn.Linear(hidden_size, hidden_size),
29
        )
30

31
    @staticmethod
32
    def timestep_embedding(t, dim, max_period=10000):
33
        half = dim // 2
34
        freqs = torch.exp(
35
            -math.log(max_period)
36
            * torch.arange(half, dtype=torch.float32) / half
37
        ).to(t.device)
38
        args = t[:, None].float() * freqs[None]
39
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
40
        return embedding
41

42
    def forward(self, t):
43
        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
44

45

46
def modulate(x, shift, scale):
47
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
48

49

50
class MmDiTBlock(nn.Module):
51
    """MMDiT Single Block: Joint Attention + 双流 MLP。"""
52

53
    def __init__(
54
        self,
55
        hidden_size,
56
        img_mod_dim,
57
        txt_mod_dim,
58
        num_heads,
59
        mlp_ratio=4.0,
60
    ):
61
        super().__init__()
62
        self.num_heads = num_heads
63
        self.head_dim = hidden_size // num_heads
64
        self.scale = self.head_dim ** -0.5
65

66
        # 1) 双流 Norm
67
        self.img_norm1 = RMSNorm(hidden_size)
68
        self.txt_norm1 = RMSNorm(hidden_size)
69

70
        # 2) 双流独立 QKV 投影
71
        self.img_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
72
        self.txt_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
73
        # (实际生产代码常用 fused QKV / tensor parallel)
74

75
        # 3) 输出投影 (双流独立)
76
        self.img_out = nn.Linear(hidden_size, hidden_size, bias=False)
77
        self.txt_out = nn.Linear(hidden_size, hidden_size, bias=False)
78

79
        # 4) MLP (双流独立)
80
        self.img_mlp = nn.Sequential(
81
            nn.Linear(hidden_size, int(hidden_size * mlp_ratio)),
82
            nn.GELU(),
83
            nn.Linear(int(hidden_size * mlp_ratio), hidden_size),
84
        )
85
        self.txt_mlp = nn.Sequential(
86
            nn.Linear(hidden_size, int(hidden_size * mlp_ratio)),
87
            nn.GELU(),
88
            nn.Linear(int(hidden_size * mlp_ratio), hidden_size),
89
        )
90

91
        # 5) AdaLN-Zero 调制参数生成
92
        # 图像侧: shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp (×2 流)
93
        self.adaLN_img = nn.Sequential(nn.SiLU(), nn.Linear(img_mod_dim, 6 * hidden_size))
94
        self.adaLN_txt = nn.Sequential(nn.SiLU(), nn.Linear(txt_mod_dim, 6 * hidden_size))
95

96
    def forward(self, img_tokens, txt_tokens, c_img, c_txt):
97
        B, Nv, D = img_tokens.shape
98
        _, Nt, _ = txt_tokens.shape
99

100
        # 调制参数: (B, 6D)
101
        img_params = self.adaLN_img(c_img).chunk(6, dim=-1)
102
        txt_params = self.adaLN_txt(c_txt).chunk(6, dim=-1)
103
        img_shift_a, img_scale_a, img_gate_a, img_shift_m, img_scale_m, img_gate_m = img_params
104
        txt_shift_a, txt_scale_a, txt_gate_a, txt_shift_m, txt_scale_m, txt_gate_m = txt_params
105

106
        # ======= Joint Attention =======
107
        # 调制 + Norm
108
        img_n = modulate(self.img_norm1(img_tokens), img_shift_a, img_scale_a)
109
        txt_n = modulate(self.txt_norm1(txt_tokens), txt_shift_a, txt_scale_a)
110

111
        # QKV
112
        qi, ki, vi = self.img_qkv(img_n).chunk(3, dim=-1)
113
        qt, kt, vt = self.txt_qkv(txt_n).chunk(3, dim=-1)
114

115
        # Multi-Head reshape
116
        qi = qi.view(B, Nv, self.num_heads, self.head_dim).transpose(1, 2)
117
        ki = ki.view(B, Nv, self.num_heads, self.head_dim).transpose(1, 2)
118
        vi = vi.view(B, Nv, self.num_heads, self.head_dim).transpose(1, 2)
119

120
        qt = qt.view(B, Nt, self.num_heads, self.head_dim).transpose(1, 2)
121
        kt = kt.view(B, Nt, self.num_heads, self.head_dim).transpose(1, 2)
122
        vt = vt.view(B, Nt, self.num_heads, self.head_dim).transpose(1, 2)
123

124
        # 拼接 → Attention
125
        q = torch.cat([qi, qt], dim=2)  # (B, h, Nv+Nt, d)
126
        k = torch.cat([ki, kt], dim=2)
127
        v = torch.cat([vi, vt], dim=2)
128

129
        attn = F.scaled_dot_product_attention(q, k, v)
130

131
        attn = attn.transpose(1, 2).reshape(B, Nv + Nt, D)
132

133
        # 拆分回两流
134
        a_img = attn[:, :Nv, :]
135
        a_txt = attn[:, Nv:, :]
136

137
        # 门控残差
138
        img_tokens = img_tokens + img_gate_a.unsqueeze(1) * self.img_out(a_img)
139
        txt_tokens = txt_tokens + txt_gate_a.unsqueeze(1) * self.txt_out(a_txt)
140

141
        # ======= MLP =======
142
        img_norm2 = RMSNorm(D)
143
        txt_norm2 = RMSNorm(D)
144

145
        img_n2 = modulate(img_norm2(img_tokens), img_shift_m, img_scale_m)
146
        txt_n2 = modulate(txt_norm2(txt_tokens), txt_shift_m, txt_scale_m)
147

148
        img_tokens = img_tokens + img_gate_m.unsqueeze(1) * self.img_mlp(img_n2)
149
        txt_tokens = txt_tokens + txt_gate_m.unsqueeze(1) * self.txt_mlp(txt_n2)
150

151
        return img_tokens, txt_tokens

1
class DualTextEmbedder(nn.Module):
2
    """CLIP-G + T5-XXL 双文本编码器组合。"""
3

4
    def __init__(self):
5
        super().__init__()
6
        self.clip = CLIPTextModel.from_pretrained("openclip-vit-bigG-14")
7
        self.t5 = T5EncoderModel.from_pretrained("t5-xxl")
8
        # 冻结两个文本编码器 (实际部署时大多冻结)
9
        for p in self.clip.parameters():
10
            p.requires_grad = False
11
        for p in self.t5.parameters():
12
            p.requires_grad = False
13

14
    def forward(self, prompt_ids):
15
        clip_out = self.clip(prompt_ids["clip_ids"]).last_hidden_state       # (B, 77, 768)
16
        t5_out   = self.t5(prompt_ids["t5_ids"]).last_hidden_state            # (B, 256, 4096)
17
        return clip_out, t5_out

1
def rf_training_step(model, x0):
2
    """Rectified Flow 训练 (SD3 / FLUX 使用)。"""
3
    B = x0.size(0)
4

5
    # 1. 采样时间步 t ∈ [0, 1]
6
    t = torch.rand(B)  # uniform
7

8
    # 2. 采样噪声
9
    noise = torch.randn_like(x0)
10

11
    # 3. 沿"直线"插值
12
    xt = (1 - t.view(-1, 1, 1, 1)) * x0 + t.view(-1, 1, 1, 1) * noise
13

14
    # 4. 模型预测速度场 (velocity)
15
    v_pred = model(xt, t)
16

17
    # 5. 真实速度: ε - x0
18
    v_target = noise - x0
19

20
    # 6. MSE
21
    loss = F.mse_loss(v_pred, v_target)
22
    return loss

1
@torch.no_grad()
2
def rf_sample(model, shape, prompt_embed, num_steps=20, method="euler"):
3
    """Rectified Flow 的 ODE 采样。"""
4
    x = torch.randn(*shape)
5

6
    # 时间步网格
7
    timesteps = torch.linspace(1.0, 0.0, num_steps + 1)
8
    dt = timesteps[0] - timesteps[1]
9

10
    for i in range(num_steps):
11
        t = timesteps[i]
12
        t_next = timesteps[i + 1]
13
        t_batch = t * torch.ones(shape[0], device=x.device)
14

15
        # 预测速度场
16
        v = model(x, t_batch, prompt_embed)
17

18
        # Euler 步 (也可以用 Heun)
19
        x = x + (t_next - t) * v
20

21
    return x

1
class FluxDoubleBlock(nn.Module):
2
    """FLUX 的 Double Block: MMDiT 块 + 独立 MLP。"""
3

4
    def __init__(self, hidden_size):
5
        super().__init__()
6
        self.mmdit = MmDiTBlock(hidden_size)       # 联合注意力
7
        self.img_mlp = nn.Sequential(
8
            nn.Linear(hidden_size, 4 * hidden_size),
9
            nn.GELU(),
10
            nn.Linear(4 * hidden_size, hidden_size),
11
        )
12
        self.txt_mlp = nn.Sequential(
13
            nn.Linear(hidden_size, 4 * hidden_size),
14
            nn.GELU(),
15
            nn.Linear(4 * hidden_size, hidden_size),
16
        )
17

18
    def forward(self, img, txt, c):
19
        # 1) 联合注意力 (双流)
20
        img, txt = self.mmdit(img, txt, c, c)
21

22
        # 2) 独立 MLP (独立流)
23
        img = img + self.img_mlp(img)
24
        txt = txt + self.txt_mlp(txt)
25

26
        return img, txt
27

28

29
class FluxSingleBlock(nn.Module):
30
    """FLUX 的 Single Block: 真把两流拼接。"""
31

32
    def __init__(self, hidden_size, num_heads):
33
        super().__init__()
34
        # 单一 QKV / 单一 MLP: 两流完全共享
35
        self.qkv = nn.Linear(hidden_size, 3 * hidden_size)
36
        self.proj = nn.Linear(hidden_size, hidden_size)
37
        self.mlp = nn.Sequential(
38
            nn.Linear(hidden_size, 4 * hidden_size),
39
            nn.GELU(),
40
            nn.Linear(4 * hidden_size, hidden_size),
41
        )
42

43
    def forward(self, img, txt):
44
        # 把两流拼接成一条序列
45
        x = torch.cat([img, txt], dim=1)
46
        # 单 QKV + 单 MLP
47
        q, k, v = self.qkv(x).chunk(3, dim=-1)
48
        x = self.proj(F.scaled_dot_product_attention(q, k, v))
49
        x = x + self.mlp(x)
50

51
        Nv = img.shape[1]
52
        return x[:, :Nv, :], x[:, Nv:, :]

1
阶段 1: 预训练
2
  - 数据: LAION, COYO 等 (10B+ 图文对)
3
  - 分辨率: 256 → 512 → 1024 (渐进式增长)
4
  - 优化: AdamW, lr=1e-4, beta=(0.9, 0.95), wd=0
5
  - Batch size: 1024-4096
6
  - 时间调度: Rectified Flow (uniform [0,1])
7
  - 文本编码器: 部分微调 / 完全冻结
8

9
阶段 2: 微调 (可选)
10
  - 数据: 高质量美学数据集
11
  - 主题: DPO / 人类反馈
12

13
阶段 3: 蒸馏 (FLUX Schnell)
14
  - 蒸馏到 1-4 步 ODE 求解
15
  - 损失: 对抗 + LPIPS + Flow Matching

1
class RoPE2D(nn.Module):
2
    """2D RoPE for image. 行坐标 + 列坐标各拿一半维度。"""
3

4
    def __init__(self, head_dim, base=10000):
5
        super().__init__()
6
        assert head_dim % 4 == 0, "head_dim must be divisible by 4"
7
        self.head_dim = head_dim
8
        self.base = base
9

10
    def _rotate_half(self, x):
11
        d = x.shape[-1]
12
        half = d // 2  # 前一半: x 坐标, 后一半: y 坐标
13
        x_rot, y_rot = x[..., :half], x[..., half:]
14
        # 旋转 90°
15
        return torch.cat([-x_rot, -y_rot], dim=-1)
16

17
    def forward(self, q, k, grid_h, grid_w):
18
        B, H, N, D = q.shape
19
        half = D // 2
20

21
        # 生成 sin / cos
22
        device = q.device
23
        # 简化的实现, 实际生产常用高效 fused RoPE
24
        ys = torch.arange(grid_h, device=device).repeat_interleave(grid_w)
25
        xs = torch.arange(grid_w, device=device).repeat(grid_h)
26
        # compute frequency → apply rotation
27
        ...
28
        return q_rot, k_rot

1
class MMDiT(nn.Module):
2
    """Multimodal Diffusion Transformer (SD3 风格)。"""
3

4
    def __init__(
5
        self,
6
        img_size=64,            # latent 分辨率
7
        patch_size=2,
8
        in_chans=4,             # VAE latent 通道
9
        hidden_size=1024,
10
        depth=12,
11
        num_heads=16,
12
        mlp_ratio=4.0,
13
        txt_seq_len=256,
14
        txt_dim=4096,           # T5 输出维度
15
    ):
16
        super().__init__()
17
        self.img_size = img_size
18
        self.patch_size = patch_size
19
        self.num_patches = (img_size // patch_size) ** 2
20

21
        # 1) 图像侧 Patchify
22
        self.patch_embed = nn.Conv2d(in_chans, hidden_size, patch_size, patch_size)
23
        # 2D Sin-Cos 位置编码
24
        self.pos_embed = nn.Parameter(
25
            get_2d_sincos_pos_embed(hidden_size, int(self.num_patches ** 0.5)),
26
            requires_grad=False,
27
        )
28

29
        # 2) 文本侧 Projection
30
        self.txt_proj = nn.Linear(txt_dim, hidden_size)
31
        self.txt_pos = nn.Parameter(torch.zeros(1, txt_seq_len, hidden_size))
32
        nn.init.trunc_normal_(self.txt_pos, std=0.02)
33

34
        # 3) 条件注入
35
        self.t_embedder = TimestepEmbedder(hidden_size)
36
        # 文本向量还会与 timestep 相加
37

38
        # 4) N 个 MMDiT block
39
        self.blocks = nn.ModuleList([
40
            MmDiTBlock(hidden_size, hidden_size, hidden_size, num_heads, mlp_ratio)
41
            for _ in range(depth)
42
        ])
43

44
        # 5) Final Norm + AdaLN + Final Layer
45
        self.norm_final = RMSNorm(hidden_size)
46
        self.final_adaLN = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
47
        self.final_layer = nn.Linear(hidden_size, patch_size ** 2 * in_chans)
48

49
        # 6) Zero-init (AdaLN-Zero 必做)
50
        self._init_weights()
51

52
    def _init_weights(self):
53
        # AdaLN 的 Linear 全部 zero-init
54
        for block in self.blocks:
55
            nn.init.zeros_(block.adaLN_img[1].weight)
56
            nn.init.zeros_(block.adaLN_img[1].bias)
57
            nn.init.zeros_(block.adaLN_txt[1].weight)
58
            nn.init.zeros_(block.adaLN_txt[1].bias)
59
        nn.init.zeros_(self.final_adaLN[1].weight)
60
        nn.init.zeros_(self.final_adaLN[1].bias)
61
        nn.init.zeros_(self.final_layer.weight)
62
        nn.init.zeros_(self.final_layer.bias)
63

64
    def unpatchify(self, x):
65
        """(B, S, p*p*C) → (B, C, H, W)."""
66
        p = self.patch_size
67
        H = W = int(x.shape[1] ** 0.5)
68
        x = x.reshape(x.shape[0], H, W, p, p, -1)
69
        x = torch.einsum("bhwpqc->bchpwq", x)
70
        return x.reshape(x.shape[0], -1, H * p, W * p)
71

72
    def forward(self, z_t, t, txt_tokens):
73
        """
74
        z_t: (B, C, H, W) VAE latent (噪声后)
75
        t:   (B,)          timestep ∈ [0, 1]
76
        txt_tokens: (B, N_t, txt_dim)  T5 输出
77
        return: v_pred  (B, C, H, W)
78
        """
79
        B = z_t.shape[0]
80

81
        # 1) Patchify + 位置编码
82
        img = self.patch_embed(z_t).flatten(2).transpose(1, 2)  # (B, N_v, D)
83
        img = img + self.pos_embed.unsqueeze(0)
84

85
        # 2) 文本投影
86
        txt = self.txt_proj(txt_tokens)                          # (B, N_t, D)
87
        txt = txt + self.txt_pos
88

89
        # 3) 时间步 + 文本 (共享)
90
        c = self.t_embedder(t) + txt.mean(dim=1, keepdim=True).expand(-1, img.shape[1] + txt.shape[1], -1).reshape(B, -1).mean(dim=-1, keepdim=True)
91
        c_img = c
92
        c_txt = self.t_embedder(t)  # 文本侧也独立一个调制信号 (MMDiT 现代做法)
93
        # 实践中 c_txt 与 c_img 同源
94
        c_img = c_txt = c.squeeze() if c.dim() > 1 else c
95

96
        # 4) N 个 MMDiT block
97
        for block in self.blocks:
98
            img, txt = block(img, txt, c_img, c_txt)
99

100
        # 5) 最终 AdaLN + Linear 解码
101
        shift, scale = self.final_adaLN(c).chunk(2, dim=-1)
102
        img = modulate(self.norm_final(img), shift, scale)
103
        img = self.final_layer(img)
104
        img = self.unpatchify(img)
105

106
        return img

1
效率:
2
  ├── 共享 QKV (双流合一)
3
  ├── 量化 (FP8/INT4)
4
  ├── 蒸馏到 1-4 步 (FLUX Schnell 路线)
5
  └── 移除 CFG, 用引导替代
6

7
能力:
8
  ├── 视频 + 多帧联合注意力
9
  ├── 任意分辨率 / 长宽比 (RoPE)
10
  ├── 多语种 prompt
11
  ├── 文本编辑图像 (instruction-based inpainting)
12
  └── 4D 内容生成 (时空动态)
13

14
统一:
15
  ├── 与 LLM 共架构 (语言 + 视觉扩散统一预训练)
16
  ├── 与 AudioLDM / MusicLDM 统一 (音频 + 图像)
17
  └── 与 RL/VLA 融合 (机器人, 决策)