FLOW MATCHING FOR GENERATIVE MODELING 阅读笔记

论文提出了一种新的生成模型。论文的目的是给定一个目标分布，有目标分布的一定量的样本，但是不知道目标分布的概率密度函数，学习一个模型能生成服从目标分布的新样本。
Flow Matching (FM)是一种训练连续标准化流Continuous Normalizing Flow (CNF)的方法。
FM是一种通用的方法。FM可以用于训练扩散路径，用FM训练扩散路径更稳定。FM也可以用于训练其他路径，一个例子是训练最优传输(OT)位移插值定义的条件概率路径，这些路径比扩散路径更有效，提供更快的训练和采样，从而获得更好的泛化效果。

核心的思想是把无条件估计问题的转换为有条件的问题的来学习。作者说是从denoised score matching得到的启发：

We first show that we can construct such target vector fields through per-example (i.e., conditional) formulations. Then, inspired by denoising score matching, we show that a per-example training objective, termed Conditional Flow Matching (CFM), provides equivalent gradients and does not require explicit knowledge of the intractable target vector field.

连续标准化流

数据点$x\in {\mathbb{R}}^d$，时变概率密度路径$p:[0,1]\times {\mathbb{R}}^d\to {\mathbb{R}}_{>0}$，时变向量场$v_t:[0,1]\times {\mathbb{R}}^d\to {\mathbb{R}}^d$。
流flow把一个分布映射成另一个分布，可以通过常微分方程用$v_t$构建flow $\varphi :[0,1]\times {\mathbb{R}}^d\to {\mathbb{R}}^d$：
$$\begin{gathered} \frac{d{\varphi }_t(x)}{dt}=v_t({\varphi }_t(x)) \\ {\varphi }_0(x)=x\text{\hspace{1em}(1)} \end{gathered}$$时变向量场可以用神经网络$v_t(x;\theta )$来建模，这样构建的flow${\varphi }_t$叫做连续标准化流（Continuous Normalizing Flow，CNF）。CNF通常用于把一个简单的分布$p_0$变成一个复杂的分布$p_1$，其符合push-forward方程：
$$p_t(x)=[{\varphi }_t{]}_{\star }{p}_0(x)=p_0({\varphi }_t^{-1}(x))\det [\frac{\partial {\varphi }_t^{-1}}{\partial x}(x)]$$我们的目标是采样服从复杂目标分布的样本，方法是首先随机采样服从简单分布的噪声样本$x\sim \mathcal{N}(0,I)$，然后使用ODE求解器在区间$t\in [0,1]$上使用训练得到的向量场$v_t$求解方程(1)得到服从目标分布的样本${\varphi }_1(x)$。所以主要的问题是如何学习$v_t(x;\theta )$。

Flow Matching（FM）

用$x_1$表示服从未知的目标分布$q(x_1)$的随机变量，我们不知道$q(x_1)$的密度函数，但可以获得服从$q(x_1)$的样本。用$p_t$表示概率密度路径，$p_0$服从标准高斯分布，$p_1$近似$q$。
Flow Matching的训练目标是学习$v_t$，损失函数是$${\mathcal{L}}_{FM}(\theta )={\mathbb{E}}_{t,p_t(x)}\Vert v_t(x;\theta )-u_t(x){\Vert }^2$$流匹配的损失函数很简单，但在实践中没法使用，因为我们不知道如何定义合适的$p_t$和$u_t$。

Conditional Flow Matching（CFM）

为了解决上面的问题，考虑条件流匹配。条件流匹配的损失函数是$${\mathcal{L}}_{CFM}(\theta )={\mathbb{E}}_{t,q(x_1),p_t(x|x_1)}\Vert v_t(x;\theta )-u_t(x|x_1){\Vert }^2$$与流匹配的目标不同，条件流匹配的目标允许我们轻松地对无偏估计进行采样，只要我们可以从 $p_t(x|x_1)$ 有效地采样并计算$u_t(x|x_1)$，这两者都可以很容易地完成，因为它们是对每个样本定义的。
论文中证明了优化CFM目标等价于优化FM目标(从期望的角度)。所以，剩下的问题是如何设计合适的条件概率路径$p_t(x|x_1)$和向量场$u_t(x|x_1)$。

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

上面的讨论是通用的，并没有规定条件概率路径和条件向量场的形式。为了简单，作者讨论的是高斯条件概率路径：
$$p_t(x|x_1)=\mathcal{N}(x|{\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 }$。
有无数的向量场可以产生给定的概率路径，这里作者讨论的是最简单的典型变换。
考虑条件flow：
$${\psi }_t(x)={\sigma }_t(x_1)x+{\mu }_t(x_1)$$对应的条件向量场可以通过求解方程得到，并有封闭解：
$$u_t(x|x_1)=\frac{{\sigma }_t^{\prime }(x_1)}{{\sigma }_t(x_1)}(x-{\mu }_t(x_1))+{\mu }_t^{\prime }(x_1)$$优化的损失函数是$${\mathcal{L}}_{CFM}(\theta )={\mathbb{E}}_{t,q(x_1),p(x_0)}\Vert v_t({\psi }_t(x_0);\theta )-u_t({\psi }_t(x_0)|x_1){\Vert }^2$$