3.1.1 可学习的动力学模型#

给定一个先验分布 $p_{prior}(x)$ 和一个真实数据分布 $p_{data}(x)$，其中对于真实数据分布 $p_{data}(x)$ 我们只能得到服从于它的样本，但是得不到分布的概率密度函数，而先验分布 $p_{prior}(x)$ 我们是知道概率密度函数的。流匹配需要实现从先验分布到真实数据分布的转换。

正如 DDPM 模型中所做的那样，可以设计一套转换的动力学过程，来完成这一任务。在 DDPM 中，前向过程和逆向过程的动力学过程都是具有解析形式的。考虑到神经网络具有强大的拟合能力，是否可以考虑使用神经网络来学习得到一个合适的速度场呢？答案是肯定的。

3.1.2 数学形式#

流匹配模型假定存在有速度场 $u^{target}(x_t, t)$，使得样本的演化遵循

并且保证它诱导的中间分布 $p_t(x)$ 满足边界条件

根据先前的分析，这样定义的速度场 $u^{target}(x_t, t)$ 和它所诱导出的中间概率分布族 $\{p_t(x)\}_{t\in[0,1]}$ 满足 Liouville 方程

那么只要求得了速度场 $u^{target}(x_t, t)$，就可以通过 ODE 的动力学过程实现从 $p_{prior}(x)$ 到 $p_{data}(x)$ 的转换。实际上，速度场 $u^{target}(x_t, t)$ 并不一定存在解析形式，因此往往是采用神经网络学习的方式来得到这样的满足条件的速度场的，参数化为 $u_\theta(x_t, t)$。

满足条件的速度场是否唯一？

答案是不唯一。

1. 多种概率转移路径: 在连续时间动力学系统中，两端分布的约束并不足以唯一确定中间的演化过程。换句话说，存在无数种不同的中间分布族 $\{p_t(x)\}_{t\in[0,1]}$ 可以实现从 $p_{prior}(x)$ 到 $p_{data}(x)$ 的转换。即有无数条可能的概率路径 (probability trajectories) 连接起这两个分布。

2. Liouville 方程对速度场的约束不足: Liouville 方程只给出了速度场的散度约束，而没有给出完整的速度场信息。即使确定了中间分布族 $\{p_t(x)\}_{t\in[0, 1]}$，对应的速度场也可能有无数个。

3.1.3 初始目标函数#

直觉上，我们希望让模型预测的速度场 $u_\theta(x_t, t)$ 接近真实的速度场 $u^{target}(x_t, t)$。这里假设我们知道在每个时间点 $t$ 下，样本 $x_t$ 的真实速度场 $u^{target}(x_t, t)$，虽然它的闭式解未必存在。使用边缘流匹配回归损失函数

其中，$u^{target}(x_t, t)$ 满足 Liouville 方程 $\frac{\partial p_t(x)}{\partial t} + \nabla_x \cdot \left[p_t(x) u^{target}(x_t, t)\right] = 0$

边缘流匹配损失函数的困难

1. $u^{target}(x_t, t)$ 未知: 实际应用中，我们不知道 $u^{target}(x_t, t)$ (否则就不需要学习了)
2. $p_t(x_t)$ 设计上的困难: 直接设计出合适的、有闭式形式的，且拟合出的速度场有强泛化性能的 $p_t(x_t)$ 并不容易。即使设计有闭式的 $p_t(x_t)$，也不能保证直接求出对应 $u^{target}(x, t)$ 有闭式解。

综上考虑，需要转换这个目标函数。

如果 $p_t(x)$ 已知，能否用 Liouville 方程替代作为损失函数用于拟合 $u^{target}(x, t)$ 呢？

• 理论上可以，但非常困难：在这个条件下， Liouville 方程本身不能唯一确定 $u^{target}(x, t)$，只能确定其散度 $\nabla_x \cdot (p_t(x) u^{target}(x, t))$。换言之，Liouville 方程是”弱监督”信号。
• 工程上计算难度大: 计算 $\nabla_x\cdot \left[p_t(x) u_\theta(x, t)\right]$ 在数据高维时难度大。

3.2.1 边缘性定理#

边缘性定理 (Marginalization Theorem) 的数学形式为

它表明边缘速度场 $u^{target}(x_t, t)$ 可以通过条件速度场 $u^{target}(x_t, t | z)$ 的加权平均来表示。它的推导过程详见附录 A.2。

3.2.2 条件流匹配目标函数#

引入条件流匹配目标函数 (Conditional Flow Matching Objective)

可以证明，条件流匹配目标函数 $\mathcal{L}_{CFM}(\theta)$ 与边缘流匹配目标函数 $\mathcal{L}_{FM}(\theta)$ 在最优解上是等价的。这意味着，通过最小化 $\mathcal{L}_{CFM}(\theta)$，我们可以找到一个神经网络 $u_\theta(x_t, t)$，它能够有效地逼近真实的边缘速度场 $u^{target}(x_t, t)$。

证明: 首先，我们将 $\mathcal{L}_{CFM}(\theta)$ 的期望形式展开

根据贝叶斯定理，有 $p_t(x_t, z) = p_t(x_t | z) p_{data}(z) = p(z | x_t) p_t(x_t)$。代入上式可得

对于内层的期望，我们可以将其展开

其中，$u_\theta(x_t, t)$ 和 $\mathbb{E}_{z \sim p(z|x_t)}[u^{target}(x_t, t | z)]$ 相对于 $z$ 是常数，因此第一项前面的期望符号去掉。此外，根据边缘性定理，$u^{target}(x_t, t) = \mathbb{E}_{z \sim p(z | x_t)}\left[u^{target}(x_t, t | z)\right]$，替换 $\mathbb{E}_{z \sim p(z | x_t)}\left[u^{target}(x_t, t | z)\right]$。因此，$\mathcal{L}_{CFM}(\theta)$ 可以重写为

上式表明，$\mathcal{L}_{CFM}(\theta)$ 等于 $\mathcal{L}_{FM}(\theta)$ 加上一个与模型参数 $\theta$ 无关的方差项。因此，最小化 $\mathcal{L}_{CFM}(\theta)$ 等价于最小化 $\mathcal{L}_{FM}(\theta)$。

这个转换的优势在于，$\mathcal{L}_{CFM}$ 的计算不需要边缘分布 $p_t(x_t)$ 或边缘速度场 $u^{target}(x_t, t)$，而是依赖于条件概率路径 $p_t(x_t|z)$ 和条件速度场 $u^{target}(x_t, t|z)$，这两者都可以被设计成简单的、具有解析表达式的形式，从而使得训练变得可行。

既然有了条件速度场的 $u^{target}(x_t, t | z)$ 解析形式，那直接用解析形式生成数据不就行了？为什么还需要训练神经网络？

• 条件速度场需要确定的真实数据点: 条件速度场 $u^{target}(x_t, t | z)$ 是基于特定的真实数据点 $z$ 定义的。它的作用相当于是为数据的演化提供了一个方向信息。然而，推理过程中，我们不会有真实的数据点 $z$ 可供使用，因此无法直接利用条件速度场来生成数据。
• $u^{target}(x_t, t) = \mathbb{E}_{z \sim p(z | x_t)}\left[u^{target}(x_t, t | z)\right]$ 的一种加权理解: 边缘速度场实际上相当于是对所有可能的真实数据点 $z$ 的条件速度场的加权平均。这个加权过程隐含了对数据分布的整体理解，而不仅仅是单个数据点的演化方向。让 $u_\theta(x_t, t)$ 去拟合这个加权平均，可以让模型学会把握数据的整体模式。

3.2.3 条件概率路径与条件速度场#

为什么需要条件概率路径？

直接用设置有简单闭式表达式的边际概率路径 (marginal probability path) $\{p_t(x)\}_{t\in[0,1]}$ 训练其对应的速度场 $v_\theta(x, t)$ 是非常困难的。为了解决这个问题，流匹配模型引入了条件概率路径 (conditional probability path) 的概念。其核心思想是，不直接对整个分布的演化进行建模，而是为每一对噪声样本 $x_0 \sim p_0(x)$ 和数据样本 $x_1 \sim p_1(x)$ 定义一个简单的、确定的路径。

条件概率的边界条件

设 $z \sim p_{data}(z)$ 来自于真实数据分布。为了实现 $p_0(x_0) = p_{prior}(x_0), p_1(x_1) = p_{data}(x_1)$，条件概率族 $\{p_t(x_t | z)\}_{t\in[0,1]}$ 同样需要满足一定的边界条件。

• 先验条件分布：$p_0(x_0 | z) = p_0(x_0)$，即 $x_0$ 与 $z$ 独立
• 数据条件分布：$p_1(x_1 | z) = \mathcal{N}(x_1 | z, \sigma_1^2 I_d)$，其中 $\sigma_1$ 是很小的正数，用于将分布紧紧包围在数据点 $z$ 附近。事实上，当 $\sigma_1 \to 0$ 时，就是 Dirac Delta 分布 $p_1(x_1 | z) = \delta_z(x_1) = \begin{cases}\infty, & x_1 = z \\ 0, & x_1 \neq z\end{cases}$

条件概率路径的设计

为了便于实现，将条件概率 $p_t(x_t | z)$ 设置为一个简单的、具有闭式表达式的路径族

这个形式就将条件概率路径的设计转换成了对均值函数 $\mu_t(z)$ 和标准差函数 $\sigma_t$ 的设计。现在我们只需要保证 $\mu_t(z)$ 和 $\sigma_t$ 满足边界条件就可以了

条件速度场

根据条件概率路径的定义，可以计算出对应的条件速度场。条件概率 $p_t(x_t | z)$ 和条件速度场 $u^{target}(x_t, t | z)$ 满足条件 Liouville 方程：

当条件概率路径被设计为高斯分布 $p_t(x_t | z) = \mathcal{N}(x_t | \mu_t(z), \sigma_t^2 I_d)$ 时，可以推导出对应的条件速度场具有一个简单的解析形式：

其中 $\dot{\mu_t}(z) = \frac{\partial \mu_t(z)}{\partial t}$，$\dot{\sigma_t} = \frac{d\sigma_t}{dt}$。

推导过程:

(i) 对数概率的推导

对于高斯分布 $p_t(x_t | z) = \mathcal{N}(x_t | \mu_t(z), \sigma_t^2 I_d)$，其对数概率为

对时间 $t$ 求导，得到

其中 $\dot{\mu_t}(z) = \frac{\partial\mu_t(z)}{\partial t}$，$\dot{\sigma}_t = \frac{d\sigma_t}{dt}$。

(ii) Liouville 方程的对数形式

另一方面，将 Liouville 方程两边除以 $p_t$ 得到

其中有利用性质，$\nabla\cdot (a\vec{v}) = a \nabla\cdot \vec{v} + \vec{v} \cdot \nabla a$。

(iii) 联立求解条件速度场

对于高斯分布

代入上式

将两个 $\frac{\partial \log p_t}{\partial t}$ 的表达式联立，得到关于 $u^{target}(x_t, t|z)$ 的方程

可以发现当 $u^{target}(x_t, t|z)$ 取以下形式时方程成立

(iv) 代入验证

将这两项相减，正好得到 $\frac{\partial \log p_t(x_t|z)}{\partial t}$ 的表达式，证明了该速度场形式的正确性。关于方程的求解，可以参考附录 C.1。

这个结果非常重要，因为它表明，我们可以设计出简单的均值函数 $\mu_t(z)$ 和标准差函数 $\sigma_t$，得到一个具有闭式解的条件速度场 $u^{target}(x_t, t | z)$。这使得条件流匹配的目标函数 $\mathcal{L}_{CFM}(\theta)$ 变得完全可计算。

3.3.1 VP 过程的条件向量场 (Variance Preserving Conditional Vector Field)#

VP 过程是一种常见的扩散过程，它的特点是在前向过程 ($p_{data}(x_0) \to p_{noise}(x_1)$) 中，噪声的方差 $\sigma_t\, (t \in [0, 1])$ 随着时间的推移而增加，且满足 $\sigma_0 = 0$，$\sigma_1 < 1$。VP 过程对应的逆过程 ($p_{noise}(x_1)\to p_{data}(x_0)$) 的概率路径为

其中 $\alpha_t = e^{-\frac{T(t)}{2}}$，$T(t) = \int_0^t \beta(s)ds$，$\beta(t)$ 是一个非负的、连续的噪声尺度函数。这个条件概率路径的均值和标准差分别为 $\mu_t(z) = \alpha_t z, \sigma_t = \sqrt{1 - \alpha_t^2}$，代入可以算得条件概率场为

3.3.2 VE 过程的条件向量场 (Variance Exploding Conditional Vector Field)#

VE 过程是另一种常见的扩散过程，它的特点是在前向过程 ($p_{data}(x_0) \to p_{noise}(x_1)$) 中，噪声的方差 $\sigma_t\, (t \in [0, 1])$ 随着时间的推移而增加，且满足 $\sigma_0 = 0$，$\sigma_1 \gg 1$。VE 过程对应的逆过程 ($p_{noise}(x_1)\to p_{data}(x_0)$) 的条件概率路径为

因此，VE 过程的均值和标准差分别为 $\mu_t(z) = z, \sigma_t$，代入可以算得条件概率场为

3.3.3 最优传输条件向量场 (Optimal Transport Conditional Vector Field)#

最优传输 (Optimal Transport, OT) 过程是一种特殊的概率路径，它通过线性插值的方式连接起噪声分布和数据分布。OT 过程对应的条件概率路径为 ($p_{noise} \to p_{data}$)

因此，OT 过程的均值和标准差分别为 $\mu_t(z) = tz, \sigma_t = 1 - t$，代入可以算得条件概率场为

可以发现，OT 过程的路径相较于 VP 和 VE 过程更加直接和线性，这使得它在某些应用中具有更好的性能。

3.4.1 训练算法#

Algorithm Flow Matching Training Procedure

3.4.2 推理算法#

Algorithm Flow Matching Sampling Procedure