引言
从这一章开始深入到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满足归一性：
\[\begin{equation} \int p_X(x)\mathrm{d}x=1 \end{equation}\]一个最常见的PDF就是$d$维的高斯分布：
\[\begin{equation} \begin{aligned} p_X(x) &=\mathcal{N}(x\mid \mu,\sigma^2I)\\ &=\left( 2\pi\sigma^2 \right)^{-\frac d2} \exp\left( -\frac{\|x-\mu\|_2^2}{2\sigma^2} \right) \end{aligned} \end{equation}\]其中，$\mu\in\mathbb{R}^d,\sigma\in\mathbb{R}^+$分别表示均值和标准差。
可以证明，随机向量 $X$ 的【期望】等于在2-范数意义下最接近 $X$ 的一个常向量 $z$ 。（证明见附录）
\[\begin{equation} \begin{aligned} \mathbb{E}[X] &:=\int xp_X(x)\mathrm{d}x\\ &=\arg\min_{z\in\mathbb{R}^d} \int p_X(x) \|x-z\|^2_2\mathrm{d}x\\ \end{aligned} \end{equation}\]1.2. 条件概率密度与条件期望
给定两个随机向量 $X,Y\in\mathbb{R}^d$，以及这两个随机向量的联合PDF $p_{X,Y}(x,y)$。我们可以定义两个边缘分布：
\[\begin{equation} \begin{aligned} p_X(x)&=\int p_{X,Y}(x,y)\mathrm{d}y\\ p_Y(y)&=\int p_{X,Y}(x,y)\mathrm{d}x\\ \end{aligned} \end{equation}\]基于此，我们还可以定义如下的条件分布：
\[\begin{equation} p_{X\mid Y}(x\mid y)=\frac{p_{X,Y}(x,y)}{p_Y(y)} \end{equation}\]计算条件分布的一个重要手段是贝叶斯定理：
\[\begin{equation} p_{Y\mid X}(y\mid x)=\frac{p_{X\mid Y}(x\mid y)p_Y(y)}{p_X(x)} \end{equation}\]条件期望$\mathbb{E}[X\mid Y]=g_\star(Y)$，是一个在2-范数意义下最接近$X$的随机向量：
\[\begin{equation} \begin{aligned} g_\star &:=\arg\min_{g:\mathbb{R}^d\mapsto\mathbb{R}^d} \mathbb{E}\left[ \|X-g(Y)\|_2^2 \right]\\ &=\arg\min_{g:\mathbb{R}^d\mapsto\mathbb{R}^d}\int\int p_{X,Y}(x,y)\cdot \|x-g(y)\|_2^2 \mathrm{d}x\mathrm{d}y\\ &=\arg\min_{g:\mathbb{R}^d\mapsto\mathbb{R}^d}\int \left[\int p_{X\mid Y}(x\mid y)\cdot \|x-g(y)\|_2^2 \mathrm{d}x \right]p_Y(y)\mathrm{d}y\\ \end{aligned} \end{equation}\]当$Y=y\in\mathbb{R}^d$时，这个条件期望变为：
\[\begin{equation} \begin{aligned} \mathbb{E}[X\mid Y=y]&=g_\star(y)\\ &=\int xp_{X\mid Y}(x\mid y)\mathrm{d}x \end{aligned} \end{equation}\]公式$(8)$的证明和公式$(3)$非常类似。
两个随机向量的条件期望有如下的重要性质，称为tower property：
\[\begin{equation} \begin{aligned} \mathbb{E}[\mathbb{E}[X\mid Y]] &=\int \left(\int x\cdot p_{X\mid Y}(x\mid y) \mathrm{d}x \right)p_Y(y)\mathrm{d}y\\ &=\int \int x\cdot p_{X, Y}(x, y) \mathrm{d}x\mathrm{d}y\\ &=\int x\cdot p_{X}(x) \mathrm{d}x\\ &=\mathbb{E}[X] \end{aligned} \end{equation}\]1.3. 微分同胚 (Diffeomorphisms) 和前推 (push-forward)
首先，我们定义 $C^r(\mathbb{R}^m,\mathbb{R}^n)$ 为一个函数集合 ${f\mid f:\mathbb{R}^m\mapsto\mathbb{R}^n} $，且这些函数具有连续的 $r$ 阶偏导数：
\[\begin{equation} \frac{\partial^r f_k}{\partial x_{i_1}\cdots\partial x_{i_r}} \end{equation}\]其中，$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) —— 有如下形式：
\[\begin{equation} p_Y(y)=p_X(\psi^{-1}(y))|\det \partial_y\psi^{-1}(y)| \end{equation}\]其中，$\partial_y\psi^{-1}(y)$ 表示 $\psi^{-1}(y)$ 的雅可比矩阵。这个公式其实就是多元积分的换元公式。
因此我们可以定义一个前推算子 $\sharp$：
\[\begin{equation} [\psi_\sharp p_X](y):=p_X(\psi^{-1}(y))|\det \partial_y\psi^{-1}(y)| \end{equation}\]二、流的数学表达
经过这么多铺垫，我们终于可以进入正题。
回顾在前一篇中提到的，生成式建模的总体目标就是把初始分布 $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$ 上：
\[\begin{equation} X_t=\psi_t(X_0)\quad t\in[0,1],\ X_0\sim p \end{equation}\]$X_t$ 的马尔可夫性可以通过下面的推导证明。
任意给定 $0\le t\le s\le 1$，我们有
\[\begin{equation} \begin{aligned} X_s &=\psi_s(X_0)\\ &=\psi_s(\psi_t^{-1}(\psi_t(X_0)))\\ &=\psi_{s\mid t}(X_t) \end{aligned} \end{equation}\]其中$\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来描述：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\psi_t(x)&=u_t(\psi_t(x))\\ \psi_0(x)&=x \end{aligned} \end{equation}\]下面的定义保证了公式$(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$ ：
\[\begin{equation} u_t(x)=\frac{\mathrm{d}\psi_t}{\mathrm{d}t}(\psi_t^{-1}(x))。 \end{equation}\]2.2. 求解ODE
给定初始样本 $X_0=x_0$，求解任意时刻 $t$ 下的样本 $X_t$ 就是在估计公式(15)的解。
求解ODE初值问题的一个最简单的方法就是欧拉法，这是一种迭代求解的方法：
\[\begin{equation} X_{t+h}=X_t+hu_t(X_t)+o(h) \end{equation}\]其中 $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)中定义的前推算子得到：
\[\begin{equation} p_t(x)=[\psi_{t\sharp} p](x) \end{equation}\]其中 $p$ 是初始分布。
我们称速度场 $u_t$ 生成概率路径 $p_t$，当且仅当 $\forall t\in[0,1)\ X_t=\psi_t(X_0)\sim p_t$。
我们可以通过验证 $u_t$ 和 $p_t$ 是否满足下面的PDE，来判断速度场是否能够生成概率路径：
\[\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}$ 称为散度矩阵。
公式(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$ 可积意味着：
\[\begin{equation} \int_0^1\int \|u_t(x)\|p_t(x)\mathrm{d}x\mathrm{d}t< \infty \end{equation}\]我们对速度场所对应的ODE两边对 $t$ 积分得：
\[\begin{equation} \psi_t(x) = \int_0^t u_s(\psi_s(x))\mathrm{d}s+C \end{equation}\]代入 $\psi_0(x)=x$ 得 $C=x$。
因此，我们有：
\[\begin{equation} \begin{aligned} \mathbb{E}\|X_t\| &=\int\|\psi_t(x)\|p(x)\mathrm{d}x\\ &=\int\|x+\int_0^t u_s(\psi_s(x))\mathrm{d}s\|p(x)\mathrm{d}x\\\\ &\le \int \|x\|p(x)\mathrm{d}x + \int_0^1\int \|u_t(x)\|p_t(x)\mathrm{d}x\mathrm{d}t\\ \end{aligned} \end{equation}\]其中最后一步用到了三角不等式。
公式(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
\[\begin{equation} \frac{\mathrm{d}}{\mathrm{d}t}\log p_t(\psi_t(x))=-\textrm{div}(u_t)(\psi_t(x)) \end{equation}\]定理3的证明见附录。
我们对公式$(23)$两端从$t=0$到$t=1$积分得：
\[\begin{equation} \log p_1(\psi_1(x))=\log p_0(\psi_0(x))-\int_0^1\textrm{div}(u_t)(\psi_t(x))\mathrm{d}t \end{equation}\]计算散度 $\textrm{div}(u_t)$ 等价于计算雅可比矩阵 $\partial_x u_t(x)\in\mathbb{R}^{d\times d}$ 的迹，需要 $O(d)$ 次求导。因此，在实践中，一般使用对散度进行无偏估计。一个常用的无偏估计为Hutchinson迹估计，用于估计一个随机矩阵的迹：
\[\begin{equation} \textrm{div}(u_t)(x)=\textrm{tr}[\partial_x u_t(x)]=\mathbb{E}_Z\textrm{tr}[Z^T\partial_x u_t(x)Z] \end{equation}\]其中，$Z\in\mathbb{R}^{d\times d}$ 是任意的随机变量，满足 $\mathbb{E}[Z]=0$ 且 $\textrm{Cov}(Z,Z)=I$。
将公式$(25)$代入公式$(24)$得：
\[\begin{equation} \log p_1(\psi_1(x))=\log p_0(\psi_0(x))-\mathbb{E}_Z\int_0^1\textrm{tr}\left[Z^T\partial_x u_t(\psi_t(x))Z\right]\mathrm{d}t \end{equation}\]引入随机变量 $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：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} \begin{bmatrix} f(t)\\ g(t) \end{bmatrix} &=\begin{bmatrix} u_t(f_t)\\ -\textrm{tr}\left[Z^T\partial_x u_t(f(t))Z\right] \end{bmatrix}\\ \begin{bmatrix} f(1)\\ g(1) \end{bmatrix} &=\begin{bmatrix} x \\ 0 \end{bmatrix} \end{aligned} \end{equation}\]得到 $f(0)$ 和 $g(0)$ 后，我们就能得到似然函数的无偏估计值：
\[\begin{equation} \log p_1(x)=\log p_0(f(0))-g(0) \end{equation}\]五、流模型的训练
在得到似然估计之后，很自然地我们就可以通过最大化这个似然估计来训练一个流模型。
具体来说，设 $u_t^\theta$ 是一个参数为 $\theta$ 的速度场，我们的希望让 $p_1$ 尽可能接近我们的目标分布 $q$。
因此，我们可以通过下面的最大似然目标来优化参数 $\theta$：
\[\begin{equation} \mathcal{L}(\theta)=D_{KL}(q;p_1^\theta) =-\mathbb{E}_{Y\sim q}\log p_1^\theta (Y)+\text{constant} \end{equation}\]其中，$p_1^\theta$ 是 $X_1=\psi_1^\theta(X_0)$ 的概率密度函数，而 $\psi_t^\theta$ 是由速度场 $u_t^\theta$ 定义的流。
附录
公式(3)的证明
首先根据向量2-范数的定义，展开$|x-z|_2^2$：
\[\begin{equation} \begin{aligned} \|x-z\|_2^2 &=(x-z)^T(x-z)\\ &=(x^T-z^T)(x-z)\\ &=x^Tx-2z^Tx+z^Tz \end{aligned} \end{equation}\]我们需要最小化的函数为：
\[\mathcal{J}(z)=\int p_X(x) \left(x^Tx-2z^Tx+z^Tz\right)\mathrm{d}x\]对$z$求导得：
\[\begin{equation} \begin{aligned} \nabla_z\mathcal{J}(z) &=\nabla_z\int p_X(x) \left(x^Tx-2z^Tx+z^Tz\right)\mathrm{d}x\\ &=\int \nabla_z \left[p_X(x) \left(x^Tx-2z^Tx+z^Tz\right)\mathrm{d}x\right]\\ &=\int p_X(x) \left(-2x+2z\right)\mathrm{d}x\\ &=\int 2\cdot p_X(x)\cdot \left(z-x\right)\mathrm{d}x\\ \end{aligned} \end{equation}\]令梯度为0得：
\[\begin{equation} \begin{aligned} &\int 2\cdot p_X(x)\cdot \left(z-x\right)\mathrm{d}x=0\\ &\int zp_X(x)\mathrm{d}x=\int xp_X(x)\mathrm{d}x\\ &z=\int xp_X(x)\mathrm{d}x \end{aligned} \end{equation}\]因此，
\[\begin{equation} \arg\min_{z\in\mathbb{R}^d} \int p_X(x) \|x-z\|^2_2\mathrm{d}x=\int xp_X(x)\mathrm{d}x \end{equation}\]Q.E.D.
瞬时换元公式的证明
下面我们来证明定理3（公式(23)）
由全导数链式法则，公式$(23)$的左侧为：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\log p_t(\psi_t(x)) &=\left.\left[ \frac{\partial}{\partial t}\log p_t(y) + \nabla_y\log p_t(y) \cdot\frac{\mathrm{d}y}{\mathrm{d}t}\right]\right|_{y=\psi_t(x)}\\ &=\frac{1}{p_t(\psi_t(x))}\left[\frac{\partial p_t}{\partial t}\psi_t(x) + \nabla p_t(\psi_t(x))\cdot \frac{\mathrm{d}\psi_t(x)}{\mathrm{d}t} \right] \end{aligned} \end{equation}\]将流的ODE（公式(15)）代入上式中，得到
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\log p_t(\psi_t(x)) &=\frac{1}{p_t(\psi_t(x))}\left[\frac{\partial p_t}{\partial t}\psi_t(x) + \nabla p_t(\psi_t(x))\cdot u_t(\psi_t(x)) \right] \end{aligned} \end{equation}\]接下来我们看连续性方程（公式19）。根据散度的乘法法则：$\textrm{div}(fV)=f(\textrm{div} V)+(\nabla f)\cdot V$ 得：
\[\begin{equation} \begin{aligned} &\frac{\partial}{\partial t}p_t+p_t\cdot\textrm{div}(u_t)+\nabla p_t\cdot u_t =0\\ &\frac{\partial p_t}{\partial t}+\nabla p_t\cdot u_t =-p_t\cdot\textrm{div}(u_t) \end{aligned} \end{equation}\]注意到公式$(36)$对所有时间 $t$ 和所有位置 $x$ 均成立，因此代入 $x=\psi_t(x)$ 得：
\[\begin{equation} \begin{aligned} &\frac{\partial p_t}{\partial t}(\psi_t(x))+\nabla p_t(\psi_t(x))\cdot u_t(\psi_t(x))=-p_t(\psi_t(x))\cdot\textrm{div}(u_t(\psi_t(x))) \end{aligned} \end{equation}\]将上式代入公式$(35)$得到：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\log p_t(\psi_t(x)) &=\frac{1}{p_t(\psi_t(x))}\left[\frac{\partial p_t}{\partial t}\psi_t(x) + \nabla p_t(\psi_t(x))\cdot u_t(\psi_t(x)) \right]\\ &=\frac{1}{p_t(\psi_t(x))}\left[ -p_t(\psi_t(x))\cdot\textrm{div}(u_t(\psi_t(x))) \right]\\ &=-\textrm{div}(u_t(\psi_t(x))) \end{aligned} \end{equation}\]O.D.E.
