生成模型 (3.2)：Flow Model

2025-12-18 生成模型 #生成模型 #流模型 #连续性方程 #瞬时换元公式约 31 分钟 · 10792 字

引言

从这一章开始深入到FM的技术细节中。首先从流模型开始讲起。

流 (flows) 是一种最简单的连续时间马尔可夫过程 (Continuous Time Markov Process, CTMP)。在理论上，流模型的优势有以下三点：

流可以让我们从任意的初始分布转移到任意的目标分布
通过估计流所对应的ODE的解，我们可以较为高效地在流中进行采样（相较于求解Diffusion中的SDE）
流具有确定性的过程，因此可以得到似然函数的无偏估计。而SDE中由于是随机过程，因此只能去估计似然的下界。

同样的，以下内容完全来自于Flow Matching Guide and Code。再次感谢FAIR的大佬们

一、概率论和微分方程的一些预备知识

1.1. 随机向量

考虑欧氏空间 $\mathbb{R}^d$ ，以及定义在这个空间上的内积运算 $\langle x,y\rangle = \sum_{i=1}^dx_iy_i$ 和2-范数 $\|x\|_2=\sqrt{\langle x,x\rangle}$ 。

随机向量 (Random Vector, RV) 是一个定义在 $\mathbb{R}^d$ 上的一个随机变量 $X\in\mathbb{R}^d$ ，且具有连续的PDF $p_{X}$ 。PDF满足归一性：

(1)

一个最常见的PDF就是 $d$ 维的高斯分布：

(2)

其中， $\mu\in\mathbb{R}^d,\sigma\in\mathbb{R}^+$ 分别表示均值和标准差。

可以证明，随机向量 $X$ 的【期望】等于在2-范数意义下最接近 $X$ 的一个常向量 $z$ 。（证明见附录）

(3)

1.2. 条件概率密度与条件期望

给定两个随机向量 $X,Y\in\mathbb{R}^d$ ，以及这两个随机向量的联合PDF $p_{X,Y}(x,y)$ 。我们可以定义两个边缘分布：

(4)

基于此，我们还可以定义如下的条件分布：

(5)

计算条件分布的一个重要手段是贝叶斯定理：

(6)

条件期望 $\mathbb{E}[X\mid Y]=g_\star(Y)$ ，是一个在2-范数意义下最接近 $X$ 的随机向量：

(7)

当 $Y=y\in\mathbb{R}^d$ 时，这个条件期望变为：

(8)

公式 $(8)$ 的证明和公式 $(3)$ 非常类似。

两个随机向量的条件期望有如下的重要性质，称为tower property：

(9)

1.3. 微分同胚 (Diffeomorphisms) 和前推 (push-forward)

首先，我们定义 $C^r(\mathbb{R}^m,\mathbb{R}^n)$ 为一个函数集合 $\{f\mid f:\mathbb{R}^m\mapsto\mathbb{R}^n\}$ ，且这些函数具有连续的 $r$ 阶偏导数：

(10)

其中， $k\in[n]$ 且 $i_j\in[r]$ 。特别地，定义 $C^r(\mathbb{R}^m):=C^r(\mathbb{R}^m,\mathbb{R})$ 。

一个函数 $\psi:\mathbb{R}^d\mapsto\mathbb{R}^d$ 是一个【微分同胚】，当且仅当：

$\psi$ 是可逆函数
$\psi\in C^r(\mathbb{R}^d,\mathbb{R}^d)$ 且 $\psi^{-1}\in C^r(\mathbb{R}^d,\mathbb{R}^d)$ 。

微分同胚可以直观理解为一种可逆的、平滑的变换。在流模型中，我们必须保证所有变换都是一个位于 $C^1(\mathbb{R}^d,\mathbb{R}^d)$ 的微分同胚,这样才能在初始分布和目标分布之间来回转换。

给定一个随机向量 $X\sim p_X$ 以及一个微分同胚 $\psi\in C^1(\mathbb{R}^d,\mathbb{R}^d)$ 。我们考虑变换后的随机向量 $Y=\psi(X)\sim p_Y$ ，这个随机向量的PDF —— 也称为 $p_X$ 的一个前推 (push-forward) —— 有如下形式：

(11)

其中， $\partial_y\psi^{-1}(y)$ 表示 $\psi^{-1}(y)$ 的雅可比矩阵。这个公式其实就是多元积分的换元公式。

因此我们可以定义一个前推算子 $\sharp$ ：

(12)

二、流的数学表达

经过这么多铺垫，我们终于可以进入正题。

回顾在前一篇中提到的，生成式建模的总体目标就是把初始分布 $p$ 上的一个样本 $X_0\sim p$ 映射到目标分布 $q$ 上的一个样本 $X_1\sim q$ 。

2.0. 一些定义准备

一个 $C^r$ 流是一个与时间有关的映射函数 $\psi:[0,1]\times \mathbb{R}^d\mapsto\mathbb{R}^d$ ，且有 $\psi\in C^r([0,1]\times \mathbb{R}^d,\mathbb{R}^d)$ 。

进一步地，我们规定 $\forall t\in[0,1]$ ， $\psi(t,x)=\psi_t(x)$ 是一个关于 $x$ 的微分同胚。

流模型是一个连续时间马尔可夫过程 (CTMP)，把一个流 $\psi$ 作用在一个随机向量 $X_0$ 上：

(13)

$X_t$ 的马尔可夫性可以通过下面的推导证明。

任意给定 $0\le t\le s\le 1$ ，我们有

(14)

其中 $\psi_{s\mid t}:=\psi_s \circ \psi_t^{-1}$ 同样也是一个微分同胚。公式 $(14)$ 说明了 $t$ 时刻以后的任意状态都只取决于 $X_t$ ，也就说明了马尔可夫性。

最终，流模型的目标是找到一个流 $\psi$ ，使得 $X_1:=\psi_1(X_0)\sim q$ 。

2.1. 流与速度场的等价性

一个 $C^r$ 流同样可以由一个位于 $C^r([0,1]\times \mathbb{R}^d,\mathbb{R}^d)$ 中的速度场 $u:[0,1]\times \mathbb{R}^d\mapsto\mathbb{R}^d$ 定义。 $\psi$ 和 $u$ 之间通过下面的这个带初始值的ODE来描述：

(15)

下面的定义保证了公式 $(15)$ 对应的解的存在性和唯一性。

Theroem 1 (Picard–Lindelof)：如果速度场 $u$ 满足以下两个条件：

$u\in C^r,r\ge 1$ ;
在局部上满足利普希茨条件，即在一个邻域上满足 $\forall x,y,\|u(t,x)-u(t,y)\|\le L\cdot \|x-y\|$

那么我们有：

公式(15)存在唯一解 $\psi$ 。对于任何一个起始位置，都有且仅有一条路径。
$\psi$ 是一个微分同胚，即解是光滑的。
在 $t=0$ 时，所有起始位置都存在唯一解，即解的定义域 $\Omega=\{0\}\times\mathbb{R}^d$ 。

需要注意的是，这个定理只在【局部】上保证了解的存在性和唯一性。也就是说，只能保证一个流在有限的时间内对某个样本 $x\in\mathbb{R}^d$ 的作用是唯一且确定的。

如果要保证在整个 $t\in[0,1]$ 时间内都成立，则需要有更强的条件（如全局利普希茨性质）或者引入一些新的假设。在流模型的理论中，引入的是【可积性】的假设，来保证这个流在 $t\in[0,1]$ 的时间中几乎处处存在。

至此，我们证明了一个速度场能够唯一定义一个流。同样地，给定一个 $C^1$ 流 $\psi_t$ ，我们也可以得到出其对应的速度场 $u_t$ ：

。 (16)

2.2. 求解ODE

给定初始样本 $X_0=x_0$ ，求解任意时刻 $t$ 下的样本 $X_t$ 就是在估计公式(15)的解。

求解ODE初值问题的一个最简单的方法就是欧拉法，这是一种迭代求解的方法：

(17)

其中 $n$ 是迭代次数， $h=\frac{1}{n}$ 是迭代步长。

欧拉法的误差项 $o(h)$ 是步长 $h$ 的高阶无穷小。也就是说理论上只要迭代次数 $n$ 足够多，欧拉法就能够收敛到准确解上。

三、概率路径和连续性方程

我们将一个与时间相关的概率序列 $(p_t)_{0\le t\le 1}$ 称为一个概率路径 (probability path)。

在流模式的场景中，我们关心的是流 $\psi_t$ 对应的概率路径。对于任意时间 $t\in[0,1]$ ，概率值都可以由公式(12)中定义的前推算子得到：

(18)

其中 $p$ 是初始分布。

我们称速度场 $u_t$ 生成概率路径 $p_t$ ，当且仅当 $\forall t\in[0,1)\ X_t=\psi_t(X_0)\sim p_t$ 。

我们可以通过验证 $u_t$ 和 $p_t$ 是否满足下面的PDE，来判断速度场是否能够生成概率路径：

(19)

其中， $\textrm{div}(v)(x)=\text{tr}[\partial_{x}v(x)]\in\mathbb{R}^{d\times d}$ 称为散度矩阵。

公式(19)是物理学中一个非常重要的方程，称为 连续性方程 (Continuity Equation)。连续性方程描述了一些物理量（比如质量、电荷等）的总量不随时间变化的性质，也就是质量/电荷守恒定律。

下面的定理2是质量守恒方程的等价形式，它说明了【连续性方程的解 $u_t$ 能够生成概率路径 $p_t$ 】。

Theroem 2 (Mass Conservation)：令 $p_t$ 是一个概率路径， $u_t$ 是一个满足局部利普希茨性质的可积速度场，则我们有如下的等价表述：

$\forall t\in [0,1)$ ，公式 (19) 成立
$u_t$ 能够生成 $p_t$

这里引入了【可积性】的性质之后， $X_t$ 就能拥有一些比较良好的性质，即 $\|X_t\|$ 的期望是有界的。下面我们展示这一点。

速度场 $u_t$ 可积意味着：

(20)

我们对速度场所对应的ODE两边对 $t$ 积分得：

(21)

代入 $\psi_0(x)=x$ 得 $C=x$ 。

因此，我们有：

(22)

其中最后一步用到了三角不等式。

公式(22)说明了，只要 $\|X_0\|$ 的期望是有界的，则 $\|X_t\|$ 也是。

Note：这里的推导我也看不太明白，主要是积分变量从 $s$ 变成了 $t$ 。希望各位大佬指正。

四、瞬时换元 (Instantaneous Change of Variable, ICV)

流模型的一个显著优势在于，它的似然值是可解的 (tractable)。也就是说，我们可以计算出 $\log p_1(x)$ 的解析解。这个性质是通过在连续性方程 —— 公式(19) —— 上使用瞬时换元公式得到的。下面我们展示这一点。

我们的结论建立在下面的定理3上，定理3也被称为瞬时换元公式。

Theroem 3: Instantaneous Change of Variable

(23)

定理3的证明见附录。

我们对公式 $(23)$ 两端从 $t=0$ 到 $t=1$ 积分得：

(24)

计算散度 $\textrm{div}(u_t)$ 等价于计算雅可比矩阵 $\partial_x u_t(x)\in\mathbb{R}^{d\times d}$ 的迹，需要 $O(d)$ 次求导。因此，在实践中，一般使用对散度进行无偏估计。一个常用的无偏估计为Hutchinson迹估计，用于估计一个随机矩阵的迹：

(25)

其中， $Z\in\mathbb{R}^{d\times d}$ 是任意的随机变量，满足 $\mathbb{E}[Z]=0$ 且 $\textrm{Cov}(Z,Z)=I$ 。

将公式 $(25)$ 代入公式 $(24)$ 得：

(26)

引入随机变量 $Z$ 后，我们可以通过 vector-Jacobian product 的方法（比如torch.autograd.functional.vjp）来快速计算 $\textrm{tr}\left[Z^T\partial_x u_t(\psi_t(x))Z\right]$ 。

总结一下，在求解似然估计 $\log p_1(x)$ 时，我们实际上是从 $t=1$ 到 $t=0$ 反向模拟一下面的ODE：

(27)

得到 $f(0)$ 和 $g(0)$ 后，我们就能得到似然函数的无偏估计值：

(28)

五、流模型的训练

在得到似然估计之后，很自然地我们就可以通过最大化这个似然估计来训练一个流模型。

具体来说，设 $u_t^\theta$ 是一个参数为 $\theta$ 的速度场，我们的希望让 $p_1$ 尽可能接近我们的目标分布 $q$ 。

因此，我们可以通过下面的最大似然目标来优化参数 $\theta$ ：

(29)

其中， $p_1^\theta$ 是 $X_1=\psi_1^\theta(X_0)$ 的概率密度函数，而 $\psi_t^\theta$ 是由速度场 $u_t^\theta$ 定义的流。

附录

公式(3)的证明

首先根据向量2-范数的定义，展开 $\|x-z\|_2^2$ ：

(30)

我们需要最小化的函数为：

(31)

对 $z$ 求导得：

(32)

令梯度为0得：

(33)

因此，

(34)

Q.E.D.

瞬时换元公式的证明

下面我们来证明定理3（公式(23)）

由全导数链式法则，公式 $(23)$ 的左侧为：

(35)

将流的ODE（公式(15)）代入上式中，得到

(36)

接下来我们看连续性方程（公式19）。根据散度的乘法法则： $\textrm{div}(fV)=f(\textrm{div} V)+(\nabla f)\cdot V$ 得：

(37)

注意到公式 $(36)$ 对所有时间 $t$ 和所有位置 $x$ 均成立，因此代入 $x=\psi_t(x)$ 得：

(38)

将上式代入公式 $(35)$ 得到：

(39)

O.D.E.