论文阅读笔记——FLOW MATCHING FOR GENERATIVE MODELING

Flow Matching 论文
扩散模型：根据中心极限定理，对原始图像不断加高斯噪声，最终将原始信号破坏为近似的标准正态分布。这其中每一步都构造为条件高斯分布，形成离散的马尔科夫链。再通过逐步去噪得到原始图像。
Flow matching 采取直接将已知分布（如白噪声）转换为真实数据分布来生成数据，并且 Flow 是基于 Normalizing Flow，故而是可微双射。生成过程中变化的概率密度构成一个集合，称为概率密度路径 $p_t$ ，T 为路径长度。初始数据 $x_0\sim p_0(x_0)$，目标数据 $x_T\sim p_T(x_T)$。
从 $x_0$ 到 $x_T$ 的过程可以表示为：$x_T=\varphi (x_0)={\varphi }_T\circ \cdots \circ {\varphi }_{t+1}\circ {\varphi }_t\circ \cdots {\varphi }_1(x_0)$
且对中间任意时间步 $x_t$ 有：
$$\begin{array}{r}{x}_t={\varphi }_t(x_{t-1})\\ x_{t-1}={\varphi }_t^{-1}(x_t)\end{array}$$
根据概率密度函数的变量变换关系可得：（行列式为时刻 t 对应的流 ${\varphi }_t$ 的 Jacobian 行列式）
$$\begin{array}{rl}{p}_t(x_t)& =p_{t-1}(x_{t-1})\det \left[\frac{\partial x_{t-1}}{\partial x_t}\right]\\ & =p_{t-1}({\varphi }_t^{-1}(x_t))\det \left[\frac{\partial {\varphi }_t^{-1}}{\partial x_t}(x_t)\right](1)\end{array}$$
那么就可以从初始数据分布 $p_0$ 推导到目标数据分布 $p_T$。行列式的本质是空间缩放的度量，相当于每次变换时都对概率密度进行归一化，采用更简洁的前推方程为：
$$p_t=[{\varphi }_t{]}_{\ast }{p}_0$$

向量场建模

这可以通过 Neural Ordinary Differential Equations(NODE) 对 Jacobian 行列式中的常微分方程（ODE）建模，求出 ${\varphi }_t$。为了实现这一点，需要将离散的时间步 t = { t i } i = 1 T t=\{t_i\}_{i=1}^T t={ti​}i=1T​ 映射到连续时间变量 $t\in [0,1]$ ，这样将 $p_t$ 定义为连续时间和数据点的笛卡尔积：$p:[0,1]\times {\mathbb{R}}^d->{\mathbb{R}}_{>0}$ 且 $\int p_t(x)dx=1$，这就是 CNF 建模。
虽然 CNF 对 Flow 进行了建模，但同时也面临 Jacobian 行列式计算性能低下、训练时需要进行模拟、难以高效采样的问题。Flow Matching 通过回归概率路径的向量场，在训练时规避了复杂计算，并具有更高的采样性能，而且向量场基本上是完全精确的，与扩散模型使用的变分推断等近似方法相比在似然计算上更具优势。而重新审视 Flow，数据点在时间上的变换可以用 Flow 的梯度表示，构成了关于时间的向量场。
$$\begin{array}{rl}\frac{d}{dt}{\varphi }_t(x)& =v_t({\varphi }_t(x))(2)\\ {\varphi }_0(x)& =x\end{array}$$
对于给定向量场 ${\varphi }_t$ 满足式（1）时，可以说向量场 $v_t$ 生成概率密度路径 $p_t$，类比物理学的连续性方程，则有：
$$\frac{d}{dt}{p}_t+\nabla \cdot v_t{p}_t=0(3)$$
选取合适的 ODE 求解器，那么可以得到目标函数：
$${\mathcal{L}}_{FM}(\theta )={\mathbb{E}}_{t,p_t(x)}||v_t(x)-u_t(x)|{|}^2(4)$$

条件流匹配

尽管我们已经推导出了 ${\mathcal{L}}_{FM}$ 但是目标向量 $u_t$ 还是未知的，无法学习，可以通过用易于访问的混合分布来构造真实概率路径。通过目标数据样本 $x_1\sim q(x_1)$ 定义一个条件概率 $p_t(x|x_1)$，使得 $p_0(x|x_1)=p(x)$ 。用这个条件概率和真实分布 $q(x_1)$ 边缘化 $p_t$ 有：
$$p_t(x)=\int p_t(x|x_1)q(x_1)d{x}_1(5)$$
同样“边缘化”向量场 $u_t$，有：
$$u(x)=\int u_t(x|x_1)\frac{p_t(x|x_1)q(x_1)}{p_t(x)}(6)$$
该公式是连接条件向量场和边缘向量场的桥梁，论文证明只要条件向量场 $u_t(x|x_1)$ 能生成对应的条件概率路径 $p_t(x|x_1)$，对于任何分布 $q(x_1)$，上述定义的边缘向量场 $u_t(x)$ 能够生成对应的边缘概率路径 $p_t(x)$。并且证明了除与 $\theta$ 无关的常数外，${\mathcal{L}}_{CM}$ 和 ${\mathcal{L}}_{FM}$ 相等，即 ${\nabla }_{\theta }{\mathcal{L}}_{CM}(\theta )={\nabla }_{\theta }{\mathcal{L}}_{CFM}(\theta )$。
由此论文提出了基于条件概率路径 $p_t(x|x_1)$ 和条件向量场 $u_t(x|x_1)$ 的目标函数：
$${\mathcal{L}}_{\mathbf{CFM}}(\theta )={\mathbb{E}}_{t,q(x_1),p_t(x|x_1)}\Vert v_t(x)-u_t(x|x_1){\Vert }^2(7)$$
Conditional Flow Matching可以选择任意的条件概率路径，只要满足边界条件即可，这里针对一般高斯条件概率路径：$p_t(x|x_1)=\mathcal{N}(x|{\mu }_t(x_1),{\sigma }_t(x_1{)}^{2}I)$。

• 当 t=0 时，${\mu }_0(x_1)=0{\sigma }_0(x_1)=1$ 确保所有的条件概率路径都会收敛到相同的标准高斯分布。
• 当 t=1 时，${\mu }_1(x_1)=x_1{\sigma }_1(x_1)={\sigma }_{min}$ 以 $x_1$ 为中心的高斯分布。
  对于一个概率路径，存在无限多个向量场可以生成它，这里采用高斯分布的标准变换：
  $${\psi }_t(x)={\sigma }_t(x_1)x+{\mu }_t(x_1)(8)$$
  代入式（2）则有：
  $$\frac{d}{dt}{\psi }_t(x)=u_t({\psi }_t(x)|x_1)$$
  用 $x_0$ 重参数化 $p_t(x|x_1)$ 代入式（7）则有：
  $${\mathcal{L}}_{\mathbf{CFM}}(\theta )={\mathbb{E}}_{t,q(x_1),p_0(x_0)}\Vert v_t({\psi }_t(x_0))-\frac{d}{dt}{\psi }_t(x_0){\Vert }^2(9)$$
  根据式（8）可知 ${\varphi }_t$ 是可逆的仿射变换，故而可以得到条件向量场为：
  $$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)(10)$$

值得注意的是，这里选择高斯概率密度路径只是可选方式之一，实际上可以根据需要设计任何合理的路径，SD3 也是 Conditional Flow Match 的一种应用。

扩散模型

Variance Exploding（VE）

根据反向时间对称性，从噪声到数据的逆向过程中的条件 $p_t$ 为：
$$p_t(x)=\mathcal{N}(x|x_1,{\sigma }_{1-t}^{2}I)$$
且 ${\sigma }_0=0{\sigma }_1>>1$。因此 ${\mu }_t(x_1)=x_1$ 和 ${\sigma }_t(x_1)={\sigma }_{1-t}$ 代入式（10）：
$$u_t(x|x_1)=-\frac{{\sigma }_{1-t}^{{}^{\prime }}}{{\sigma }_{1-t}}(x-x_1)(11)$$
对于 VP-SDE （DDPM）的条件 $p_t$ 有：
$$p_t(x|x_1)=\mathcal{N}(x|{\alpha }_{1-t}{x}_1,(1-{\alpha }_{1-t}^2)I),\mathrm{where}{\alpha }_t=e^{-\frac{1}{2}T(t)},T(t)={\int }_0^t\beta (s)ds$$
其中 $\beta$ 是噪声尺度函数。因此：${\mu }_t(x_1)={\alpha }_{1-t}{x}_1$ 和 ${\sigma }_t(x_1)=\sqrt{1-{\alpha }_{1-t}^2}$ 代入式（10）得到 $u_t$ 解析式：
$$u_t(x|x_1)=\frac{{\alpha }_{1-t}^{{}^{\prime }}}{1-{\alpha }_{1-t}^2}({\alpha }_{1-t}x-x_1)=-\frac{T^{{}^{\prime }}(1-t)}{2}\left[e^{-T(1-t)}x-\frac{e^{-\frac{1}{2}T(1-t)}{x}_1}{1-e^{-T(1-t)}}\right]$$

Optimal Transport

由于 Flow Matching 不依赖扩散过程，可以构建一个最优传输，将条件 $p_t$ 的 $\mu$ 和 $\sigma$ 构建为简单的随时间的线性变换：
$$\begin{array}{r}{\mu }_t(x)=t{x}_1\\ {\sigma }_t(x)=1-(1-{\sigma }_{min})t\end{array}$$
代入式（10）得到条件向量场 $u_t$：
$$u_t(x|x_1)=\frac{x_1-(1-{\sigma }_{min})x}{1-(1-{\sigma }_{min})t}$$
根据式（8）对应的条件流 $\psi$ 有：
$${\psi }_t(x)={\sigma }_t(x)x+{\mu }_t(x)=(1-(1-{\sigma }_{min})t)x+t{x}_1(12)$$
根据式（2）和式（9）重参数化 $\psi$ 得到目标函数为：
$${\mathcal{L}}_{\mathbf{CFM}}(\theta )={\mathbb{E}}_{t,q(x_1),p_0(x_0)}\Vert v_t({\psi }_t(x_0))-(x_1-(1-{\sigma }_{min})x_0){\Vert }^2(13)$$

Reference：

1. https://zhuanlan.zhihu.com/p/741939590
2. https://zhuanlan.zhihu.com/p/685921518