Shengtao Yao

1969-12-31

Flow Matching 笔记

1. 核心问题

我们有一批数据 $x \in \mathbb{R}^d$（比如图片），想学习如何从简单的噪声分布生成这些数据。

2. 速度场与 ODE

定义一个速度场 $u_t(x): [0,1] \times \mathbb{R}^d \to \mathbb{R}^d$，概率流由常微分方程定义：

$dx = u_t(x) \, dt \tag{1}$

直观理解

想象一群粒子在水中漂流：

$u_t(x)$ 就是”水流”——告诉每个位置 $x$ 在时刻 $t$ 应该往哪走、走多快
$t=0$：粒子从噪声分布 $p_0$（通常是标准高斯）出发
$t=1$：粒子到达目标分布 $p_1$（真实数据）

边界条件：$x_0 \sim p_0$，$x_1 \sim p_1$

3. Flow Map 形式

按照 Lipman 的原始论文，(1) 式还可以写成：

$\frac{d}{dt}\psi_t(x) = u_t(\psi_t(x)) \tag{2}$

符号解释

符号	含义
$x$	初始位置（固定，$t=0$ 时刻的起点）
$\psi_t(x)$	起点 $x$ 在时刻 $t$ 的位置（轨迹）
$u_t(\cdot)$	时刻 $t$ 的速度场

直观理解：追踪一只蚂蚁

想象你在追踪一只蚂蚁的移动：

t=0        t=0.5       t=1
 ●    →     ●     →    ●
 x         ???        目标
(起点)                (终点)

$x$ = 蚂蚁的出发点（固定不变的标签）
$\psi_t(x)$ = 这只蚂蚁在时刻 $t$ 的实际坐标

时刻	符号	含义
$t=0$	$\psi_0(x) = x$	蚂蚁在起点
$t=0.5$	$\psi_{0.5}(x)$	蚂蚁走到一半的位置
$t=1$	$\psi_1(x)$	蚂蚁到达终点

公式的含义：

$\frac{d}{dt}\psi_t(x) = u_t(\psi_t(x))$

翻译成人话：蚂蚁位置的变化速度 = 速度场在蚂蚁当前位置的值

即：蚂蚁按照速度场 $u_t$ 的指示移动。

为什么要引入 $\psi$？ 因为 $x$ 被”占用”了——它专门表示起点。我们需要一个新符号来表示”从 $x$ 出发后，随时间变化的位置”。

4. 两个公式的关系

(1) 和 (2) 本质相同，只是视角不同：

	公式 (1)	公式 (2)
视角	欧拉视角	拉格朗日视角
$x$ 的含义	$x(t)$，随时间变化	固定的初始条件
类比	站在河边，看水流过固定点	跟着一滴水走，记录轨迹

后续推导主要以 (2) 式为主，因为 $x$ 作为初始条件更清晰。

5. Conditional Flow Matching

核心问题

我们想学 $u_t(x)$（无条件速度场），但直接学有困难：

我们只有数据样本 $x_1$，没有”真正的”速度场
这是个”因果倒置”问题：要先有速度场才能积分得到轨迹，但我们连速度场都没有

解决方案：引入条件分布

思路：给定一个目标样本 $x_1$，构造一个条件分布 $p_t(x|x_1)$，对应的条件速度场 $u_t(x|x_1)$。

边界条件

$p(t=0, x|x_1) = p(t=0, x)$：初始时与 $x_1$ 无关，就是先验分布（噪声）
$p(t=1, x|x_1) \sim \mathcal{N}(x_1, \sigma_{min}^2 I)$：终止时集中在 $x_1$ 附近的一个很小区域

时刻	条件	含义
$t=0$	$p(t=0, x\	x_1) = p(t=0, x)$	初始时与 $x_1$ 无关，就是噪声
$t=1$	$p(t=1, x\	x_1) \sim \mathcal{N}(x_1, \sigma_{min}^2 I)$	终止时集中在 $x_1$ 附近

直观理解

1
2
3

t=0                              t=1
噪声分布 ───条件速度场指引───→  集中在 x₁ 附近
(与x₁无关)                      (几乎就是 x₁)

类比：

无条件：一群蚂蚁乱跑，最后形成某个分布（但我们不知道怎么指挥）
条件：告诉每只蚂蚁”你的目标是 $x_1$”，它就有明确的方向了

为什么这样做？

条件速度场 $u_t(x|x_1)$ 变得 tractable（可计算），因为：

我们知道起点（噪声）
我们知道终点（$x_1$）
可以设计一条简单的路径连接它们

6. 条件分布的构造和代理目标

6.1 边缘分布与条件分布的关系

引入隐变量 $z$（可以理解为目标样本 $x_1$），边缘分布可以通过条件分布积分得到：

$p_t(x) = \int p_t(x|z) q(z) \, dz \tag{5}$

其中：

$p_t(x|z)$ 是条件分布
$q(z)$ 是隐变量分布（数据分布）
$p_t(x|z)$ 对应的条件速度场为 $u_t(x|z)$

6.2 连续性方程

概率分布随时间演化需要满足连续性方程（质量守恒）：

$\frac{\partial}{\partial t} p_t(x) = -\nabla \cdot (u_t(x) \cdot p_t(x)) \tag{3}$

从一维车流理解

想象一条公路上的车流：

      ← 区间 [x, x+Δx] →
──────|═══════════════|──────→ x
      x              x+Δx

      车密度 p(x)    车速 u(x)

问：这个区间内车的数量如何变化？

答：变化 = 流入 - 流出

从左边进来的流量 = $u(x) \cdot p(x)$（速度 × 密度）
从右边出去的流量 = $u(x+\Delta x) \cdot p(x+\Delta x)$

所以（一维情况）：

$\frac{\partial p}{\partial t} = -\frac{\partial (u \cdot p)}{\partial x}$

推广到高维，把 $\frac{\partial}{\partial x}$ 换成散度 $\nabla \cdot$，就得到 (3) 式。

符号解释

符号	含义
$p_t(x)$	时刻 $t$、位置 $x$ 的概率密度（”有多少粒子在这”）
$u_t(x)$	时刻 $t$、位置 $x$ 的速度（”粒子往哪走”）
$u_t \cdot p_t$	概率通量（”单位时间流过的概率质量”）
$\nabla \cdot (\cdot)$	散度（”净流出量”）
右边的负号	散度为正 = 净流出 → 密度减少

为什么叫”质量守恒”？

粒子不会凭空产生或消失：

某处密度增加 → 一定是从别处流过来的
某处密度减少 → 一定是流到别处去了

条件分布也满足类似的连续性方程：

$\frac{\partial}{\partial t} p_t(x|z) = -\nabla \cdot (u_t(x|z) \cdot p_t(x|z))$

6.3 从条件速度场推导无条件速度场

将条件连续性方程代入边缘分布的时间导数，可以得到：

$\frac{\partial}{\partial t} p_t(x) = -\nabla \cdot \int u_t(x|z) p_t(x|z) q(z) \, dz \tag{6}$

对比无条件连续性方程，我们希望：

$\nabla \cdot \int u_t(x|z) p_t(x|z) q(z) \, dz = \nabla \cdot (u_t(x) p_t(x))$

由此得到无条件速度场的表达式：

$u_t(x) = \int \frac{u_t(x|z) p_t(x|z)}{p_t(x)} q(z) \, dz \tag{7}$

直观理解：无条件速度场是所有条件速度场的加权平均，权重是各条件分布在该点的贡献。

6.4 为什么 (7) 式仍然难以计算？

虽然我们有了 $u_t(x)$ 的表达式，但积分仍然难以计算。

解决方案：不直接计算 $u_t(x)$，而是考虑 loss 函数的梯度是否等价。

6.5 两个 Loss 函数

Flow Matching Loss（我们真正想优化的）：

$L_{FM}(\theta) = \mathbb{E}_{t, x \sim p_t(x)} \| v_\theta(t, x) - u_t(x) \|^2 \tag{4}$

Conditional Flow Matching Loss（我们实际能算的）：

$L_{CFM}(\theta) = \mathbb{E}_{t \sim U(0,1), x \sim p_t(x|z), z \sim q(z)} \| v_\theta(t, x) - u_t(x|z) \|^2 \tag{8}$

6.6 核心定理：梯度相等

$\nabla_\theta L_{FM}(\theta) = \nabla_\theta L_{CFM}(\theta) \tag{9}$

证明思路：

展开两个 loss 对 $\theta$ 的梯度：

$\nabla_\theta L_{FM} = \nabla_\theta \mathbb{E}_{t, p_t(x)} [v_\theta^2] - 2 \nabla_\theta \mathbb{E}_{t, p_t(x)} [v_\theta \cdot u_t(x)]$ $\nabla_\theta L_{CFM} = \nabla_\theta \mathbb{E}_{t, p_t(x|z), q(z)} [v_\theta^2] - 2 \nabla_\theta \mathbb{E}_{t, p_t(x|z), q(z)} [v_\theta \cdot u_t(x|z)]$

关键在于第二项相等。利用 $u_t(x)$ 的定义 (7) 式：

$\mathbb{E}_{p_t(x)}[v_\theta \cdot u_t(x)] = \int v_\theta(t,x) \left[ \int \frac{u_t(x|z) p_t(x|z)}{p_t(x)} q(z) dz \right] p_t(x) dx$ $= \int \int v_\theta(t,x) \cdot u_t(x|z) \cdot p_t(x|z) \cdot q(z) \, dx \, dz = \mathbb{E}_{p_t(x|z), q(z)}[v_\theta \cdot u_t(x|z)]$

因此两个 loss 的梯度相等！

6.7 实际意义

这个结论告诉我们：

想要的	实际做的
优化 $L_{FM}$（拟合无条件速度场）	优化 $L_{CFM}$（拟合条件速度场）
需要知道 $u_t(x)$（不可行）	只需要知道 $u_t(x\	z)$（可行！）

结论：我们可以用 $L_{CFM}$ 作为代理目标，训练出的 $v_\theta(t, x)$ 可以直接当作 $u_t(x)$ 来使用！