291 lines
12 KiB
Markdown
291 lines
12 KiB
Markdown
---
|
||
title: 03-FM-RF-FlowMatching与RectifiedFlow
|
||
draft: false
|
||
tags:
|
||
- Flow-Matching
|
||
- Rectified-Flow
|
||
- 扩散模型
|
||
- 生成模型
|
||
---
|
||
|
||
# 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**
|
||
|
||
1. **初始化**：使用线性路径训练初始向量场 $v_\theta^{(0)}$
|
||
2. **迭代优化**：
|
||
- 给定当前向量场 $v_\theta^{(k)}$，解 ODE 生成样本轨迹
|
||
- 估计轨迹的"曲率"（Curvature）
|
||
- **重新定义路径**：令新路径为上一条轨迹的直线插值
|
||
- 在新路径上训练更新的向量场 $v_\theta^{(k+1)}$
|
||
3. **收敛**：当路径曲率足够小时停止
|
||
|
||
**为什么路径变直后，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. **采样效率**：路径直线化后，理论上可实现 1 步生成；即使不完全直线化，也可在 5-10 步内完成，远少于 DDPM 的 50+ 步。
|
||
|
||
2. **数学简洁性**：损失函数退化为简单的 MSE 回归，无须处理 KL 散度或变分下界。
|
||
|
||
3. **灵活性**：路径可以任意设计，包括从最优传输理论推导出的"最优路径"。
|
||
|
||
---
|
||
|
||
## 五、训练实战与潜在问题
|
||
|
||
### 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，这一演进反映了一个根本性的数学转向：**从微分方程的数值解转向统计优化的函数拟合**。这不仅降低了计算复杂度，也为生成模型的理论统一提供了新的视角。
|
||
|
||
---
|
||
|
||
**延伸阅读**：
|
||
|
||
1. Lipman et al., "Flow Matching for Generative Modeling" (2022)
|
||
2. Albergo & Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants" (2023)
|
||
3. Liu et al., "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" (2023)
|
||
4. Stable Diffusion 3 Paper, "Flow Matching at Scale" (2024) |