12 KiB
title, draft, tags
| title | draft | tags | ||||
|---|---|---|---|---|---|---|
| 03-FM-RF-FlowMatching与RectifiedFlow | false |
|
Flow Matching 与 Rectified Flow：从 ODE 求解到向量场回归的范式突破
一、从 CNF 的痛点引入
1.1 CNF 训练的计算瓶颈
在上一份笔记（NF-CNF）中，我们已经了解到 Continuous Normalizing Flow 通过将离散流叠加为连续流，成功将行列式计算从 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 核心哲学转向：为什么不直接拟合向量场？
问题的根源在于：CNF 将向量场学习与ODE 求解耦合在一起。训练阶段，我们必须"解 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
直观上，p_t 描述了在时刻 $t$，样本在从噪声向数据演化的过程中所处的概率分布状态。
直觉：如果将每个样本视为一个沿路径运动的粒子，p_t 描述的是在时刻 t 所有粒子的空间分布密度。
2.2 条件向量场与边际向量场
对于从 x_0 \sim p_0 到 x_1 \sim p_1 的演化，条件向量场 (Conditional Vector Field) u_t(x_t | x_1) 描述的是：给定终点为 x_1 时，粒子在时刻 t 所受到的速度场驱动。
边际向量场 v_t(x_t) 则是所有条件向量场的加权平均：
v_t(x_t) = \mathbb{E}_{x_1 \sim p_1(x_1)} \left[ u_t(x_t | x_1) \cdot w(x_t, x_1) \right]
其中 w 是与路径定义相关的权重函数。
2.3 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_1 \sim p_{data}, \, x_t \sim p_t(\cdot | x_1)} \left[ \| v_\theta(x_t, t) - u_t(x_t | x_1) \|^2 \right]
其中 x_t 是根据预设路径 p_t 采样得到的中间状态。
2.4 关键定理：条件匹配等价于边际匹配
定理（Flow Matching 替换定理）：
设 u_t(x_t | x_1) 为条件向量场，v_t(x_t) 为对应的边际向量场。则最小化 CFM 损失等价于最小化边际流匹配损失：
\mathbb{E}_{t, x_t} \| v_\theta(x_t, t) - v_t(x_t) \|^2 = \mathbb{E}_{t, x_t, x_1} \| v_\theta(x_t, t) - u_t(x_t | x_1) \|^2
证明思路：
对每个 $x_1$，边际向量场 v_t(x_t) 可以表示为条件向量场的期望：
v_t(x_t) = \mathbb{E}_{x_1} \left[ u_t(x_t | x_1) | x_t \right]
因此，对于固定时刻 t 和 $x_t$，误差 \| v_\theta - v_t \|^2 可以展开为：
\| v_\theta - \mathbb{E}[u_t | x_t] \|^2 = \mathbb{E}_{x_1} \left[ \| v_\theta - u_t(x_t | x_1) \|^2 | x_t \right] - \| \mathbb{E}[u_t | x_t] - v_\theta \|^2
第二项是常数。因此，最小化 \mathbb{E} \| v_\theta - v_t \|^2 等价于最小化 $\mathbb{E} | v_\theta - u_t |^2$。
直觉：这就像是说，"所有人对正确答案的误差平方的平均"最小化，等价于"每个人对正确答案的误差平方"最小化。条件匹配提供了足够的监督信号。
三、Rectified Flow (RF) 与路径直线化
3.1 线性插值路径
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 是这条直线上的一个插值点。
3.2 1-Rectified Flow 的速度场推导
沿线性路径，速度场 \frac{d x_t}{d t} 为：
\frac{d x_t}{d t} = x_1 - x_0
关键发现：在直线插值下，条件向量场 u_t(x_t | x_1) 简化为一个与 t 和 x_t 无关的常数向量 $x_1 - x_0$！
这意味着我们可以用一个与时间无关的向量场来驱动样本从噪声演化为数据。
数学验证：
对 x_t = (1-t)x_0 + t x_1 求导：
\frac{d x_t}{d t} = -x_0 + x_1 = x_1 - x_0
得证。
3.3 1-Step 生成公式
路径完全直线化后，ODE 求解退化为单步仿射变换：
x_1 = x_0 + v_\theta(x_0)
其中 v_\theta 是学习到的速度场。
物理意义：这相当于在噪声 x_0 上直接加一个"位移向量" $v_\theta(x_0)$，该位移向量指向对应数据点 x_1 的方向。与 DDPM 需要多步迭代不同，Rectified Flow 可以通过单步前向传播直接从噪声映射到数据。
3.3 Re-flow 流程：迭代掰直路径
然而，线性插值生成的路径未必是最优的——它可能导致粒子在演化过程中经过低密度区域（即"路径交叉"问题）。
**Re-flow（Rectification）**是一种迭代训练策略，用于逐步"掰直"流线：
Algorithm: Re-flow
- 初始化：使用线性路径训练初始向量场
v_\theta^{(0)} - 迭代优化：
- 给定当前向量场 $v_\theta^{(k)}$，解 ODE 生成样本轨迹
- 估计轨迹的"曲率"（Curvature）
- 重新定义路径：令新路径为上一条轨迹的直线插值
- 在新路径上训练更新的向量场
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.4 曲率衰减的物理直觉
路径的"曲率"衡量的是粒子速度场的变化剧烈程度。当曲率较高时，相邻粒子可能沿完全不同的路径演化，导致最终分布与目标分布产生偏差。
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 | 向量场回归 |
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 耦合问题 (Coupling) 与最优传输缓解
问题根源：
在训练时，起点 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) 性质。
解决方案：Unbalanced Flow Matching 与最优传输映射：
最优传输（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 是所有满足边际约束的联合分布。
计算上的挑战：精确求解 OT 需要 Sinkhorn 等近似算法，复杂度为 $O(N^2)$；在大规模训练中，这可能成为新的瓶颈。
折衷方案：在训练早期使用独立采样（简单），在后期引入 OT 指导的路径（高效）。
六、应用案例：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。
七、总结
Flow Matching 将连续流模型的训练目标从"求解 ODE"重新定义为"向量场回归"。通过预先设计概率路径，训练过程退化为简单的监督学习。Rectified Flow 通过 Re-flow 迭代，显著降低了路径曲率，实现了少步采样。
从 CNF 到 FM，再到 RF，这一演进反映了一个根本性的数学转向：从微分方程的数值解转向统计优化的函数拟合。这不仅降低了计算复杂度，也为生成模型的理论统一提供了新的视角。
延伸阅读：
- Lipman et al., "Flow Matching for Generative Modeling" (2022)
- Albergo & Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants" (2023)
- Liu et al., "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" (2023)
- Stable Diffusion 3 Paper, "Flow Matching at Scale" (2024)