类似连续时间步版本的 Diffusion。有一个非常优美的结论。

发现数学内容更多的论文往往名字更短

[2210.02747] Flow Matching for Generative Modeling（2022 年 10 月）

参考的博客：通俗易懂的Flow Matching原理解读（附核心公式推导和源代码） – 知乎

也参考了一下：笔记｜扩散模型（一八）：Flow Matching 理论详解 | 極東晝寢愛好家

啥是流？要干啥？

有一组图片。生成模型希望学习之后能生成这些图片。

生成模型的任务：有一个初始分布 $p_0$（Gaussian 分布之类的，总之不包含待生成图片集的任何信息），有一个目标分布 $p_1$（就是希望生成的那些图片），希望学一个映射把 $p_0$ 变成 $p_1$。

一步到位的生成不太靠谱；DDPM 系列中，这样的映射是分 $T$ 个时间步进行的，每步是一个小映射，因此是离散的，我们考虑如果把这个过程变成连续的呢？

现在假设我们对每个 $t\in [0,1]$ 都有一个分布 $p_t$；取一个初始分布中的 $x$，存在一条以 $x$ 为起点的连续路径 ${\psi }_t(x),t\in [0,1]$ 最终变到目标分布里。所有这些路径构成一个流，我们考虑其速度场 $u(t,x)$：

$$dx=u(t,x)dt$$

假设我们已经学出这个场了，生成方式就是随机取一个 $x_0\sim p_0$，然后沿 $u$ 的方向流动：

$$x_{t+\mathrm{\Delta }t}=x_t+u(t,x_t)\mathrm{\Delta }t$$

（或者它的连续形式，为了在理论上讨论的话）

我们希望模型学的就是这个 $u$：

$$L_{FM}={\mathbb{E}}_{t\in [0,1],x_t\sim p_t}\parallel u_{\theta }(t,x_t)-u(t,x_t){\parallel }^2$$

两个问题：其一，我们根本不知道 $u$ 是啥（应该存在很多 $u$，但我们目前也构造不出任何一个实例）；其二，我们不知道 $p_t$，从而也没法取 $x_t$。

但是我们还根本没用上我们的图片集 $p_1$ 啊！真实的训练过程肯定是取一张图片 $z\sim p_1$ 然后干点什么事情……

论文的想法是这样的：

• 考虑条件流 $u(t,x_t|z)$，这个流会把初始分布 $p_0$ 变成一个单点 $z$，这样的 $u$ 是容易构造的；
• 发现用条件流作优化目标，优化效果和无条件流是完全等价的；
• 构造一个具体的条件流。

这三个都完成之后，训练过程就变成了取 $z$，算 $u(t,x_t|z)$，直接学。

Step 1. $u$

这是论文中的定理 1，也是上面说的第一步。

假设有 $z\sim q(z)$，考虑条件分布 $p(t,x|z)$ 和边缘分布 $p(t,x)$ 之间的关系：

$$p(t,x)=\int p(t,x|z)q(z)dz$$

以下将 $p(t,x)$ 和 $p_t(x)$ 混用。它们都表示无条件流。

注意到流动中的分布 $p_t$ 始终满足全概率为 $\int p(t,x)dx=1$ 类似流体的质量守恒；我们从流体力学里面借过来一个连续性方程，这是 $p_t$ 需要满足的条件：

$$\frac{\partial p(t,x)}{\partial t}=-\mathrm{\nabla }\cdot (p(t,x)u(t,x))$$

连续性方程是描述某种物理量守恒的基本方程，广泛应用于流体力学、电磁学等领域。其核心思想是某一物理量在系统中的变化率等于流入和流出该系统的净流量。

现在把之前那个关系式对 $t$ 求导，跟连续性方程对比一下：

$$\frac{\partial p(t,x)}{\partial t}=\int q(z)\frac{\partial p(t,x|z)}{\partial t}dz$$

条件分布也是分布，从而也有连续性方程：

$$\frac{\partial p(t,x|z)}{\partial t}=-\mathrm{\nabla }\cdot (p(t,x|z)u(t,x|z))$$

现在把两个连续性方程都代进去

$$\mathrm{\nabla }\cdot (p(t,x)u(t,x))=\int q(z)(\mathrm{\nabla }\cdot (p(t,x|z)u(t,x|z)))dz$$

最后得到

$$u(t,x)=\int q(z)\frac{p(t,x|z)u(t,x|z)}{p(t,x)}dz$$

完工！

我们在干什么？

如果我们找了一个条件流 $p(t,x|z)$（从而也有 $u(t,x|z)$）并且知道 $q(z)$，则可以确定无条件流 $u(t,x)$。这个式子将被用在 Step 2 的等价性证明中。

Step 2. $L_{CFM}=L_{FM}$

这是论文中的定理 2，也是上面说的第二步。

$L_{FM}$ 是以无条件流为目标的优化：

$$L_{FM}={\mathbb{E}}_{t\in [0,1],x_t\sim p_t}\parallel u_{\theta }(t,x_t)-u(t,x_t){\parallel }^2$$

取一个 $z\sim q(z)$，人工选定一个条件流 $p(t,x|z)$，$L_{CFM}$ 是以条件流为目标的优化：

$$L_{CFM}={\mathbb{E}}_{t\in [0,1],z\sim q(z),x_t\sim p(t,x|z)}\parallel u_{\theta }(t,x_t)-u(t,x_t|z){\parallel }^2$$

怎么样才能说它们等价呢？我们考察

$${\mathrm{\nabla }}_{\theta }{L}_{CFM}={\mathrm{\nabla }}_{\theta }\mathbb{E}\parallel u_{\theta }(t,x_t){\parallel }^2-2{\mathrm{\nabla }}_{\theta }\mathbb{E}[u_{\theta }(t,x_t)\cdot u(t,x_t|z)]$$

和

$${\mathrm{\nabla }}_{\theta }{L}_{FM}={\mathrm{\nabla }}_{\theta }\mathbb{E}\parallel u_{\theta }(t,x_t){\parallel }^2-2{\mathrm{\nabla }}_{\theta }\mathbb{E}[u_{\theta }(t,x_t)\cdot u(t,x_t)]$$

只要它们相等就行啦！

好、好离谱的操作，但真的好有道理……

看别的写博客的人都对此没啥看法，难道这是个常规操作吗……

它们相等只需要后面那项相等即可。我们利用 Step 1 的最终结果：

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }\mathbb{E}[u_{\theta }(t,x_t)\cdot u(t,x_t)]& ={\mathrm{\nabla }}_{\theta }\iint u_{\theta }(t,x_t)u(t,x_t)p(t,x_t)d{x}_{t}dt\\ & ={\mathrm{\nabla }}_{\theta }\iint u_{\theta }(t,x_t)\int q(z)\frac{p(t,x_t|z)u(t,x_t|z)}{p(t,x_t)}p(t,x_t)dzd{x}_{t}dt\\ & ={\mathrm{\nabla }}_{\theta }\iiint q(z)u_{\theta }(t,x_t)p(t,x_t|z)u(t,x_t|z)dzd{x}_{t}dt\\ & ={\mathrm{\nabla }}_{\theta }\mathbb{E}[u_{\theta }(t,x_t)\cdot u(t,x_t|z)]\end{array}$$

完工！

我们在干什么？

我们证明了：取一个 $z\sim q(z)$，人工选定一个条件流 $p(t,x|z)$，以 $u(t,x|z)$ 为优化目标和以无条件 $u(t,x)$ 为优化目标是等价的。

Step 3. $p:=\mathcal{N}$

Gaussian 分布是个好东西。我们假设 $p(t,x|z)=\mathcal{N}({\mu }_t,{\sigma }_t^2)$。$t=0$ 时 $x_0\sim \mathcal{N}({\mu }_0,{\sigma }_0^2)$。我们取条件流为

$${\psi }_t(x_0|z)={\mu }_t+{\sigma }_t(\frac{x_0-{\mu }_0}{{\sigma }_0})$$

为啥能这么取？

初始分布 $p_0$ 是 Gaussian 分布，条件目标分布是单点，可以看成方差为某个小量 ${\sigma }_{ϵ}$ 的 Gaussian 分布，因此中间插一堆 $\mu$ 和 $\sigma$ 线性变化的 Gaussian 分布是连续的；然后上面那个 ${\psi }_t$ 确实是一个流。

总之能取出来就行

现在来算 $u$。

$$u(t,{\psi }_t(x_0)|z)=\frac{d}{dt}{\psi }_t(x_0|z)=\frac{d{\mu }_t}{dt}+\frac{x_0-{\mu }_0}{{\sigma }_0}\frac{d{\sigma }_t}{dt}$$

用 $x$ 代替 ${\psi }_t(x_0)$，右边按上面的定义式换元：

$$u(t,x|z)=\frac{d{\mu }_t}{dt}+\frac{x-{\mu }_t}{{\sigma }_t}\frac{d{\sigma }_t}{dt}$$

按刚刚说的，让 $\mu$ 和 $\sigma$ 线性变化，${\mu }_t={\mu }_0+t(z-{\mu }_0)$，${\sigma }_t={\sigma }_0-t{\sigma }_0$，再假设初始分布是标准正态分布，${\mu }_0=0,{\sigma }_0=1$代入，得到

$$u(t,x|z)=\frac{z-(1-{\sigma }_{ϵ})x}{1-(1-{\sigma }_{ϵ})t}$$

一个不错的目标。

回到 Step 2，我们注意到，这个证明中用到的 $z$ 其实并不要求取自目标分布 $p_1$ 而可以是随便什么东西。现在假设我们在初始分布中随机采样了一个 Gaussian 噪声 $x_0$，在目标分布中随便取了一个图片 $x_1$。

令 $p_t(x|x_0,x_1)=\mathcal{N}((1-t)x_0+t{x}_1,{\sigma }_{ϵ})$，代入上面得到

$$u_t(x|x_0,x_1)=x_1-x_0$$

wow.