Contents

Flow Matching

论文地址:https://arxiv.org/pdf/2210.02747

Normalizing Flow

在正式介绍 Flow Matching 之前,先介绍一下 flow 方法。我没找到“最早”提出这个技术的论文或教材,我看到的比较早的一篇来自于 Google DeepMind 的这篇论文(ICML,2015)。核心的思路是:使用一系列可逆的函数,来将一个简单的概率分布变换成某个很复杂的概率分布。这一系列可逆的函数,就叫做 Normalizing Flow。

最初的 Normalizing Flow 方法里,概率分布的变换,是由一个个的函数描述的,但是这篇论文对 flow 进行了连续参数化,引入了一个“时间”变量 ,让概率分布随时间从 连续变化到 ,中间的这个变化过程,使用神经网络来进行建模。假设样本空间里的样本点 是一个含时的位置变换函数。 是一个向量场,描述坐标随时间的演化。 是一个含时的概率分布。为了记号简单,时间参数也可以写在脚标上: 对样本空间做了变换之后,概率密度也发生变化(证明见 #概率换元): 这个方法就叫做 Continuous Normalizing Flow(CNF)。 生成了概率路径 。我们希望 是一些简单的分布,如高斯分布; 是我们想建模的复杂分布。我们用神经网络建模 ,训练后就得到了概率变换的 CNF,这一 CNF 就帮助我们从简单的概率分布变换成复杂的难以建模的概率分布。这里还存在着一些物理图像:把 看作坐标随时间的变换函数,0 时刻在 位置的点,在 时刻的位置是 。那么它对时间求导得到的 就是速度向量场。因为质量不会凭空增减,一个小空间内的质量变化一定和离开这一小空间的粒子质量相互抵消,就得到了连续性方程 其中 表示向量场 的散度场。当我们想要验证一个向量场 生成了某个概率路径 时,只需要验证上述连续性方程即可。

Flow Matching

假设存在某一个概率分布 ,它非常复杂,我们难以刻画它的性质,但是却有从这个分布采样得到的若干样本点。例如:绘画在二维色彩空间中的概率分布难以刻画,但我们有很多这个概率的采样样本(一副画);人类语言在音频空间中的分布难以刻画,但我们有很多这个概率的采样样本(一段语音)。我们的目标,就是从这些样本点上,找到一个刻画 的方法。

CNF 建模与 Flow Matching

我们难以刻画 ,于是尝试建模 flow,帮助我们从简单的分布变换到这个复杂分布。初始的分布是一个标准高斯:,终态的分布是我们想要的分布 (或者至少需要 )。假设某个向量场 生成了这一概率路径,那么我们想要神经网络 去拟合这个向量场。即是说我们希望学习一个 flow,去拟合这个假想的 flow,这就是 Flow Matching 这个名字的由来。因此损失函数为 期望下标中的概率分布 指的是随机变量 服从概率分布: 。尽管这个式子看起来非常符合直觉,这个公式只是在理论上存在,实际上没法应用。因为我们其实不知道 是什么,因此就更不知道 是什么。

条件向量场和条件概率路径

虽然我们没办法去刻画 ,但是我们有若干样本点 。对每一个样本点 ,我们构造一个后验的高斯分布,并认为我们的样本点集合是从这些后验的高斯分布的连续混合中采样得到的。具体而言,定义条件概率路径 ,它满足 ,其中 是某个初始分布,例如说标准高斯。将这些条件概率路径进行混合,得到了边际概率路径 并且,这个边际概率路径的终点是我们想要的那个概率分布 假设条件向量场 生成了条件概率路径 ,那么将条件向量场进行“混合”,得到的向量场 ,生成了边际概率路径 (证明见 #生成条件概率路径): 相比于之前,我们已经前进了一大步。因为 都是可以从样本点集合进行构造的。例如,直接在两个概率分布之间进行插值,就得到了 ,进而就能反解出 。而我们构造的后验,可以混合出我们想要的那个概率路径。

条件 Flow Matching

这一段,可以说是整篇论文最核心最精彩的部分

我们用神经网络 去拟合条件向量场 ,进而得到了条件概率路径,这就是 Conditional Flow Matching (CFM)。条件 Flow Matching 的损失函数是: 这个损失已经是可以计算的了,因为 是我们构造得到的,而求期望的概率都是在后验意义下的,在数据集中选取一些样本点计算损失均值,即是上面的期望。而更加精妙的是,条件 Flow Matching 和普通的 Flow Matching 的梯度完全一样(证明见 #条件 Flow Matching 的梯度): 也就是说,我们按照条件 Flow Matching 的损失公式(容易计算)进行优化,就等同于按照原始 Flow Matching 的损失公式(存在,但无法计算)进行优化,于是最终的网络,就给出了从简单概率分布到无法刻画的概率分布的一条概率路径。

构造求解

接下来详细讲一下如何构造 。在整个概率路径上,我们都使用高斯分布,即 都是一些含时的连续函数,满足 ,使得 是一个标准高斯分布; 是一个充分小的数值,使得终态的高斯分布 集中在样本点 附近。在这个构造下,对应的 flow 是在给定时刻下,是非常简单的线性变换: 于是 代换一下反函数 ,就得到 这样一来, 就是一个很容易计算的函数了。

论文里提到,如果选择在 之间、 之间进行指数插值,就能给出 diffusion model 的公式,不过我不是很懂 diffusion,所以这部分也没理解。

Optimal Transport 性质

在上面的构造中,如果我们对 进行线性插值,公式会进一步简化: 对应的向量场和概率路径为: 这不仅使得对应的损失函数非常容易计算,更重要的是,空间中所有点的速度是常数,移动轨迹是直线。空间中所有点都是沿着直线直达终点,这是代价最低的一个变换,因此这个概率路径是最优的。而且这种直线移动的函数,是比较简单的神经网络的拟合任务,训练也会相对容易。这就是 Optimal Transport 性质。

需要注意的是,由于点的坐标随着时间会变化,上面公式 中的 是含时的:。它实际上是 作为初始位置,才是一个时间无关的变量:。求导的时候别算错了。实际上

附录

概率换元

做数学的喜欢用测度空间和 Push-Forward 算子来讲这些,flow matching 的论文里也用那些语言来描述的。其实就是一些积分、换元、Jacobian 矩阵。

高维随机变量 有概率分布 ,经过变换 后的概率分布为 。由于临域内的概率在变换前后会保持相等,按照积分坐标变换公式:

对于概率路径的情形,我们认为一个变换 flow 帮助我们构造出一条概率路径 。因此 。带入上述概率换元公式就得到

注意:变换是作用在样本空间的,而不是概率空间。别错写成 了。

生成条件概率路径

假设我们用到的函数都具有一致连续性,因此散度和积分也可以交换顺序,有:

条件 Flow Matching 的梯度

我们拿出 公式的期望里面的部分: 其中 是内积。两个式子中的 都是和 无关的,而且 因此两个损失函数关于 的梯度差只与中间的内积项的期望有关。假设一致连续性,因此积分可以换序,有: 也就是说 所以,优化条件版本的 Flow Matching 和优化原始的 Flow Matching 是一回事。