Flow Matching
连续归一化流
想象 空间中有一系列粒子,在 时服从分布 ,随时间流逝粒子在空间中流动,直至 时形成分布 . 于是这样的流动过程形成了从分布 到分布 的一个转换。如果我们能够为上述流体运动过程建立起模型,并控制 为某简单易采样分布而 服从数据分布,那么不就得到一个生成模型了吗?这样的生成模型称作连续归一化流 (Continuous Normalizing Flows, CNFs)[2].
概率密度路径
粒子在流动的过程中,其概率分布在不断地改变。定义概率密度路径 ,其中 表示 时刻 位置处的概率密度。称之为“路径”是因为 可以视作无限维概率分布空间中的流形上的一条路径。
显然,对任意时刻 ,概率密度 都应满足归一化条件:. 因此,区别于流体力学中一般的流(不要求密度是归一化的),我们称这种模型为归一化流;又由于时间是连续的,因此称为连续归一化流。
速度场与连续性方程
定义速度场为 ,其中 表示 时刻 位置处粒子的运动速度。流体中的粒子沿着速度场 运动,引起概率密度的变化,因此 与 之间一定存在某种关系,这个关系式称为连续性方程: 直观上,连续性方程的第一项表示单位时间内 位置处粒子的增加/减少量,第二项表示单位时间内 位置处粒子的流出/流入量,显然粒子增加量就是流入量,因此连续性方程成立。熟悉随机微分方程的读者可能会发现,连续性方程其实就是 Fokker-Planck 方程在扩散项为零时的情形。
值得注意的是,概率密度路径与速度场不是一一对应的关系,不同的速度场可以产生相同的概率密度路径。例如,给 加上散度为零的场(无源场),就得到了一个新的速度场,并且连续性方程依旧成立,所以概率密度路径不变。
流与变量替换
粒子沿着速度场运动,得到的轨迹称为流。具体而言,流 是下述常微分方程的解: 其中 表示 时刻位于 处的粒子在 时刻运动到的位置。换句话说,随着 从 0 到 1 变化, 形成了从 位置出发的粒子运动的轨迹。为了看得更清楚,可以对上式左右两边同时从 到 积分: 左边表示从 出发的粒子在 时间内的位移,右边是对速度的积分,自然也是位移。
根据流
拓展(瞬时变量替换公式):Push-forward 方程给出了
综上所述,概率密度路径、速度场和流之间的关系可表示为下图:
现在,为了构建生成模型,我们只需要构建速度场
Flow Matching
原始损失函数
设
条件速度场
具体而言,给定某特定样本
证明:我们只需要证明基于上面的方式定义的
熟悉 score function 的读者应该知道,对 score function 也有类似的关系式:
CFM 损失函数
基于条件概率路径和条件速度场,原始的 flow matching 损失函数可以改写做如下 conditional flow matching 损失函数:
证明:将
现在,只要我们设计一个可解的条件速度场
条件速度场的设计
一般形式
我们从条件概率路径
- 当
时,设 ,即与 无关的标准高斯分布; - 当
时,设 ,其中标准差 充分小,即大约是在 处的确定性分布。
直接找到对应上述条件概率路径的条件速度场是比较困难的,但是我们可以找到对应的以
可以验证上述形式的条件流
有了条件流,就可以找到对应的条件速度场
代入
如果只是为了训练,我们没有必要把条件速度场显式地解出来,但是这是可以做到的:
当
特殊情形:经典扩散
考虑 VE diffusion:
特殊情形:最优传输
相比经典的扩散路径,一个更合适的路径是将
需要强调的是,“直线”路径只在给定
总结
Flow matching 看起来蛮复杂,但是仔细想想,其实就是给已有的扩散模型框架套了个流体力学的壳,用新的体系重新说了一遍。比如条件速度场其实就是条件 score function 的线性变换,而 flow matching 的核心数学技巧也都是借用自 denoising score matching. 文章的亮点在于:
- 不要求源分布必须是高斯分布,从而可以用来学习两个分布之间的转换。
- 提出了基于最优传输的新的扩散路径,声称比传统的 VP 和 VE 路径更适合采样。
参考资料
- Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 (2022). ↩︎
- Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems 31 (2018). ↩︎
- Tor Fjelde, Emile Mathieu, Vincent Dutordoir. AN INTRODUCTION TO FLOW MATCHING. https://mlg.eng.cam.ac.uk/blog/2024/01/20/flow-matching.html ↩︎