Flow Matching

连续归一化流

想象 $$ 空间中有一系列粒子，在 $$ 时服从分布 $$，随时间流逝粒子在空间中流动，直至 $$ 时形成分布 $$. 于是这样的流动过程形成了从分布 $$ 到分布 $$ 的一个转换。如果我们能够为上述流体运动过程建立起模型，并控制 $$ 为某简单易采样分布而 $$ 服从数据分布，那么不就得到一个生成模型了吗？这样的生成模型称作连续归一化流 (Continuous Normalizing Flows, CNFs)^[2].

概率密度路径

粒子在流动的过程中，其概率分布在不断地改变。定义概率密度路径 $$，其中 $$ 表示 $$ 时刻 $$ 位置处的概率密度。称之为“路径”是因为 $$ 可以视作无限维概率分布空间中的流形上的一条路径。

显然，对任意时刻 $$，概率密度 $$ 都应满足归一化条件：$$. 因此，区别于流体力学中一般的流（不要求密度是归一化的），我们称这种模型为归一化流；又由于时间是连续的，因此称为连续归一化流。

速度场与连续性方程

定义速度场为 $$，其中 $$ 表示 $$ 时刻 $$ 位置处粒子的运动速度。流体中的粒子沿着速度场 $$ 运动，引起概率密度的变化，因此 $$ 与 $$ 之间一定存在某种关系，这个关系式称为连续性方程： $$ 直观上，连续性方程的第一项表示单位时间内 $$ 位置处粒子的增加/减少量，第二项表示单位时间内 $$ 位置处粒子的流出/流入量，显然粒子增加量就是流入量，因此连续性方程成立。熟悉随机微分方程的读者可能会发现，连续性方程其实就是 Fokker-Planck 方程在扩散项为零时的情形。

值得注意的是，概率密度路径与速度场不是一一对应的关系，不同的速度场可以产生相同的概率密度路径。例如，给 $$ 加上散度为零的场（无源场），就得到了一个新的速度场，并且连续性方程依旧成立，所以概率密度路径不变。

流与变量替换

粒子沿着速度场运动，得到的轨迹称为流。具体而言，流 $$ 是下述常微分方程的解： $$ 其中 $$ 表示 $$ 时刻位于 $$ 处的粒子在 $$ 时刻运动到的位置。换句话说，随着 $$ 从 0 到 1 变化，$$ 形成了从 $$ 位置出发的粒子运动的轨迹。为了看得更清楚，可以对上式左右两边同时从 $$ 到 $$ 积分： $$ 左边表示从 $$ 出发的粒子在 $$ 时间内的位移，右边是对速度的积分，自然也是位移。

根据流 $$ 的定义，$$ 时刻位于 $$ 位置处的粒子在 $$ 时刻的出发位置是 $$，因此根据随机变量的变量替换公式，可以知道 $$ 与 $$ 之间有关系： $$ 简记作 $$，称作 push-forward 方程。

拓展（瞬时变量替换公式）：Push-forward 方程给出了 $$ 时刻到 $$ 时刻概率密度的变化。进一步地，Neural ODE^[2] 的作者还给出了描述概率密度变化的微分方程，称作瞬时变量替换公式： $$ 推导过程：计算 $$ 对 $$ 的全导数： $$ 其中第二行是代入了连续性方程，第三行是将散度展开，最后一行是根据流的定义有 $$. 于是我们有： $$ 这就得到了瞬时变量替换公式。

综上所述，概率密度路径、速度场和流之间的关系可表示为下图：

现在，为了构建生成模型，我们只需要构建速度场 $$，使得对应的概率密度路径 $$ 满足 $$ 为某简单分布且 $$ 为数据分布即可，这就是 flow matching 要完成的事。

Flow Matching

原始损失函数

设 $$ 是一概率密度路径，满足 $$ 为某简单分布（例如标准正态分布）且 $$ 为数据分布，$$ 为对应的未知速度场。Flow matching^[1] 使用神经网络构建速度场 $$ 去近似 $$，损失函数为： $$ 然而式中真实速度场 $$ 是未知的，因此无法直接使用 $$​ 训练，怎么办呢？聪明的读者可能已经发现了，flow matching 的训练目标和遇到的问题与 score matching 一模一样，其中速度场直接对标 score function，所以把 score matching 的解决方法搬过来即可。借鉴 denoising score matching，我们每次针对单个样本去近似条件速度场，其中样本从数据集中采样，这样均摊下来就是在近似无条件的真实速度场了。

条件速度场

具体而言，给定某特定样本 $$，称 $$ 为条件概率路径。那么，边缘概率路径可以通过条件概率路径对所有样本求期望得到： $$ 设上述条件概率路径 $$ 可以由条件速度场 $$ 得到（即二者满足连续性方程），那么可以推得，速度场 $$ 与条件速度场 $$ 有如下关系： $$

证明：我们只需要证明基于上面的方式定义的 $$ 与 $$ 满足连续性方程即可。直接代入： $$ 证毕。

熟悉 score function 的读者应该知道，对 score function 也有类似的关系式： $$

CFM 损失函数

基于条件概率路径和条件速度场，原始的 flow matching 损失函数可以改写做如下 conditional flow matching 损失函数： $$ 可以证明 $$ 与 $$ 只差与 $$ 无关的常数，因此对优化问题而言是等价的。推导思路其实与 score matching 到 denoising score matching 如出一辙。

证明：将 $$ 中的平方打开： $$ 第一项保留，第二项与 $$​ 无关扔掉，第三项做如下变形： $$ 加上第一项得： $$ 看起来不太好看，配个方就得到了 CFM 损失函数： $$

现在，只要我们设计一个可解的条件速度场 $$，就可以基于 $$ 训练了。

条件速度场的设计

一般形式

我们从条件概率路径 $$​ 入手，考虑高斯分布形式： $$ 其中 $$ 为与时间步有关的均值，$$ 为与时间步有关的标准差。特别地，我们规定：

• 当 $$ 时，设 $$，即与 $$ 无关的标准高斯分布；
• 当 $$ 时，设 $$，其中标准差 $$ 充分小，即大约是在 $$ 处的确定性分布。

直接找到对应上述条件概率路径的条件速度场是比较困难的，但是我们可以找到对应的以 $$ 为条件的流 $$： $$

可以验证上述形式的条件流 $$ 与条件概率路径 $$ 的确满足变量替换公式： $$ 换句话说，只要取 $$，那么就有 $$.

有了条件流，就可以找到对应的条件速度场 $$ 了。根据流的定义有： $$

代入 $$ 得： $$ 其中 $$ 的求导表示对时间 $$ 求导。这就是 flow matching 最终用于训练的损失函数形式。

如果只是为了训练，我们没有必要把条件速度场显式地解出来，但是这是可以做到的： $$

当 $$ 与 $$ 取特殊形式时，我们可以还原出经典的扩散模型，还可以得到比扩散路径更优的最优传输路径。

特殊情形：经典扩散

考虑 VE diffusion： $$ 这对应着 $$ 的情形。代入条件速度场得： $$ 相应的条件流为： $$ 考虑 VP diffusion： $$ 这对应着 $$​ 的情形。代入条件速度场得： $$ 相应的条件流为： $$ 可以看见，VE diffusion 或 VP diffusion 的条件流对时间 $$ 不是线性关系，也就是说它们的轨迹会拐弯。

特殊情形：最优传输

相比经典的扩散路径，一个更合适的路径是将 $$ 都设置为 $$ 的线性变换： $$ 这样得到的条件速度场为： $$ 相应的条件流为： $$ 该条件流事实上是两个高斯分布 $$​ 和 $$​ 之间的最优传输 displacement map. 直观上，沿最优传输路径运动的粒子会在直线上均匀运动，而非像扩散路径一样拐弯，如下图所示：

需要强调的是，“直线”路径只在给定 $$ 的条件下成立，在真正生成数据时粒子依旧走的是曲线路径。

总结

Flow matching 看起来蛮复杂，但是仔细想想，其实就是给已有的扩散模型框架套了个流体力学的壳，用新的体系重新说了一遍。比如条件速度场其实就是条件 score function 的线性变换，而 flow matching 的核心数学技巧也都是借用自 denoising score matching. 文章的亮点在于：

1. 不要求源分布必须是高斯分布，从而可以用来学习两个分布之间的转换。
2. 提出了基于最优传输的新的扩散路径，声称比传统的 VP 和 VE 路径更适合采样。

参考资料

1. Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 (2022). ↩︎
2. Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems 31 (2018). ↩︎
3. Tor Fjelde, Emile Mathieu, Vincent Dutordoir. AN INTRODUCTION TO FLOW MATCHING. https://mlg.eng.cam.ac.uk/blog/2024/01/20/flow-matching.html ↩︎