生成式建模的流匹配.

paper：Flow Matching for Generative Modeling

0. TL; DR

Flow Matching (FM) 是一种用于训练连续归一化流 (Continuous Normalizing Flows, CNFs) 的新范式，它实现了前所未有的训练规模和效率。具体来说，作者们提出了流匹配的概念，这是一种基于回归固定条件概率路径的向量场，从而实现无模拟（simulation-free）训练CNF的方法。

FM兼容一类通用的高斯概率路径，用于在噪声和数据样本之间进行转换。将FM与扩散路径结合使用，为训练扩散模型提供了一种更稳健和稳定的替代方案。此外，FM为使用其他非扩散概率路径来训练CNF打开了大门。一个特别值得关注的实例是使用最优传输（Optimal Transport, OT）位移插值来定义条件概率路径。这些路径比扩散路径更高效，提供了更快的训练和采样速度，并带来了更好的泛化性能。

在ImageNet上使用Flow Matching训练CNF，在似然和样本质量方面都取得了比其他基于扩散的方法持续更好的性能，并允许使用现成的数值ODE求解器进行快速可靠的样本生成。

1. 背景介绍

深度生成模型旨在从未知的数据分布中进行估计和采样。近年来，由扩散模型（diffusion models）驱动的生成式建模取得了惊人的进展，尤其是在图像生成领域。然而，扩散模型也存在其固有的局限性，例如其采样过程通常需要数百甚至上千步迭代，耗时较长。

连续归一化流（Continuous Normalizing Flows, CNFs）作为一种通用的、确定性的生成模型框架，能够模拟任意的概率路径，并且原则上可以包含扩散模型所模拟的概率路径。然而，除了可以通过分数匹配（score matching）有效训练的扩散模型外，目前还没有可扩展的CNF训练算法。传统的最大似然训练方法需要昂贵的数值ODE模拟，而现有的无模拟方法要么涉及难以处理的积分，要么存在有偏梯度的问题。

为了打破这些限制，本文提出了Flow Matching (FM)，一种用于训练CNF模型的、高效的、无模拟的方法。FM使得我们可以采用更通用的概率路径来指导CNF的训练，从而摆脱了对扩散过程的依赖，可以直接在概率路径上进行操作。

2. 流匹配

流匹配（Flow Matching, FM）的核心思想是，学习一个时变的向量场$v_\theta(x,t)$，该向量场定义了一个常微分方程（ODE），其轨迹（或称“流”）能够将一个简单的先验分布$p_0$（如高斯噪声）平滑地传输到目标数据分布$p_1$。

2.1 从概率路径到向量场

假设存在一个从$t=0$到$t=1$连续演化的概率密度路径$p_t(x)$，其中$p_0$是先验分布，$p_1$是目标数据分布。这个演化过程由一个连续性方程描述，该方程将密度的变化与一个驱动其流动的速度向量场$u_t(x)$联系起来。FM的目标就是训练一个神经网络$v_\theta(x, t)$来逼近这个真实的向量场$u_t(x)$。

FM的损失函数定义为：

\[\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t, p_t(x)} \| v_t(x) - u_t(x) \|^2\]

然而，这个目标函数在实践中是难以直接优化的，因为我们既不知道中间时刻的分布$p_t(x)$，也不知道真实的向量场$u_t(x)$。

2.2 条件流匹配 (Conditional Flow Matching, CFM)

为了解决上述问题，FM引入了一个关键的技巧：条件化。其核心思想是，不直接对复杂的边际概率路径$p_t(x)$建模，而是将其分解为一系列简单的、以单个数据点$x_1$为条件的概率路径$p_t(x

x_1)$的混合：

\[p_t(x) = \int p_t(x|x_1) q(x_1) dx_1\]

其中$q(x_1)$是数据分布。

作者们证明了两个关键定理：

定理1：边际向量场$u_t(x)$可以表示为所有条件向量场$u_t(x

x_1)$的加权平均。这个定理将一个复杂的、未知的边际向量场分解为了一系列简单的、可定义的条件向量场。

定理2：在边际路径上最小化FM损失，等价于在条件路径上最小化一个更简单的CFM损失。这个CFM损失函数为：

\[\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, q(x_1), p_t(x|x_1)} \| v_t(x) - u_t(x|x_1) \|^2\]

这个CFM目标是可计算的，因为它只涉及到从已知的分布（数据分布$q(x_1)$和我们自己定义的条件路径$p_t(x

x_1)$）中采样。这使得FM成为一个高效的、无模拟的训练框架。

2.3 两种核心的条件概率路径

CFM框架的灵活性在于我们可以自由设计条件概率路径$p_t(x

x_1)$。作者们重点探讨了两种路径。

1. 扩散路径 (Diffusion Paths)

我们可以选择那些与现有扩散模型（如VP或VE SDE）的逆向过程相对应的概率路径。作者们推导出了这些路径所对应的条件向量场$u_t(x|x_1)$的解析形式。在这种设置下，FM可以被看作是训练扩散模型的一种替代方法。实验表明，使用FM来训练扩散路径，比传统的分数匹配方法更稳定、更鲁棒。

2. 最优传输路径 (Optimal Transport Paths)

作者们定义了一种简单而高效的最优传输（OT）条件路径：

路径均值：$\mu_t(x_1) = tx_1$ (从0线性插值到$x_1$)
路径标准差：$\sigma_t = 1 - (1-\sigma_{min})t$ (从1线性插值到$\sigma_{min}$)

这条路径实际上是两个高斯分布（$p_0 = \mathcal{N}(0, I)$ 和 $p_1 = \mathcal{N}(x_1, \sigma_{min}^2I)$）之间的Wasserstein-2 OT位移插值。

与弯曲的扩散路径相比，OT路径具有以下优势：

路径更直：粒子沿直线以恒定速度运动。
向量场更简单：其条件向量场$u_t(x|x_1)$的方向不随时间改变，这使得神经网络的回归任务变得更简单。

这些优良特性使得使用OT路径的FM（FM-OT）在训练速度、采样效率和最终性能上都表现出色。

3. 实验分析

作者们在CIFAR-10和ImageNet等大规模图像数据集上对FM进行了广泛的实验验证。

3.1 密度建模与样本质量

在CIFAR-10, ImageNet 32x32 和 ImageNet 64x64 上，作者们比较了FM-OT, FM-Diffusion 以及两种基于分数匹配的扩散模型基线（SM, SF）。

在所有数据集和所有指标（负对数似然NLL、FID分数、采样函数评估次数NFE）上，FM-OT都取得了最好的结果。在ImageNet 128x128上，FM-OT也取得了与当时最先进的无条件生成模型相竞争的FID分数。

3.2 训练与采样效率

在ImageNet 64x64的训练过程中，FM-OT的FID分数下降速度最快，收敛也最快。相比于需要数百万次迭代的扩散模型，FM只需要几十万次迭代就能达到很好的效果。

定性比较显示，FM-OT的采样路径能够更早地生成图像的轮廓和结构，而扩散路径则在大部分时间内都处于去噪的早期阶段，直到最后时刻才显现出图像内容。

在ImageNet 32x32上比较了低NFE下的采样性能。FM-OT在达到相同的数值误差时所需的NFE最少。FM-OT在相同的低NFE下能够达到最低的FID分数，展示了其在快速采样方面的巨大优势。

3.3 条件图像生成

作者们将FM-OT应用于图像超分辨率任务（从64x64到256x256）。结果显示，FM-OT在FID和IS这两个衡量生成质量和多样性的指标上，显著优于基于回归和扩散模型（SR3）的基线方法。

3.4 与其他方法的定性比较

在一个二维的棋盘格生成任务中，FM-OT的轨迹能够更早地形成棋盘格的结构，而扩散模型（无论是用分数匹配还是流匹配训练）的轨迹则显得更为混沌。这再次说明了OT路径的直接性和高效性。