0. TL; DR

连续归一化流（Continuous normalizing flows, CNFs）一直受到基于模拟的最大似然训练方法的限制。本文引入了广义条件流匹配（generalized conditional flow matching, CFM）技术，这是一系列用于CNF的无模拟训练目标。CFM具有一个稳定的回归目标，类似于用于训练扩散模型中随机流的目标，但它享有确定性流模型的有效推理优势。

与扩散模型和先前的CNF训练算法不同，CFM不要求源分布是高斯分布，也不需要评估其密度。该目标的一个变体是最优传输CFM（optimal transport CFM, OT-CFM），它创建了更简单的流，这些流在训练中更稳定，并导致更快的推理速度，这一点在作者的实验中得到了评估。此外，作者表明，当真实的最优传输方案可用时，作者的OT-CFM方法可以近似动态OT。

1. 背景介绍

生成式建模考虑的是近似和从一个概率分布中采样的问题。归一化流作为一种有竞争力的生成建模方法，通过构建一个可逆且可高效微分的映射，将一个固定的、简单的分布（如标准正态分布）转换为复杂的数据分布。CNF则通过一个神经普通微分方程（neural ordinary differential equation, ODE）来表达这个映射。然而，CNF的训练和扩展到大规模数据集一直存在困难。

与此同时，作为当前许多生成任务中最先进的技术，扩散模型（diffusion models）通过近似一个将简单密度转换为数据分布的随机微分方程（stochastic differential equation, SDE）来工作。扩散模型的成功部分归功于其简单的回归训练目标，该目标在训练期间无需模拟SDE。最近，流匹配（flow matching, FM）表明，CNF也可以使用类似于训练扩散模型的ODE漂移项回归来进行训练。FM被证明可以产生高质量的样本并稳定CNF的训练，但它假设源分布是高斯的。

本文的第一个主要贡献是提出了一个统一的条件流匹配（conditional flow matching, CFM）框架，用于具有任意传输映射的FM模型，从而推广了现有的FM和扩散建模方法。

CNF（ODE）和扩散（SDE）模型的一个主要缺点是，为了生成高质量的样本，需要多次通过网络进行积分，导致推理时间很长。这促使研究人员在神经ODE中强制执行最优传输（OT）属性，从而产生更“直”的流，可以在更少的神经网络评估中被准确积分。本文的第二个主要贡献是CFM的一个变体，称为最优传输CFM（OT-CFM），它通过CNF来近似动态OT。作者表明，OT-CFM不仅提高了训练和推理的效率，而且比现有的基于ODE、SDE或输入凸神经网络的神经OT模型能得到更准确的OT流。

2. 方法介绍

CFM的核心思想是通过一个统一的框架，使用回归目标来训练CNF，使其能够将任意的源分布$q_0$映射到任意的目标分布$q_1$。

2.1 生成概率路径混合的向量场

假设边际概率路径$p_t(x)$是随某个条件变量$z$变化的概率路径$p_t(x|z)$的混合：

\[p_t(x) = \int p_t(x|z)q(z) dz\]

如果概率路径$p_t(x|z)$是由向量场$u_t(x|z)$从初始条件$p_0(x|z)$生成的，那么边际概率路径$p_t(x)$是由一个“边际向量场”$u_t(x)$生成的，这个边际向量场是所有条件向量场的加权平均：

\[u_t(x) := E_{q(z)} \frac{u_t(x|z) p_t(x|z)}{p_t(x)}\]

直接计算这个边际向量场是困难的，因为它涉及到对$p_t(x)$的积分。因此作者提出了一个可行的随机回归目标，即CFM目标，用于让一个神经网络$v_\theta$去逼近这个$u_t(x)$：

\[L_{CFM}(\theta) := E_{t, q(z), p_t(x|z)} \|v_\theta(t, x) - u_t(x|z)\|^2\]

作者证明，最小化这个CFM目标等价于最小化在边际路径上的流匹配目标。这使得我们可以在只知道条件路径和条件向量场的情况下，学习到正确的边际向量场。

2.2 条件概率路径的来源

根据对条件变量$z$、条件路径$p_t(\cdot|z)$和条件向量场$u_t(\cdot|z)$的不同选择，可以衍生出CFM的多种变体。

2.2.1 从高斯分布出发的流匹配 (FM)

条件变量$z$就是一个目标数据点$x_1$。条件路径被定义为一个从标准正态分布$N(0,I)$流向以$x_1$为中心的高斯分布$N(x_1, \sigma^2)$的路径。这种方法虽然有效，但局限于源分布必须是高斯分布。

2.2.2 独立耦合的CFM (I-CFM)

将条件变量$z$定义为一个点对$(x_0, x_1)$，其中$x_0$从源分布$q_0$中采样，$x_1$从目标分布$q_1$中独立采样。条件路径被定义为一个从$x_0$流向$x_1$的线性高斯路径：

\[p_t(x|z) = N(x|tx_1 + (1-t)x_0, \sigma^2)\]

对应的条件向量场非常简单：

\[u_t(x|z) = x_1 - x_0\]

这种设置允许源分布$q_0$是任意的，而不再局限于高斯分布。

2.2.3 最优传输CFM (OT-CFM)

与I-CFM不同，这里的点对$(x_0, x_1)$不再是独立采样的，而是根据源分布$q_0$和目标分布$q_1$之间的2-Wasserstein最优传输方案$\pi$进行联合采样：

\[q(z) := \pi(x_0, x_1)\]

在这种设置下，当噪声$\sigma^2 \to 0$时，学习到的边际向量场$u_t(x)$正好是解决了$q_0$和$q_1$之间动态最优传输问题的解。这为使用CNF和无模拟训练来求解动态OT问题提供了理论基础。

在实践中，对于大规模数据集，计算完整的OT方案是不可行的。因此，作者采用了一种小批量OT（minibatch OT）的近似策略。即在每个训练批次内部，计算源样本和目标样本之间的OT方案，并据此进行采样。

2.2.4 薛定谔桥CFM (Schrödinger bridge CFM)

通过将采样方案$q(z)$替换为熵正则化的OT方案，并将条件路径$p_t(x|z)$替换为布朗桥（Brownian bridge），CFM目标可以用来训练一个ODE来匹配薛定谔桥（Schrödinger bridge）的概率流。作者证明这样得到的边际向量场生成的边际概率路径，与薛定谔桥问题的解是相同的。

3. 实验分析

作者在一系列任务上对CFM及其变体进行了广泛的实验评估。

3.1 低维数据实验

作者在四对2D分布之间进行了生成和传输实验。

作者使用归一化路径能量（NPE）来衡量模型解与真实动态OT解的接近程度。结果显示，OT-CFM在所有四对分布上的NPE都是最低的，表明它最接近真实的动态OT路径。相比之下，I-CFM和FM虽然也能生成数据，但其路径能量远高于最优值。基于模拟的CNF方法（Reg. CNF和CNF）路径能量也很高，且训练时间极长。

左图比较了不同FM变体在训练过程中的验证集误差。OT-CFM收敛速度最快，在相同的训练步数下能达到更低的误差。右图比较了在推理（生成样本）时，不同方法在给定不同函数评估次数（NFE）下的样本质量（FID）。OT-CFM在相同的NFE下能达到更高的样本质量，或者说，达到相同的样本质量需要更少的NFE。这证明了OT路径更“直”，更容易被ODE求解器积分。

SB-CFM在复现薛定谔桥流的任务上，其与真实样本的平均Wasserstein距离显著低于DSB方法，且训练速度快得多。

3.2 应用于单细胞插值

作者将CFM应用于单细胞轨迹插值任务。在这项任务中，模型需要根据前后时间点的数据，来预测被隐藏的中间时间点的细胞分布。

在三个不同的单细胞数据集（Cite, EB, Multi）上，作者比较了不同方法的插值误差（EMD）。结果显示，OT-CFM在所有三个数据集上的平均表现都是最好的，优于I-CFM, DSB以及早期的TrajectoryNet和Reg. CNF等方法。