Flow matching介绍

Flow matching

论文：FLOW MATCHING FOR GENERATIVE MODELING

参考blog：Flow Matching For Generative Modeling-CSDN博客

深入解析Flow Matching技术 - 知乎 (zhihu.com)

1.概率密度的变量变换定理

给定一个随机变量$z$及其概率密度函数$\pi (z)$，通过一个一对一的映射函数$f$构造一个新的随机变量$x=f(z)$，如果存在逆函数$z=f^{-1}(x)$，那么我们可以求出新变量$x$的概率密度函数：

定理的证明过程借助了概率密度函数的积分不变性，对于很小的区间$dz$，其概率$P(z-0.5dz=<Z<=z+0.5dz)=\pi (z)dz$，将区间$dz$变换到新区间$dx$时，应该有$P(x-0.5dx<=X<=x+0.5dx)=p(x)dx=\pi (z)dz$，即概率不变。

将定理从单随机变量推广到多随机变量，其中det表示行列式，J是雅可比矩阵。

2.流模型

流模型的思想就是把一个简单的分布经过若干次的可逆的变量变换之后，得到复杂的目标数据分布。（这就和概率密度变量变换定理有了联系，因为通过概率密度的可逆变量变换$f$，可以把简单分布$\pi (z)$变换成复杂分布$p(x)$，本质上流模型就是多个可逆变换的叠加作用）。

流模型通过最大化对数似然，来优化生成器$G$：

$G^{\ast }=\arg {\max }_G{\sum }_{i=1}^m\log P_G(x_i)$，

而根据变量变换定理，有：

$$\begin{gathered} P_G(x_i)=\pi (z_i)|\det (J_{G^{-1}})|, \\ {z}_i=G^{-1}(x_i),x_i=G(z_i) \end{gathered}$$

对数似然：

$log{P}_G(x_i)=log\pi (G^{-1}(x_i))+log|det(J_{G^{-1}})|$

要训练一个好的生成器我们只需要训练一个/多个网络，完成从分布$\pi (z)$到$P_{data}(x)$的变换就可以了。关键在于如何求$G^{-1}$以及计算行列式$|det(J_{G^{-1}})|$，这就是流模型的研究。

在实际中，由于可逆神经网络存在数学上的诸多限制，其单个网络的表达能力有限，我们一般需要堆叠多层网络来得到一个生成器，这也是 “流模型” 这个名称的由来。当然在这里不赘述，更多的可以看参考blog，因为不是我们主要讨论的。

3.CNF 连续归一化流（Continuous Normalizing Flow）

上面提到的堆叠多个网络，实际上就是用离散的有限个变量变换来完成简单分布到复杂数据分布（分布1->分布2->…->分布n,分布之间的变化是离散的，突变的）。而连续归一化流就是扩展为连续的情形。在连续时间内，可逆的变量变换映射函数$f$随着时间也不断变化，随着时间的推移，简单分布就像水流一样流向复杂数据分布（分布1->分布n，分布之间的变换是连续的，光滑的）。当离散的变换足够多，每次变换的差距足够小，就可以近似看成是连续的分布了。此话帮助理解。

首先我们来熟悉一些基础概念：

1.假设数据$x$是$d$维的，处于$d$维空间上，$x=(x^1,x^2,...,x^d)\in R^d$ (对于图片来说$d=H\ast W$)

2.概率密度路径p，$p:[0,1]\times R^d\to R_{>0}$，是一个关于时间的概率密度函数。注意：因为数据在不断随着时间进行变量变换，数据分布也随着时间在不断变化，所以其概率密度函数不是一成不变的，而是随着时间变化的，所以概率密度路径$p$将时间作为参数，当时间$t$确定时，$p_t(x)$就是$t$时刻数据的概率密度函数。而所有的连续$p_t(x)$一起构成了概率密度路径$p$。

3.关于时间的向量场v，$v:[0,1]\times R^d\to R^d$，向量场是表示某个时刻$t$，某个数据样本$x$的瞬时变化方向和速度。

为了方便理解，考虑1维空间上的数据，并且令$v=-x$。考虑数据点$x=2$（先明确一个概念，数据的值其实就是它在空间中的位置，位置的变化其实就是数据的变量变换，数形结合思想），那么$v(x)=-2$，说明它在向原点移动（数值上的变化趋势为减小），并且瞬时速度大小为2。推广到更为复杂的一般情况，向量场$v$不仅与$x$有关，还与$t$有关，例如$v=-x+sin(t)$，即使在向量场同一位置上，在不同的时间也会给数据点的变换带来不同的瞬时速度，增强了表达能力和复杂性。最后将1维空间推广为$d$维空间上的真实数据，就比较好理解了。

4.流$\varphi$: $\varphi :[0,1]\times R^d\to R^d$，我们先考虑这样一个问题，假如我已经知道了向量场$v$任意时刻$t$下任意位置$x$的瞬时速度，那么在时刻$t=0$时，初始位置为$x=x_0$的数据在向量场内移动，在时刻$t=t_1$时，他的位置$x_1$应该是多少呢？流$\varphi$就被定义成解决这个问题的函数。$x_1={\varphi }_{t_1}(x_0)$

特别的，当给定一个时刻$t$时，${\varphi }_t$其实就是一个变量变换的映射函数，将时刻0的任意数据变换到时刻$t$，也就是将分布$p_0$变换成分布$p_t$。当给定一个数据$x$时，$\varphi (x)$其实表示了单个数据$x$随时间$t$的移动轨迹。

显然，向量场$v$就已经很复杂了，流$\varphi$几乎很难表示出来。但是向量场$v$和流$\varphi$存在对应关系，

5.流和向量场之间的关系可以用ODE来表示：
$$\begin{gathered} \frac{d}{dt}{\varphi }_t(x)=v_t({\varphi }_t(x)) \\ {\varphi }_0(x)=x\text{\hspace{1em}(1,2)} \end{gathered}$$
对于初始时刻的数据点$x$，它在$t$时刻的位置即为${\varphi }_t(x)$，左边其实就是时刻$t$附近很小的时间段$dt$内，数据点${\varphi }_t(x)$的位移除以时间，右边就是数据点${\varphi }_t(x)$瞬时速度。在$dt$无穷小时，他们应该相等。**从该ODE可以看出，一个向量场$v$定义了一个流$\phi $。**因为知道了所有位置所有时刻的瞬时速度情况，任意数据任意时刻的位置就确定了。

了解完概率密度路径，向量场和流之后，再来看CNF。CNF考虑在连续时间上，把简单的分布$p_0$变成复杂数据分布$p_1$，push forward equation如下
$$\begin{gathered} p_t=[{\varphi }_t{]}_{\ast }{p}_0 \\ [{\varphi }_t{]}_{\ast }{p}_0(x)=p_0({\varphi }_t^{-1}(x))\det [\frac{\partial {\varphi }_t^{-1}}{\partial x}(x)]\text{\hspace{1em}(3,4)} \end{gathered}$$
从之前流的概念，我们知道式子(3)就是在表达：特别的，当给定一个时刻$t$时，${\varphi }_t$其实就是一个变量变换的映射函数，将时刻0的任意数据变换到时刻$t$，也就是将分布$p_0$变换成分布$p_t$。那么$p_t$的概率密度公式是什么呢？应用变量变换定理，就得到(4)。

我们认为：一个向量场$v_t$生成了一个概率密度路径$p_t$，如果该向量场定义的流${\varphi }_t$满足公式(3)(4)。（逻辑链：向量场定义流，流控制分布变换，生成概率密度路径，当我们拥有一个$v_t$时，我们实际上也已经拥有了对应的流和概率密度路径，一切的一切我们都有了，并且这三者之间满足(1)(2)(3)(4)的内在关系）。

是不是所有的$v_t$都能如愿以偿的生成一个概率密度路径$p_t$呢？其实不是。如果$v_t$生成的流${\varphi }_t$不能满足(3)(4)就不行。

在什么情况下会不满足(3)(4)？因为(3)(4)的推导使用了变量变换定理，而变量变换定理需要满足概率密度函数的积分不变性，所以当流${\varphi }_t$对概率密度函数进行变换时，如果不能保证该性质，或者流变换不可逆，那么就不可行。

有没有什么准则可以帮助我们判断什么样的$v_t$是可以生成概率路径$p_t$？论文的附录B对此作出说明。

连续性方程1

$$\begin{gathered} \frac{d}{dt}{p}_t(x)+div(p_t(x)v_t(x))=0 \\ div=\sum _{i=1}^d\frac{\partial }{\partial x^i} \end{gathered}$$

对连续性方程在此不作推导，但是可以提供直观的感性理解。

等式的第一项是概率密度路径$p$对时间的微分，代表概率密度路径$p$随着时间的变化趋势以及瞬时的变化速度（类比一元函数的梯度），更确切的说是在$t$时刻，数据点$x$的概率密度值的变化趋势和变化速度，可能是增大，也可能是减小。

第二项是概率密度路径乘以向量场的散度。让我们先了解散度的概念。

散度概念

假设有一个向量场，空间中每个点都对应一个向量$v_t(x)$，表示该点的水流速度。对于一个曲面来说，曲面的通量就是单位时间内经过曲面的水体积。对曲面进行微分，切成很多个小面，将水流速度$v_t(x)$投影到该面的法线方向上，再乘以小面的面积就可以得到小面的通量，最后所有小面进行积分就是曲面的通量。对于一个闭合曲面来说，其通量大于0时，说明从总体来看水是从曲面内空间流向曲面外的空间，通量小于0时，说明水从曲面外流向曲面内。当该闭合曲面慢慢缩小，导致其闭合的体积趋于一个点时，该曲面的通量就变成了该点的散度。该点的散度就可以理解当前时刻单位时间内，从该点是流出水（大于0），还是流入水（小于0），（建议搜索一下散度）。（散度的定义式在此也不赘述，可以自行搜索，nabla算子）

连续性方程2

回到第二项，考虑一个多维空间内，每个位置$x$都有一个值$p_t(x)$，代表数据可能出现在该位置的概率密度。$p_t(x)v_t(x)$就是在该多维空间的概率密度函数$p_t(x)$上叠加了一个向量场$v_t(x)$，让概率密度随时间变动，是整体数据分布的流动。那么散度就表示对于位置$x$，是流入概率密度还是流出概率密度，其绝对值即为流入/流出的速度。

• 如果概率密度在一个点增加（第一项大于零），那么散度项必须反映出有概率密度流入该点（第二项小于零）。
• 如果概率密度在一个点减少（第一项小于零），那么散度项反映出有概率密度流出该点（第二项大于零）。

这也是连续性方程的物理意义：在一个封闭系统中，概率密度的局部变化是由流入或流出该点的概率密度引起的。

OK,解释完了连续性方程为什么成立之后，我们就知道假如$v_t$想要生成$p_t$，必须满足该连续性方程，保证概率密度的局部变化是由流入或流出该点的概率密度引起，而不是凭空多出来的概率密度（保证概率的性质）。

4.Flow Matching

设$x_1$表示一个按照某个未知数据分布$q(x_1)$进行分布的随机变量。我们假设我们只能获取来自$q(x_1)$的数据样本，但无法直接访问其密度函数。此外，我们令$p_t$为一个目标概率路径，使得初始分布$p_0$是一个简单的分布，例如标准正态分布，并且$p_1$的分布与$q(x_1)$近似相等。流匹配目标（Flow Matching Objective）旨在匹配这个目标概率路径，这将允许我们从$p_0$流动到$p_1$。

$t\sim \mathcal{U}[0,1]x\sim p_t(x)$，$v_t(x)$是向量场预测模型，$u_t(x)$是决定概率密度路径$p_t$的向量场。我们要匹配概率密度路径，实际上匹配向量场就可以了。不过$L_{FM}(\theta )$本身无法计算，因为:

1.我们不知道什么样的$p_t(x)，u_t(x)$才是合适的（可以把$p_0$分布变为近似$q(x_1)$的分布称为合适的）

2.有无穷多个合适的$p_t(x)，u_t(x)$都合适

3.已知一个$p_t(x)$，我们无法求得$u_t(x)$的闭式解

我们应该如何构建$p_t(x),u_t(x)$? 论文给出了回答：借助条件概率密度路径和条件向量场构建。

给定一个数据样本$x_1$，我们定义条件概率密度路径$p_t(x|x_1)$：

• 满足$p_0(x|x_1)=p(x)$，也就是$p_0(x)$和样本$x_1$无关，是一个简单分布。
• 满足$p_1(x|x_1)=\mathcal{N}(x|x_1,{\sigma }^{2}I)$，是一个在$x=x_1$附近的分布，也就是说在$t=1$时要大致符合数据分布，即$p_1(x)=q(x)$。

边缘概率密度路径$p_t(x)$可以被很好地表示：
$$p_t(x)=\int p_t(x|x_1)q(x_1)d{x}_1\text{\hspace{1em}(6)}$$
特别地，当$t=1$时，有：
$$p_1(x)=\int p_1(x|x_1)q(x_1)d{x}_1\approx q(x)\text{\hspace{1em}(7)}$$
类似地，我们通过对条件向量场进行边缘化，定义边缘向量场：
$$u_t(x)=\int u_t(x|x_1)\frac{p_t(x|x_1)q(x_1)}{p_t(x)}d{x}_1\text{\hspace{1em}(8)}$$
其中，$u_t(|x_1)$是生成$p_t(|x_1)$的条件向量场。

这里有一个小问题，从$u_t(|x_1)$推导的$u_t(x)$的定义是否正确？即$u_t(x)$是否真的能生成$p_t(x)$?在附录A中给出了证明，只要它们满足连续性方程就可以了。具体过程论文也讲述的比较清楚，在此不赘述。

定理1：给定$p_t(x|x_1)$以及$u_t(x|x_1)$，对于任意数据分布$q(x_1)$，(8)定义的$u_t(x)$可以生成$(6)$定义的$p_t(x)$。

5.Conditional Flow Matching

问题依然没有解决，由于(6)(8)的定义内存在难以处理的积分操作，我们仍然不能计算${\mathcal{L}}_{FM}(\theta )$，但是研究人员发现了一个更为简单的优化目标，可以达到相同的最优解。

$t\sim \mathcal{U}[0,1],x_1\sim q(x_1),x\sim p_t(x|x_1)$，带有条件的概率密度和向量场都能很轻松地实现计算。

定理2：假设所有的$p_t(x)>0$，那么$L_{FM}$和$L_{CFM}$之多相差一个与$v_t(x)$参数$\theta$不相关的常数项，即它们对$\theta$进行求导的结果是一样的，${\nabla }_{\theta }{L}_{FM}(\theta )={\nabla }_{\theta }{L}_{CFM}(\theta )$

定理证明在附录A有详细解释。

分别证明前两项对应相等，第三项都是与$\theta$无关的常数项。

至此，我们把损失函数优化成了$L_{CFM}$。我们只要设计合理的条件概率路径，以及条件向量场就可以计算了。

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

$$\begin{gathered} p_t(x|x_1)=\mathcal{N}({\mu }_t(x_1),{\sigma }_t(x_1{)}^{2}I) \\ {\mu }_0(x_1)=0,{\sigma }_0(x_1)=1 \\ {\mu }_1(x_1)=x_1,{\sigma }_1(x_1)={\sigma }_{min}\text{\hspace{1em}(10)} \end{gathered}$$

我们定义一个如上的高斯条件概率路径，满足数据分布从简单分布流动至数据分布。考虑如下对应流。
$${\psi }_t(x)={\sigma }_t(x_1)x+{\mu }_t(x_1),x\sim \mathcal{N}(0,I)\text{\hspace{1em}(11)}$$
流本质上是一个变量变换映射函数，对单个数据点来说，流${\psi }_t(x)$把$x$变成${\psi }_t(x)$，对整个分布来说，把$p_0(x|x_1)=p_0(x)$变成了$p_t(x|x_1)$这个分布。把这种关系带入(4)，则有：

那么产生该流的向量场$u_t(x|x_1)$会是什么？看(1)，流和向量场的ODE关系：

观察原来的损失函数

我们改变一下抽样方法，我们原来是从$p_t(x|x_1)$中直接采样$x$，现在我们改为先从标准高斯$p_0(x_0)$中采样一个$x_0$，然后对$x_0$进行流的变量映射得到$x={\psi }_t(x_0)={\sigma }_t(x_1)x_0+{\mu }_t(x_1)$，这两种采样方法的$x$概率密度函数是一样的，可以相互等价。所以，有：

${\mathcal{L}}_{CFM}(\theta )={\mathbb{E}}_{t,q(x_1),p_0(x_0)}||v_t({\psi }_t(x_0))-{\mu }_t({\psi }_t(x_0)|x_1)|{|}^2$

将(13)带入可得：

我们可以求解${\mu }_t(x|x_1)$的闭式解：
KaTeX parse error: Undefined control sequence: \part at position 51: …psi_t(x)=\frac{\̲p̲a̲r̲t̲ ̲\sigma_t(x_1)}{…
进行换元，令$a={\psi }_t(x)$，则有$a={\sigma }_t(x_1)x+{\mu }_t(x_1)$，$x=(a-{\mu }_t(x_1))/{\sigma }_t(x_1)$，代入有
$${\mu }_t(a|x_1)=\frac{{\sigma }_t^{\prime }(x_1)}{{\sigma }_t(x_1)}(a-{\mu }_t(x_1))+{\mu }_t^{\prime }(x_1)$$
再把a还原成x得：

该向量场生成高斯路径$p_t(x|x_1)$。

定理3：假如存在(10)定义的条件高斯路径$p_t(x|x_1)$，以及对应的流变量映射(11)${\psi }_t(x)$，则存在一个独特的向量场(15)${\mu }_t(x|x_1)$能生成该高斯路径。

7.实例化一个$p_t(x),{\mu }_t(x)$

对于VF和VP两种扩散模型在flow matching视角下的描述，看论文原文。这里直接讲论文使用的新方法最优传输路径。
$$p_t(x)=\mathcal{N}(t{x}_1,[1-(1-{\sigma }_{min})t{]}^{2}I)$$

允许均值和方差随时间进行线性变化是简单且直观的路径，实际上也是两个高斯分布之间的最优映射。