引言
在上一章中，我们介绍了流模型。流模型通过模拟速度场对应的ODE，来得到参数的最大似然估计。
然而在实际优化过程中，计算这个最大似然以及其梯度需要非常精准的ODE模拟，只有完全正确的模拟结果才能得到无偏的梯度估计。
下面要介绍的流匹配，则是一种无需模拟的方法。
同样的，以下内容完全来自于Flow Matching Guide and Code。再次感谢FAIR的大佬们。
一、概览
给定一个初始分布 $p$ 和目标分布 $q$，流匹配的终极目的是训练一个流模型 $u^\theta_t$，使得 $u_t^\theta$ 能够生成概率路径 $p_t$，且 $p_0=p,p_1=q$。
因此，流匹配主要有以下几步：
- 确定一个已知的初始分布 $p$ 和一个未知的目标分布 $q$。
- 构建一条概率路径 $p_t$，满足 $p_0=p,p_1=q$。
- 学习一个速度场 $u^\theta_t$，使其能够生成 $p_t$。
- 从学到的模型中采样，生成一个目标分布中的样本。
学习速度场时，我们一般最小化下面的回归损失：
\[\begin{equation} \mathcal{L}_{FM}(\theta)=\mathbb{E}_{X_t\sim p_t}D(u_t(X_t),u^\theta_t(X_t)) \end{equation}\]其中，$D:\mathbb{R}^d\times\mathbb{R}^d\mapsto\mathbb{R}$ 是一个定义在 $\mathbb{R}^d$ 上的度量，常用的如L2范数。
这里有两个问题仍未解决：如何构建概率路径 $p_t$、如何构建真实的速度场 $u_t$。下面我们逐个说明。
二、构建概率路径
流匹配使用了一种【条件化】的策略，极大程度上简化了概率路径以及其对应向量场的构建难度。
具体来说，当我们只考虑目标分布上的某一个样本 $X_1=x_1$ ，我们就可以生成一些【条件概率路径】 $p_{t\mid 1}(x\mid x_1)$。
这些条件概率都需要满足下面的约束：
\[\begin{equation} \begin{aligned} p_{0\mid 1}(x\mid x_1)&=p(x)q(x_1)\\ p_{1\mid 1}(x\mid x_1)&=\delta(x-x_1) \end{aligned} \end{equation}\]其中 $\delta$ 是delta函数。
进一步我们可以通过对所有条件概率路径求期望得到【边缘概率路径】：
\[\begin{equation} p_t(x)=\int p_{t\mid 1}(x\mid x_1)q(x_1)\mathrm{d}x_1 \end{equation}\]这个边缘概率路径满足：$p_0=p,p_1=q$。
正如在第一章中提到的，一种最常见的构建条件概率路径的方法是使用如下高斯分布：
\[\begin{equation} p_{t\mid 1}(x\mid x_1)=\mathcal{N}(x\mid tx_1,(1-t)^2I) \end{equation}\]三、构建速度场
得到边缘概率路径 $p_t$ 后，我们可以构建出能够生成 $p_t$ 的一个速度场 $u_t$。
由于 $p_t$ 是由多个条件概率路径 $p_{t\mid 1}(x\mid x_1)$ 组合而成，因此 $u_t$ 也是由多个【条件速度场】$u_t(x\mid x_1)$ 组合而成，且满足：$u_t(x\mid x_1)$ 能够生成 $p_t(x\mid x_1)$。
最终，我们得到的速度场为：
\[\begin{equation} u_t(x)=\int u_{t}(x\mid x_1)p_{1\mid t}(x_1\mid x)\mathrm{d}x_1 \end{equation}\]注意这里用的是后验概率 $p_{1\mid t}(x_1\mid x)$，表示当前样本 $x$ 能够生成目标样本 $x_1$ 的概率。由贝叶斯公式，这个后验概率可以用下面的公式来计算：
\[\begin{equation} p_{1\mid t}(x_1\mid x)=\frac{p_{t\mid 1}(x\mid x_1)q(x_1)}{p_t(x)} \end{equation}\]3.1. 合法性证明
下面我们将说明，在引入一些假设后，公式$(4)$得到的速度场能够生成公式$(2)$中的概率路径。
我们所用到的工具是在第二章中介绍到的【质量守恒定律】。质量守恒定律是说，给定一个满足局部利普希茨性质的可积速度场 $u_t$ 以及一条概率路径 $p_t$，$u_t$ 能够生成 $p_t$ 当且仅当满足下面的公式$(6)$：
\[\begin{equation} \frac{\partial}{\partial t}p_t(x)+\textrm{div}(p_tu_t)(x)=0 \end{equation}\]其中，$\textrm{div}(v)(x)=\text{tr}[\partial_{x}v(x)]\in\mathbb{R}^{d\times d}$ 称为散度矩阵。公式$(6)$也被称为连续性方程。
不失一般性地，我们可以把上面引入的目标样本 $x_1$ 推广为任意的随机向量 $z\in\mathbb{R}^m$，则对应的概率路径变为：
\[\begin{equation} p_t(x)=\int p_{t\mid Z}(x\mid z)p_Z(z)\mathrm{d}z \end{equation}\]生成的速度场变为：
\[\begin{equation} \begin{aligned} u_t(x)&=\int u_{t}(x\mid z)p_{Z\mid t}(z\mid x)\mathrm{d}z \end{aligned} \end{equation}\]且满足：$u_t(x\mid z)$ 能够生成 $p_{t\mid Z}(x\mid z)$。
为了说明我们的结论，首先需要引入下面的假设：
- $p_{t\mid Z}(x\mid z)\in C^1([0,1)\times \mathbb{R}^d)$，
- $u_{t\mid Z}(x\mid z)\in C^1([0,1)\times \mathbb{R}^d,\mathbb{R}^d)$，
- $p_Z$ 具有有界支撑集，即满足 $p_Z(z)=0$ 的样本位于某些无界集合中，
- $\forall x\in\mathbb{R}^d,\forall t\in[0,1),p_t(x)>0$。
这四点假设并不难满足。首先，大多数常用的概率分布（如高斯分布）都是无限次平滑的，远远超出 $C^1$ 的要求。其次，真实世界的数据（如图片、声音、文本）几乎总是有界的。最后，我们可以让条件 $z$ 满足 $p_Z(z)>0$ 且 $p_{t\mid Z}(x\mid z)>0$，就能满足第四条假设。
有了这些准备，我们就可以证明下面的定理3（详细证明见附录）。
Theorem 3: (Marginalization Trick) 在满足上述假设的情况下，如果 $u_t(x\mid z)$ 是【条件可积】的，且能够生成条件概率路径 $p_{t\mid Z}(x\mid z)$，则 $u_t$ 能够生成 $p_t$。
这里条件可积是指：
\[\begin{equation} \int_0^1\int\int \|u_t(x\mid z)\|p_{t\mid Z}(x\mid z)p_Z(z)\mathrm{d}z\mathrm{d}x\mathrm{d}t<\infty \end{equation}\]四、Flow Matching损失函数
至此，我们证明了存在一个速度场 $u_t$ 能够生成从 $p$ 到 $q$ 的一条概率路径 $p_t$。
下面，我们就是要设计一个可微的损失函数，去学习一个速度场 $u_t^\theta$，使其尽可能接近真实的速度场 $u_t$。
然而，真实速度场 $u_t$ 非常难以计算。因为它需要对整个训练集中的目标样本 $x_1$ 进行积分（见公式$(4)$）。
为了求解这个问题，我们需要引入一类特殊的损失函数，称为Bregman散度。Bregman散度可以利用条件速度 $u_t(x\mid z)$ 对学习 $u_t^\theta$ 的梯度进行无偏的估计。
具体来说，Bregman散度衡量了两个向量 $u,v\in\mathbb{R}^d$之间的距离：
\[\begin{equation} D(u,v):=\Phi(u)-[\Phi(v)+\langle u-v,\nabla\Phi(v) \rangle] \end{equation}\]其中，$\Phi:\mathbb{R}^d\mapsto\mathbb{R}$ 是一个定义在凸集 $\Omega\subset\mathbb{R}^d$ 上的严格凸函数。
注意到，$[\Phi(v)+\langle u-v,\nabla\Phi(v) \rangle]$ 实际上就是在求 $\Phi(x)$ 在 $x=v$ 上的一阶Taylor近似，因此Bregman散度实际上度量的是函数 $\Phi$ 与其一阶Taylor估计的距离。
当我们选择 $\Phi(u)=|u|^2$时，Bregman散度就变为我们熟知的欧式距离：$D(u,v)=|u-v|^2$。
Bregman散度一个非常重要的性质是：Bregman散度的梯度满足【仿射不变性】（证明见附录）。具体来说：
\[\begin{equation} \forall a+b=1,\nabla_vD(au_1+bu_2,v) =a\nabla_vD(u_1,v)+b\nabla_vD(u_2,v) \end{equation}\]有了这种仿射不变性，我们就可以交换梯度和期望的位置：
\[\begin{equation} \forall Y\in\mathbb{R}^d,\nabla_vD(\mathbb{E}[Y],v)=\mathbb{E}[\nabla_vD(Y,v)] \end{equation}\]Flow Matching的损失函数中使用了Bregman散度，用于度量 $u_t$ 和 $u^\theta_t$ 的距离：
\[\begin{equation} \mathcal{L}_{FM}(\theta) =\mathbb{E}_{t,x_t\sim p_t}D(u_t(x_t),u^\theta_t(x_t)) \end{equation}\]正如上面提到的， $u_t$ 是难解的。因此我们考虑下面的条件损失：
\[\begin{equation} \mathcal{L}_{CFM}(\theta) =\mathbb{E}_{t,z,x_t\sim p_{t\mid Z}}D(u_t(x_t\mid z),u^\theta_t(x_t)) \end{equation}\]可以证明，公式$(13)$和公式$(14)$有着相同的梯度。因此二者对于优化问题来说是等价的。证明见附录。
标准的CFM损失中，时间一般是从均匀分布中采样的，即 $t\sim U[0,1]$。
但也有研究证明，在大规模的图像生成任务中，从其他的分布 $\omega$ 中采样时间往往会得到更好的效果。此时CFM损失变为：
\[\begin{equation} \begin{aligned} \mathcal{L}_{CFM}(\theta) &=\mathbb{E}_{t\sim \omega,z,x_t\sim p_{t\mid Z}}D(u_t(x_t\mid z),u^\theta_t(x_t))\\ &=\mathbb{E}_{t\sim U[0,1],z,x_t\sim p_{t\mid Z}}\omega(t)D(u_t(x_t\mid z),u^\theta_t(x_t))\\ \end{aligned} \end{equation}\]比如说，在Stable Diffusion 3中，使用双峰分布 (Bimodal Distribution)来采样时间，显著增加接近0或1的时间步的采样概率。
\[\begin{equation} \begin{aligned} w(t)&=p\cdot Beta(t;\alpha_1,\beta_1)+(1-p)\cdot Beta(t;\alpha_2,\beta_2)\\ Beta(t;\alpha,\beta)&=\frac{t^{\alpha-1}\cdot(1-t)^{\beta-1}}{B(\alpha,\beta)}\\ B(\alpha,\beta)&=\int_0^1x^{\alpha-1}\cdot(1-x)^{\beta-1}\mathrm{d}x \end{aligned} \end{equation}\]附录
Apx1. 定理3的证明
首先我们证明 $u_t$ 和 $p_t$ 满足连续性方程。
我们有：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}p_t(x) &=\int\frac{\mathrm{d}}{\mathrm{d}t}p_{t\mid Z}(x\mid z)p_Z(z)\mathrm{d}z\\ &=-\int\textrm{div}_x\left[ u_t(x\mid z)p_{t\mid Z}(x\mid z) \right]p_Z(z)\mathrm{d}z\\ &=-\textrm{div}_x\left[\int u_t(x\mid z)p_{t\mid Z}(x\mid z) p_Z(z)\mathrm{d}z\right]\\ &=-\textrm{div}_x\left[\left(\int \frac{u_t(x\mid z)p_{t\mid Z}(x\mid z)p_Z(z)}{p_t(x)} \mathrm{d}z\right)p_t(x)\right]\\ &=-\textrm{div}_x\left[\left(\int u_t(x\mid z)p_{Z\mid t}(z\mid x) \mathrm{d}z\right)p_t(x)\right]\\ &=-\textrm{div}_x\left[u_t(x)p_t(x) \right]\\ \end{aligned} \end{equation}\]其中第一、三行交换了积分和求导（散度）的顺序。这是可行的，因为 $p_{t\mid Z}(x\mid z)$ 和 $u_t(x\mid z)$ 都是 $C^1$ 函数，且$p_Z$具有有界的支撑集。第二行成立是因为 $u_t(x\mid z)$ 能够生成 $p_{t\mid Z}(x\mid z)$。第四行成立是因为 $p_t(x)>0$。第五行成立是因为贝叶斯公式。第六行成立则是代入了公式$(4)$。
然后我们再证明 $u_t$ 是一个满足局部利普希茨性质的可积速度场。
由于 $C^1$ 函数满足局部利普希茨条件，因此我们只需说明 $u_t$ 是一个 $C^1$ 函数。这是成立的，因为 $p_{t\mid Z}(x\mid z)$ 和 $u_t(x\mid z)$ 都是 $C^1$ 函数且 $p_t(x)>0$。
$u_t$ 的可积性是因为 $u_t(x\mid z)$ 的条件可积性：
\[\begin{equation} \int_0^1\int \|u_t(x)\|p_t(x)\mathrm{d}x\mathrm{d}t \le \int_0^1\int\int \|u_t(x\mid z)\|p_{t\mid Z}(x\mid z)p_Z(z)\mathrm{d}z\mathrm{d}x\mathrm{d}t<\infty \end{equation}\]第一个不等号的成立是根据Jensen不等式：
\[\begin{equation} \begin{aligned} \|u_t(x)\| &=\left\| \int u_{t}(x\mid z)p_{Z\mid t}(z\mid x)\mathrm{d}z \right\|\\ &\le \int \left\| u_{t}(x\mid z)\right\| p_{Z\mid t}(z\mid x)\mathrm{d}z\\ &=\int \left\| u_{t}(x\mid z)\right\| \frac{p_{t\mid Z}(x\mid z)p_Z(z)}{p_t(x)} \mathrm{d}z\\ \|u_t(x)\|p_t(x)&\le \int \left\| u_{t}(x\mid z)\right\| p_{t\mid Z}(x\mid z)p_Z(z)\mathrm{d}z \end{aligned} \end{equation}\]因此，我们证明了 $u_t$ 是一个满足局部利普希茨性质的可积速度场，且 $u_t$ 和 $p_t$ 满足连续性方程。
根据质量守恒定律， $u_t$ 能够生成 $p_t$。
Apx2. Bregman散度的梯度仿射不变性
下面我们证明公式 $(11)$。
首先我们推导Bregman散度的梯度。
\[\begin{equation} \begin{aligned} \nabla_v\langle u-v,\nabla\Phi(v) \rangle &=\nabla_v[(u-v)^T\nabla\Phi(v)]\\ &=[\nabla_v(u-v)]^T\nabla\Phi(v) + [\nabla_v\nabla\Phi(v)]^T(u-v)\\ &=-\nabla\Phi(v)+(\nabla^2\Phi(v))^T(u-v)\\ &=-\nabla\Phi(v)+\nabla^2\Phi(v)(u-v)\\ \end{aligned} \end{equation}\]最后一步利用了Hessian矩阵的对称性。
因此：
\[\begin{equation} \begin{aligned} \nabla_vD(u,v) &=0-\nabla\Phi(v)-\nabla_v\langle u-v,\nabla\Phi(v) \rangle\\ &=-\nabla^2\Phi(v)(u-v)\\ &=\nabla^2\Phi(v)(v-u) \end{aligned} \end{equation}\]公式 $(11)$的等式左边
\[\begin{equation} \begin{aligned} \text{LHS} &=\nabla_vD(au_1+bu_2,v)\\ &=\nabla^2\Phi(v)(v-au_1-bu_2) \end{aligned} \end{equation}\]公式 $(11)$的等式右边
\[\begin{equation} \begin{aligned} \text{RHS} &=a\nabla_vD(u_1,v)+b\nabla_vD(u_2,v)\\ &=a[\nabla^2\Phi(v)(v-u_1)]+b[\nabla^2\Phi(v)(v-u_2)]\\ &=\nabla^2\Phi(v)[a(v-u_1)+b(v-u_2)]\\ &=\nabla^2\Phi(v)[(a+b)v-au_1-bu_2]\\ &=\nabla^2\Phi(v)[v-au_1-bu_2]\\ &=\text{LHS} \end{aligned} \end{equation}\]得证。
Apx3. FM Loss和CFM Loss的等价性
下面我们证明公式$(13)$和公式$(14)$有着相同的梯度。
\[\begin{equation} \begin{aligned} \nabla_\theta\mathcal{L}_{FM}(\theta) &=\nabla_\theta\mathbb{E}_{t,x_t\sim p_t}D(u_t(x_t),u^\theta_t(x_t))\\ &=\mathbb{E}_{t,x_t\sim p_t}\nabla_\theta D(u_t(x_t),u^\theta_t(x_t))\\ &=\mathbb{E}_{t,x_t\sim p_t}\nabla_v D(u_t(x_t),u^\theta_t(x_t))\nabla_\theta u_t^\theta(x_t)\\ &=\mathbb{E}_{t,x_t\sim p_t}\nabla_v D( \mathbb{E}_{z\sim p_{Z\mid t}(\cdot\mid x_t)}[u_t(x_t\mid z)] ,u^\theta_t(x_t))\nabla_\theta u_t^\theta(x_t)\\ &=\mathbb{E}_{t,x_t\sim p_t,z\sim p_{Z\mid t}(\cdot\mid x_t)}\nabla_v D(u_t(x_t\mid z) ,u^\theta_t(x_t))\nabla_\theta u_t^\theta(x_t)\\ &=\mathbb{E}_{t,x_t\sim p_t,z\sim p_{Z\mid t}(\cdot\mid x_t)}\nabla_\theta D(u_t(x_t\mid z) ,u^\theta_t(x_t))\\ &=\mathbb{E}_{t,z,x_t\sim p_{t\mid Z}(\cdot\mid z)}\nabla_\theta D(u_t(x_t\mid z) ,u^\theta_t(x_t))\\ &=\nabla_\theta\mathbb{E}_{t,z,x_t\sim p_{t\mid Z}(\cdot\mid z)} D(u_t(x_t\mid z) ,u^\theta_t(x_t))\\ &=\nabla_\theta\mathcal{L}_{CFM}(\theta) \end{aligned} \end{equation}\]这里面使用到的证明技巧：
- 第三、第六行用到了求导链式法则，变换了求导变量
- 第四行应用了公式$(8)$
- 第五行应用了Bregman散度的梯度放射不变性，即公式$(11)$
- 第七行应用了贝叶斯公式
