🗽Flow Matching for Generative Modeling (Part 1)

type
status
date
slug
summary
tags
category
icon
password
论文作者在 2024 NeurIPS 上做了报告,我觉得跟着原作者思路梳理这篇文章并完整复现是最好的。
此外, 也为 Flow Matching 专门出了 blog。
在学的过程中还发现了 Lilian Weng 的 blog。
 
🗽
Denoising Diffusion Probabilistic Models (Part 1)
中曾为 Generative Model 建模。可以用一张图更清晰直观地描述
notion image
直接对想要的分布做采样是十分困难的,所以先对一个简单的分布(比如高斯分布)做采样,再通过变换得到采样结果。
这里的难点在于 怎么对应,或者说 怎么变换到 。通常我们不是直接从 source 映射到 target 的,而是生成一连串的样本,逐步从 source 变化到 target,每一个样本都在某种意义上比前一个样本更好。我们通常将时间 到时间 的转移用函数 来描述。我们前面讲过的 Diffusion Model 的转移是一个从 source 到 target 连续但不平滑的过程(不平滑源于随机噪声),而 Flow Matching 是一个平滑的过程。这些转移的过程是所谓的连续时间马尔科夫链,满足马尔科夫性质,即其未来的分布仅取决于其当前状态而不取决于其过去的分布。
生成模型算法的目标是学习 transition function,用一个网络帮助我们从 状态转移到 状态(报告里把 限定在了 范围内, 是 source , 是 target )。
notion image
边缘概率路径(marginal probability path)是指在变换过程中一系列特定 下的概率分布 (也就是边缘概率分布)串起来形成的路径。
notion image
直观来说,Flow Matching 的思路是拖动空间,给空间内各个点一个移动的速度,重塑空间但不丢失信息(双射)地让 变换到
这种对空间的扭动(Warping)本质是一个函数变换,写作数学形式是
这就像是由粒子组成的空间,Flow 是用粒子移动的局部速率 Velocity 来描述的,这样就避免了直接描述难以参数化的 Flow 。二者实际是等价的。显然, Velocity 是 Flow 的导数
从 Flow 到 Velocity 是相对简单的微分,从 Velocity 到 Flow (sample 时)就需要解微分方程(ODE)了。这里的关键点在于 Velocity 是线性的,而 Flow 是非线性的(其实不太懂……)。
我们定义如下术语:
Velocity generates if
在 Flow Matching 中,我们希望训练一个 Velocity ,该 Velocity 可以生成一个从 source 开始到 target 的边缘概率路径。
 
第一个问题是当我们有了一个 Velocity ,我们要怎么去对目标分布采样。
方法很简单。我们对简单的源分布做采样,然后让这个采样点根据 Velocity 到目标分布对应的采样点。从数学上来说,这是一个从 不断解常微分方程的过程。
notion image
 
第二个问题是:我们要如何训练处这个 Velocity 。
最简单的版本是直接假设 training loop 的每一轮迭代从 的变换是线性的、均匀的,这里的 是完全随机的。对 source 和 target 做线性组合计算出 ,输入到我们的网络中,我们尝试回归的是这样一个式子
 
notion image
一个简单的 for loop
为什么这样做有效
报告中有一段简单的讲解。感觉大致的意思像是从特殊情况( 集中在一点)推一般情况。
notion image
报告中给出了所谓 The Marginalization Trick
Theorem: The marginal velocity generates the marginal probability path.
 
 
到此,我们已经对 Flow Matching 这个模型有了一个大致的概念。接下来我们从论文中补充细节。

 

Introduction

… an efficient simulation-free approach to training CNF models, allowing the adoption of general probability paths to supervise CNF training.
 
论文中将 Velocity 描述为 “a target vector field that generates a desired probability path”,并称 Velocity 可以通过 per-example (i.e., conditional) formulations 来建立。由此构建出来的 Conditional Flow Matching (CFM) 拥有与原矢量场相同的梯度。
 

Preliminaries: Continues Normalize Flows

先定义几个关键数学描述
data points:
the probability density path:
the vector field:
flow:
其中,vector field 和 flow 间由前文提到的过的常微分方程联系
通过神经网络拟合 ,由此得到 Continues Normalize Flows (CNF) ,用于从源分布到目标分布到变形。该过程可以描述为
其中,运算符 定义为
CNF
Normalizing Flows(NF) 可以被看作是一连串的变量替换的迭代过程,每次替换都遵循概率密度函数的变量变换原则。
notion image
在传统的 NF 中,变换通常是通过一系列可逆的离散函数来定义的,而在 CNF 中,这种变换是连续的
这个式子的理解
前半部分 类似对 做拉伸变形,后半部分是一个雅各比矩阵,可以衡量并抵消前半部分变形对总概率的影响,类似归一化。
数学推导并不太复杂,在知乎和 Lilian Weng‘s Blog 有介绍。
 

Flow Matching

直观的,FM 尝试回归的是
与所有生成模型面对的问题一样,在实际生成前我们并不知道
论文的创新点是通过仅定义 per sample (wtf)来构建 这两个通常是未知的量。
 

Constructing , From Conditional Probability Paths and Vector Fields

一个朴素的想法是用很多个简单的概率路径拼出我们想要的概率路径。
 
对于一个给定的数据样本 ,我们希望 Conditional Probability Path 时刻呈现初始的简单分布 ,在 时刻高度集中到 附近。
我们用 来近似集中到 点的效果,其中 是很小的数。这样就得到了一条简单的概率路径。
 
将所有可能的 对应的条件概率路径拼起来的操作称为边际化 Marginalization(妙不可言)。对每个 乘以对应的概率权重再求积分,就能实现拼接
时,我们称 为 mixture distribution
如图是一个例子,两个数据样本 得到 conditional flows
notion image
 
进一步,我们可以类似地定义 marginal vector field 。一个直接的想法是直接对所有 加权平均,但显然这样没有考虑 本身服从的概率分布。正确的做法应该是根据 加权求平均,于是
根据贝叶斯公式(又来!):
于是
为什么这个式子是对的
也就是想证明通过条件向量场 适当加权确实可以得到正确的边际向量场
证明会使用 continuity function (对的,就是普物里那玩意儿)
具体证明可以看 Cambridge blog。
 
现在我们理一下各个量之间的关系
notion image
左侧 是我们期望的转变,这一转变用 Flow 来描述。根据上面的推导, 加权平均得到, 加权平均得到。
notion image
直接拟合 Flow 是困难的,我们引入了 Velocity 。前面提到, 可以通过 根据 加权平均,也就是这张图中间的路径。
总结来说,论文通过 的关系构建了条件向量场和边缘向量场间的桥梁。
The marginal vector field generates the marginal probability path.
Theorem 1. Given vector fields that generate conditional probability paths , for any distribution , the marginal vector field generates the marginal probability path , i.e. and satisfy the continuity equation.
感觉这一大段讲的不是太清楚 💦
 

Conditional Flow Matching

前面提到, FM 面对的问题是缺少先验知识 ,因为尽管存在许多概率路径使得变换后的概率分布与目标分布相同或接近,但通常无法得到一个的封闭形式的向量场 来生成期望的概率路径
对此,论文提出 CMF ,这是一种基于样本的方法,为每个样本独立地定义概率路径和向量场,而不是依赖一个全局的、封闭的向量场,然后通过适当的聚合方式(前文介绍的加权平均),将这些个体样本的概率路径和向量场聚合起来生成需要的
 
CMF 尝试回归的是
这样做的逻辑是 有相同的梯度,即
Theorem 2. Assuming that for all and , then, up to a constant independent of , and are equal. Hence, .
 

Conditional Probability Paths and Vector Fields

且不提上面一堆是不是完全理解了,我们的目标已经转变为了构建
条件概率路径有很多选择,论文给出的是 Gaussian conditional probability paths
条件概率路径应当满足如下两个边界条件:
  • 时,
  • 时, 近似聚焦在
 
一个概率路径可以由无数种向量场生成。为避免无散度项(divergence free component)增加额外计算量,论文使用最简单的向量场,即如下对标准高斯分布的变换
这是一个仿射变换, 推动最初的高斯噪声 变换到 ,用如下标记描述
求解如下 ODE 就能得到决定 conditional probability paths 的
由此,我们也得出了最终的损失函数
 
论文中的 Theorem 3 叙述了条件向量场的封闭解
Theorem 3. Let be a Gaussian probability path, and  its corresponding flow map. Then, the unique vector field that defines  has the form:
Consequently, generates the Gaussian path .
 

Special Instance of Gaussian Conditional Probability Paths

  • Example 1: Diffusion conditional VFs
论文揭示了扩散模型是 FM 的一个特例。
noise data 的 Variance Exploding (VE) 扩散路径可以描述为
其中,   是一个递增函数,满足   且 。由上式可以得到, 。带入 Theorem 3 就可以得到
data noise 的 Variance Preserving (VP) 扩散路径可以描述为
其中, 。由上式可以得到, 。带入 Theorem 3 就可以得到
tmd 看不懂
 
  • Example 2: Optimal Transport conditional VFs
将条件概率路径的均值和方差简单定义为与时间线性相关的函数
带入 Theorem 3 得到
 
总的来说, Diffusion 的轨迹是曲线,而 OT 的轨迹是直线,所以 OT 更高效。
notion image
notion image

 
实在是现在能力有限,感觉就搞懂了 60% 。这篇论文写的我真不好评价💦
非常详细的数学逻辑和讲解可以细看 Cambridge blog 和 Lilian Weng‘s blog 。
还有一篇文章对比了 Diffusion 和 Flow Matching ,看起来也很不错
希望 coding 部分逻辑会理的更顺一些。
 

Coding

 
上一篇
Denoising Diffusion Probabilistic Models (Part 2)
下一篇
Flow Matching for Generative Modeling (Part 2)
Loading...