e2hang/Notes

Fork 0

Files

e2hang 2f5211d8fe New Notes

2026-05-16 17:16:51 +08:00

16 KiB

Raw Permalink Blame History

title, draft, tags

title

draft

Flow Matching for Generative Modeling：向量场回归的范式突破

一、从 CNF/SDE 的痛点引入

1.1 CNF 训练的计算瓶颈

连续标准化流（CNF）通过将离散流叠加为连续流，成功将行列式计算从 O(D^3) 降为迹算子 $O(D)$。然而，CNF 仍面临一个根本性的训练代价：必须通过 ODE Solver 进行前向传播。

具体而言，CNF 的训练目标是最大化对数似然：

\log p_\theta(x) = \log p_0(z(0)) - \int_0^T \text{tr}\left( \frac{\partial f(x(t), t; \theta)}{\partial x(t)} \right) d t

这要求我们从数据 x 逆向积分到噪声 $z(0)$，再通过 Adjoint Method 反向传播。整个过程需要高精度的 ODE 求解，NFE (Number of Function Evaluations) 可能达到数百甚至数千。

1.2 Score SDE 的瓶颈

Score SDE 虽然提供了连续时间的理论框架，但训练时仍需要：

能够从 p_t 采样的能力（这需要 SDE 求解）
或者通过去噪得分匹配等技巧间接训练

1.3 核心哲学转向：为什么不直接拟合向量场？

问题的根源：CNF 和 Score SDE 都将向量场学习与ODE/SDE 求解耦合在一起。训练阶段必须"解 ODE"才能计算损失函数。

Flow Matching 的核心思想：

能否预先定义一条从噪声到数据的固定路径（Probability Path），然后直接拟合驱动粒子沿这条路径运动的向量场，而不经过 ODE 求解？

答案：是。

二、Flow Matching（FM）的数学推导

2.1 概率路径（Probability Paths）的定义

设 p_0 为噪声分布（通常为 $\mathcal{N}(0, I)$），p_1 为数据分布。Flow Matching 引入一个随时间演化的边际分布族 $p_t: [0,1] \times \mathbb{R}^D \rightarrow \mathbb{R}^+$，满足：

p_{t=0} = p_0, \quad p_{t=1} = p_1

直观理解：

t=0 时，所有样本都是纯噪声，服从 p_0
t=1 时，所有样本都是真实数据，服从 p_1
p_t 描述了在从噪声到数据的演化过程中，时刻 t 时所有样本的概率密度分布

物理直觉：将每个样本视为一个沿路径运动的粒子，p_t 描述的是在时刻 t 所有粒子的空间分布密度。

2.2 条件向量场与边际向量场

条件向量场（Conditional Vector Field）：u_t(x_t | x_0, x_1)

对于从 x_0 \sim p_0 到 x_1 \sim p_1 的演化，条件向量场描述的是：给定起点 x_0 和终点为 x_1 时，粒子在时刻 t 所受到的速度场驱动。

具体而言，条件向量场满足：

\frac{d x(t)}{d t} = u_t(x(t) | x_0, x_1)

边际向量场（Marginal Vector Field）：v_t(x_t)

边际向量场是所有条件向量场的加权平均：

v_t(x_t) = \mathbb{E}_{x_0 \sim p_0, x_1 \sim p_1(x_1)} \left[ u_t(x_t | x_0, x_1) \right]

物理直觉：v_t(x) 描述的是在时刻 $t$、位置 x 处，所有样本受到的平均速度向量。

2.3 条件概率路径的构造

线性插值路径

最简单的方式是线性插值：

x_t = (1-t) x_0 + t x_1, \quad t \in [0,1]

线性插值路径的速度场为：

\frac{d x_t}{d t} = x_1 - x_0

关键观察：条件向量场 u_t(x_t | x_0, x_1) = x_1 - x_0 是一个与 t 和 x_t 无关的常数向量！

高斯路径（扩散路径）

对于 DDPM 风格的前向过程：

x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon

沿此路径的速度场（对 t 求导时需注意 \epsilon 与 x_0 的独立性）：

u_t(x_t | x_0, x_1) = \frac{d x_t}{d t}

2.4 Conditional Flow Matching（CFM）损失函数

核心目标：直接回归边际向量场 $v_t$。

然而，直接监督 v_t 是困难的——因为我们无法在训练时获取 p_t 的解析形式。

关键洞察：条件向量场 u_t(x_t | x_1) 与边际向量场 v_t(x_t) 之间存在某种等价性，使得我们可以通过监督条件向量场来间接优化边际向量场。

CFM 损失函数：

\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1), \, x_0 \sim p_0, \, x_1 \sim p_{data}, \, x_t \sim p_t(\cdot | x_0, x_1)} \left[ \| v_\theta(x_t, t) - u_t(x_t | x_0, x_1) \|^2 \right]

其中 x_t 是根据预设路径 p_t 采样得到的中间状态。

2.5 关键定理：条件匹配等价于边际匹配

定理（Flow Matching Replacement Theorem）：

设 u_t(x_t | x_0, x_1) 为条件向量场，v_t(x_t) 为对应的边际向量场。则最小化条件流匹配损失等价于最小化边际流匹配损失：

\mathbb{E}_{t, x_t} \| v_\theta(x_t, t) - v_t(x_t) \|^2 = \mathbb{E}_{t, x_t, x_0, x_1} \| v_\theta(x_t, t) - u_t(x_t | x_0, x_1) \|^2 - \mathbb{E}_{t, x_t} \| v_t(x_t) - \mathbb{E}[u_t | x_t] \|^2

证明：

对每个 $x_0, x_1$，边际向量场 v_t(x_t) 可以表示为条件向量场的期望：

v_t(x_t) = \mathbb{E}_{x_0, x_1 \sim p}\left[ u_t(x_t | x_0, x_1) \mid x_t \right]

因此，对于固定时刻 t 和 $x_t$，误差 \| v_\theta - v_t \|^2 可以展开为：

\| v_\theta - \mathbb{E}[u_t | x_t] \|^2 = \mathbb{E}_{x_0, x_1} \left[ \| v_\theta - u_t(x_t | x_0, x_1) \|^2 | x_t \right] - \| \mathbb{E}[u_t | x_t] - v_\theta \|^2

第二项是常数（非负），因为 $\mathbb{E}[u_t | x_t] = v_t$。因此，最小化 \mathbb{E} \| v_\theta - v_t \|^2 等价于最小化 $\mathbb{E} | v_\theta - u_t |^2$。

直觉：这就像是说，"所有人对正确答案的误差平方的平均"最小化，等价于"每个人对正确答案的误差平方"最小化。条件匹配提供了足够的监督信号。

三、Rectified Flow（RF）与路径直线化

3.1 线性插值路径的局限

虽然线性插值路径是最简单的选择，但直接使用线性插值可能不是最优的——它可能导致粒子在演化过程中经过低密度区域（即"路径交叉"问题）。

3.2 1-Rectified Flow 的速度场推导

Rectified Flow 起源于一个最简单的直觉：直线是最短的路径。

定义线性插值路径：

x_t = (1-t) x_0 + t x_1, \quad t \in [0,1]

其中 $x_0 \sim p_0$（噪声），$x_1 \sim p_1$（数据）。

直观理解：将噪声样本 x_0 和数据样本 x_1 看作高维空间中的两个点，直接用直线连接它们。x_t 是这条直线上的一个插值点。

速度场推导：

对 x_t = (1-t)x_0 + t x_1 求导：

\frac{d x_t}{d t} = -x_0 + x_1 = x_1 - x_0

关键发现：在直线插值下，条件向量场 u_t(x_t | x_0, x_1) 简化为一个与 t 和 x_t 无关的常数向量 $x_1 - x_0$！

这意味着我们可以用一个与时间无关的向量场来驱动样本从噪声演化为数据。

3.3 单步生成公式

路径完全直线化后，ODE 求解退化为单步积分：

x_1 = x_0 + \int_0^1 v_\theta(x_t, t) d t

其中 v_\theta 是学习到的速度场。当路径完全直线化时，v_\theta(x_t, t) \approx x_1 - x_0 为常数向量，积分退化为简单的仿射变换。

物理意义：在噪声 x_0 上直接加一个"位移向量" $\int_0^1 v_\theta(x_t, t) d t$，该位移向量指向对应数据点 x_1 的方向。与 DDPM 需要多步迭代不同，Rectified Flow 可以通过单步前向传播直接从噪声映射到数据。

3.4 Re-flow 流程：迭代路径直线化

问题：线性插值生成的路径未必是最优的——它可能导致粒子在演化过程中经过低密度区域（即"路径交叉"问题），或经过不直的路径。

**Re-flow（Rectification）**是一种迭代训练策略，用于逐步"掰直"（直线化）路径：

Algorithm: Re-flow

初始化：使用线性路径训练初始向量场 v_\theta^{(0)}
迭代优化：
- 给定当前向量场 $v_\theta^{(k)}$，解 ODE 生成样本轨迹 \{x_t\}
- 重新定义路径：令新路径为上一条轨迹端点间的直线插值，即用生成样本作为新终点
- 在新路径上训练更新的向量场 v_\theta^{(k+1)}
收敛：当路径足够直（相邻轨迹终点趋于一致）时停止

为什么路径变直后，ODE Solver 只需要 1 个 Step？

对于完全直线化的路径，ODE 变为：

\frac{d x}{d t} = v(x) = \text{const}

此时，Euler Method 的单步更新即可精确求解：

x(1) = x(0) + v \cdot 1 = x(0) + (x_1 - x_0)

这正是从噪声到数据的直接映射，无需迭代。

3.5 路径直线化的物理直觉

路径的"弯曲程度"衡量的是粒子速度场的变化剧烈程度。当路径弯曲时，相邻粒子可能沿完全不同的路径演化，导致最终分布与目标分布产生偏差。

Re-flow 通过反复"直线化"路径，逐步消除这种偏差。实验表明，经过 2-3 次 Re-flow 迭代后，路径显著变直，采样质量大幅提升。

四、与扩散模型的深层联系

4.1 扩散模型作为 Flow Matching 的特例

标准的 DDPM 定义了以下前向过程（噪声注入）：

x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)

可以将其重写为 Flow Matching 的路径形式：

x_t = \alpha(t) x_0 + \beta(t) \epsilon

其中 $\alpha(t) = \sqrt{\bar{\alpha}_t}$，$\beta(t) = \sqrt{1 - \bar{\alpha}_t}$。

关键区别：

特性	DDPM	Flow Matching
路径定义	噪声调度 `\bar{\alpha}_t` 预设	任意可设计（包括线性）
训练目标	预测噪声 `\epsilon` 或 `x_0`	预测速度场 `v_\theta`
采样器	需要多步（通常 20-50 步）	可单步（路径直线化后）
数学框架	变分推断 + ELBO	向量场回归

重要结论：Flow Matching 不是"另一套完全不同的东西"，而是一个更一般的训练视角。扩散路径只是它覆盖的一类特例。

4.2 Flow Matching 的优势

采样效率：路径直线化后，理论上可实现 1 步生成；即使不完全直线化，也可在 5-10 步内完成，远少于 DDPM 的 50+ 步。
数学简洁性：损失函数退化为简单的 MSE 回归，无须处理 KL 散度或变分下界。
灵活性：路径可以任意设计，包括从最优传输理论推导出的"最优路径"。

五、训练实战与潜在问题

5.1 时间步采样策略

t 的采样分布对模型性能有显著影响。常用策略：

均匀采样（Uniform）：

t \sim \mathcal{U}(0, 1)

优点：简单；缺点：在 $t \approx 0$（噪声端）和 $t \approx 1$（数据端）处，分布变化剧烈，采样不均衡。

Logit-Normal 采样：

t = \frac{1}{1 + e^{-z}}, \quad z \sim \mathcal{N}(0, \sigma^2)

优点：使 t 更集中在中间区域（$t \approx 0.5$），这正是路径最复杂、信息最丰富的区域。

实践建议：对于大多数图像生成任务，Logit-Normal（$\sigma = 1.0$）表现优于均匀采样。

5.2 耦合问题与最优传输缓解

问题根源：

在训练时，起点 x_0 \sim p_0 和终点 x_1 \sim p_1 通常是从各自分布中独立采样的。这种独立性可能导致路径交叉——即两个不同的起点在演化过程中可能交汇到同一个终点，或反之。

数学上，这对应于联合分布 p(x_0, x_1) = p_0(x_0) p_1(x_1) 的非最优传输（Non-optimal Transport） 性质。

解决方案：最优传输映射：

最优传输（Optimal Transport, OT）理论提出，将 x_0 和 x_1 按某种最优方式配对，然后让每个粒子沿配对后的直线运动。

具体而言，求解 Monge 问题：

\min_{\pi \in \Pi(p_0, p_1)} \int c(x_0, x_1) d\pi(x_0, x_1)

其中 c(x_0, x_1) = \| x_0 - x_1 \|^2 是代价函数，\Pi 是所有满足边际约束的联合分布。

最优传输函数的作用：当边际分布 p_0 和 p_1 不平衡（如非均匀加权）时，OT 函数通过调整路径权重来保持分布的边际一致性。它不仅仅是配对，更重要的是在耦合点之间引入最优传输映射的曲率信息，使路径自动适应分布的几何结构。

计算上的挑战：精确求解 OT 需要 Sinkhorn 等近似算法，复杂度为 $O(N^2)$；在大规模训练中，这可能成为新的瓶颈。

折衷方案：在训练早期使用独立采样（简单），在后期引入 OT 指导的路径（高效）。

5.3 训练不稳定性

问题：某些路径配置下，训练可能不稳定，特别是当 t \approx 0 时，速度场变化剧烈。

解决方案：

使用更大的 batch size
采用学习率 warmup
对极端 t 值使用更小的学习率权重

六、应用案例：Stable Diffusion 3 与 Flux.1

6.1 Stable Diffusion 3（SD3）

SD3 采用了 Rectified Flow + Flow Matching 框架，替代了原有的 DDPM 采样器。

核心改进：

路径重新设计：使用 Re-flow 迭代训练，将采样步数从 50 降至 4-8 步
多模态条件融合：通过改进的 CFG 机制，在速度场空间进行条件引导，结合 Flow Matching 的向量场预测
文字渲染能力的提升：更直的路径使模型在生成包含文字的图像时，笔画更清晰、语义更一致

6.2 Flux.1

Flux.1 是目前最强的开源图像生成模型之一，其核心技术栈包括：

DiT（Diffusion Transformer） 架构
Rectified Flow 训练目标
CFG 引导：利用 classifier-free guidance 在速度场空间进行引导，而非噪声空间

Flux.1 证明了 Flow Matching 框架在大规模（12B 参数）模型上同样有效，且采样速度显著优于传统 DDPM。

七、与其他模型的关系总结

7.1 模型谱系图

生成模型
├── 基于似然
│   ├── VAE
│   ├── Normalizing Flow
│   │   └── CNF → FFJORD
│   └── Diffusion (DDPM)
│       └── Score SDE
└── 基于回归
    └── Flow Matching → Rectified Flow

7.2 核心公式对比

模型	核心方程	训练目标
NF	`x = f_K \circ \cdots \circ f_1(z)`	`\log\\|\det J_k\\|`
CNF	`dx/dt = f_\theta(x,t)`	`\int \text{tr}(\partial f/\partial x) dt`
FFJORD	`dx/dt = f_\theta(x,t)`	`\int v^T (\partial f/\partial x) v \, dt`
Score SDE	`dx = [f - gg^T \nabla\log p] dt + g d\bar{W}`	Score Matching
Flow Matching	`dx/dt = v_\theta(x,t)`	`\mathbb{E}\\|v_\theta - u_t\\|^2`

7.3 采样效率对比

模型	典型采样步数	能否单步
DDPM	50-1000	否
Score SDE	50-1000	否
CNF	10-100	否
Flow Matching	1-10	是（路径直线化）
Rectified Flow	1-4	是（Re-flow后）

八、公式速查

8.1 核心定义

概率路径：