深入理解 Flow Matching：生成建模的统一理论框架

1
2015-2019: Neural ODE (Chen et al., 2018) — "连续化神经网络"
2
  │  把离散残差网络 → 微分方程
3
  │
4
2019-2021: Score Matching → Noise Conditional Score Network (NCSN)
5
  │  学 ∇log p(x) — 朗之万采样
6
  │
7
2020: DDPM (Ho et al.) — 加噪 → 去噪, 但理论不统一
8
  │
9
2022: Flow Matching (Lipman et al., Jun 2022)
10
  │  ★ 统一了: DDPM / NCSN / 任意路径
11
  │
12
2022: Conditional Flow Matching (CFM) — 条件版本
13
  │
14
2023: Rectified Flow (Liu et al.) — 最优传输路径
15
  │  ★ 被 SD3 / FLUX / MMDiT 采用
16
  │
17
2024+: Flow Matching 成为扩散模型训练的理论基础

1
@torch.no_grad()
2
def ode_sample(model, num_steps=50, device="cuda"):
3
    """
4
    x_{t-1} = x_t - (1/num_steps) * v_θ(x_t, t)
5
    t 从 1 → 0 (反向)
6
    """
7
    x = torch.randn(batch_size, dim, device=device)
8

9
    t_grid = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
10

11
    for i in range(num_steps):
12
        t = t_grid[i]
13
        t_batch = t.expand(x.shape[0])
14

15
        v = model(x, t_batch)  # 学好的向量场
16
        dt = t_grid[i] - t_grid[i + 1]
17

18
        # Euler 积分
19
        x = x - dt * v
20

21
    return x  # 近似从 p_1 逆流到 p_0

1
def cfm_loss(model, x0, noise, t):
2
    """
3
    条件流匹配损失 = MSE(预测速度, 真实速度)
4

5
    Rectified Flow: v_t = ε - x0
6
    DDPM VP-SDE:    v_t = -σ² ∇_x log p_t(x|...)
7
    """
8
    # 1) 构造条件路径
9
    xt = (1 - t) * x0 + t * noise  # RF 直线
10

11
    # 2) 解析计算条件向量场 (速度)
12
    true_v = noise - x0             # RF: d/dt[(1-t)x0 + tε] = ε - x0
13

14
    # 3) 模型预测
15
    pred_v = model(xt, t)
16

17
    # 4) MSE
18
    return ((pred_v - true_v) ** 2).mean()

1
# ODE 采样 (Flow Matching 标准)
2
@torch.no_grad()
3
def ode_sample(model, xt, num_steps=50):
4
    """确定性反向 ODE 积分。"""
5
    dt = 1.0 / num_steps
6
    for i in reversed(range(num_steps)):
7
        t = torch.full((xt.shape[0],), i / num_steps, device=xt.device)
8
        v = model(xt, t)
9
        xt = xt - dt * v   # 反向积分
10
    return xt
11

12

13
# SDE 采样 (DDPM Langevin 类型)
14
@torch.no_grad()
15
def sde_sample(model, xt, num_steps=50, beta_fn=None):
16
    """随机微分方程采样（DDPM 风格）。"""
17
    dt = 1.0 / num_steps
18
    for i in reversed(range(num_steps)):
19
        t = torch.full((xt.shape[0],), i / num_steps, device=xt.device)
20

21
        # 预测噪声
22
        eps = model(xt, t)
23

24
        # 漂移项 + 扩散项
25
        drift = -beta_fn(t) / 2 * xt
26
        diffusion = torch.sqrt(beta_fn(t)) * torch.randn_like(xt)
27

28
        xt = xt - drift * dt + diffusion * np.sqrt(dt)
29

30
    return xt

1
import torch
2
import torch.nn as nn
3
from torch.utils.data import DataLoader
4

5

6
def flow_matching_train_step(model, batch, sigma_min=1e-3, sigma_max=50.0):
7
    """
8
    通用 Flow Matching 训练步。
9
    支持任意路径 (通过 path_fn 指定)。
10
    """
11
    x0 = batch  # (B, D) 真实数据
12

13
    # 1. 采样时间步 t ∈ [0, 1]
14
    t = torch.rand(x0.shape[0], device=x0.device)
15

16
    # 2. 采样噪声
17
    eps = torch.randn_like(x0)
18

19
    # 3. 构造路径 (可替换为任意 φ)
20
    # 默认: Rectified Flow (直线)
21
    xt = (1 - t.view(-1, 1)) * x0 + t.view(-1, 1) * eps
22

23
    # 4. 解析速度场 (由 φ 的定义决定)
24
    # RF: v = ε - x0
25
    true_v = eps - x0
26

27
    # 5. 模型预测
28
    pred_v = model(xt, t)
29

30
    # 6. MSE 损失
31
    loss = ((pred_v - true_v) ** 2).mean()
32
    return loss
33

34

35
def train_flow_matching(model, dataloader, epochs=100, lr=1e-4):
36
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
37

38
    for epoch in range(epochs):
39
        for batch in dataloader:
40
            batch = batch.to(device)
41
            loss = flow_matching_train_step(model, batch)
42
            optimizer.zero_grad()
43
            loss.backward()
44
            optimizer.step()
45

46
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
47

48

49
# 使用不同路径的示例
50
class FlowMatchingPath:
51
    """可配置的 Flow Matching 路径。"""
52

53
    @staticmethod
54
    def rectified_flow(x0, eps, t):
55
        """直线: x_t = (1-t)x_0 + t ε"""
56
        return (1 - t) * x0 + t * eps, eps - x0
57

58
    @staticmethod
59
    def ve_sde(x0, eps, t, sigma_min=0.01, sigma_max=50.0):
60
        """方差爆炸 SDE 的路径。"""
61
        sigma_t = sigma_min * (sigma_max / sigma_min) ** t
62
        xt = x0 + sigma_t * eps
63
        # 速度场: dσ/dt * ε
64
        dsigma = sigma_min * (sigma_max / sigma_min) ** t * torch.log(
65
            sigma_max / sigma_min
66
        )
67
        true_v = dsigma.view(-1, 1) * eps
68
        return xt, true_v
69

70
    @staticmethod
71
    def goeginger(x0, eps, t):
72
        """GoGGinger 路径: 更稳定的插值。"""
73
        lambda_t = torch.sinh(3 * t) / torch.sinh(3)
74
        xt = (1 + lambda_t) / 2 * x0 + (1 - lambda_t) / 2 * eps
75
        true_v = 3 / 2 * (torch.cosh(3 * t) / torch.sinh(3)) * (x0 - eps)
76
        return xt, true_v

1
def consistency_distillation_step(model, student, x0, num_steps=6):
2
    """
3
    从教师模型蒸馏到少步学生模型。
4
    适用于 RF 直线路径。
5
    """
6
    eps = torch.randn_like(x0)
7
    t = torch.rand(x0.shape[0])
8

9
    # 噪声图
10
    xt = (1 - t) * x0 + t * eps
11

12
    # 教师输出 (teacher 是多步模型)
13
    with torch.no_grad():
14
        v_teacher = model(xt, t)
15
        x0_teacher = xt - t.view(-1, 1) * v_teacher
16

17
    # 学生输出 (student 是少步模型，直接预测 x0)
18
    x0_student = student(xt, t)
19

20
    # 蒸馏损失: 学生应该预测和教师一样的 x0
21
    loss = ((x0_student - x0_teacher) ** 2).mean()
22
    return loss

1
def cfm_loss_detailed(model, x0, eps, t):
2
    """
3
    完整推导的 CFM 损失，包含权重选项。
4
    """
5
    # 可选: 重要性采样时间步 (与 σ² 相关)
6
    lambda_t = 1.0  # 或 lambda(t)
7

8
    # 路径
9
    xt = (1 - t) * x0 + t * eps
10

11
    # 速度
12
    true_v = eps - x0
13

14
    # 预测
15
    pred_v = model(xt, t)
16

17
    # 加权 MSE
18
    return (lambda_t * (pred_v - true_v) ** 2).mean()

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

6

7
class FlowMatchingUNet(nn.Module):
8
    """简化的 Flow Matching UNet (与 DDPM 相同架构, 不同目标)。"""
9

10
    def __init__(self, dim=64, channels=3):
11
        super().__init__()
12
        self.time_embed = nn.Sequential(
13
            nn.Linear(64, dim * 4),
14
            nn.SiLU(),
15
            nn.Linear(dim * 4, dim * 4),
16
        )
17
        # ... 标准 UNet 组件
18
        self.final = nn.Conv2d(dim, channels, 3, padding=1)
19

20
    def forward(self, xt, t, cond=None):
21
        """
22
        xt: (B, C, H, W) — 任意时间步的中间状态
23
        t:  (B,)           — 时间步 ∈ [0, 1]
24
        """
25
        # 1) 时间嵌入
26
        t_emb = self._timestep_embedding(t * 1000, 64)
27
        t_emb = self.time_embed(t_emb)
28

29
        # 2) 噪声图
30
        h = self.init_conv(xt)
31

32
        # 3) 下采样 + 中间层 + 上采样 (标准 UNet)
33
        # ... (省略中间层)
34

35
        # 4) 输出速度场 v_θ(x_t, t)
36
        return self.final(h)
37

38

39
class FlowMatchingTransformer(nn.Module):
40
    """基于 DiT 的 Flow Matching Transformer。"""
41

42
    def __init__(self, hidden_size=1024, depth=12):
43
        super().__init__()
44
        self.time_embed = TimestepEmbedder(hidden_size)
45
        self.blocks = nn.ModuleList([TransformerBlock(hidden_size) for _ in range(depth)])
46

47
    def forward(self, xt, t, cond=None):
48
        """
49
        返回速度场 v_θ(xt, t) = ε - x0
50
        与 DiT/MMDiT 完全兼容!
51
        """
52
        t_emb = self.time_embed(t)
53
        h = self.patch_embed(xt)
54

55
        for block in self.blocks:
56
            h = block(h, t_emb)
57

58
        return self.unpatchify(self.final_layer(h))

1
# MMDiT (SD3/FLUX) 中的 Flow Matching
2
# 1) backbone: DiT → MMDiT (双流)
3
# 2) 路径: DDPM → Rectified Flow (直线)
4
# 3) 目标: 噪声 ε → 速度 v = ε - x0
5

6
# MMDiT 的 forward 就是:
7
def mmdit_forward(xt, t, txt_tokens):
8
    v_pred = mmdit(xt, t, txt_tokens)  # 预测速度场
9
    return v_pred
10

11

12
def mmdit_training_loss(xt, x0, eps, t, txt_tokens):
13
    true_v = eps - x0                  # 速度场 (解析)
14
    pred_v = mmdit_forward(xt, t, txt_tokens)
15
    return F.mse_loss(pred_v, true_v)   # Flow Matching 目标

1
论文:
2
  - Flow Matching (Lipman et al., 2022): "Flow Matching for Causal Inference"
3
  - Conditional Flow Matching (Tong et al., 2023)
4
  - Rectified Flow (Liu et al., 2022): "Flow Straight and Fast"
5
  - RF-LR (2024): "Scaling Rectified Flow for Image Understanding" (SD3 论文)
6
  - Consistency Models (Song et al., 2023)
7

8
代码:
9
  - facebookresearch/flow_matching (官方实现)
10
  - stabilityai/sd3-ref (Rectified Flow + MMDiT)
11
  - black-forest-labs/FLUX (FLUX.1)
12

13
扩展:
14
  - GFlowNets: 离散的 Flow Matching
15
  - Diffusion Schrödinger Bridge: 最优传输的另一种视角