引言
最近看到Qwen-Edit，感觉效果特别帅。借此机会想系统学习一下Flow Matching的数学原理，这里记录一下。
后面内容完全是在翻译这篇论文：Flow Matching Guide and Code。感谢FAIR的大佬们，这种从1+1开始教的self-contained教程对我这种外行人真的非常重要。
PS：学物理的放过咱们cs吧，是真难懂啊。。。
一开始以为Flow Matching的作者 Ricky Tian Qi Chen 就是xgBoost的陈天奇，但发现似乎不是哈哈
一、大白话理解 (也没那么白。。。)
Flow Matching (FM) 在干什么？
- Flow Matching的终极任务是学习一个向量场 (vector field)；
- 这个向量场通过ODE的形式，定义了一个流 (flow) $\psi$；
- 一个流是定义在$\mathbb{R}^d$上的一系列可逆变换，这些变换是在时间上连续的；
- Flow Matching就是在学习一个流，将初始分布（一般是简单分布如$\mathcal{N}(0,1)$）中的某个样本$X_0\sim p$映射到目标分布上，即$X_1:=\psi(X_0)\sim q$。
二、FM的形式化定义
给定一个定义在$\mathbb{R}^d$上的目标分布$q$，以及从$q$中采样得到的训练集，我们的任务是学习一个模型，用来生成$q$上的一些新样本。
为了完成这个任务，Flow Matching搭建了一条概率路径${p_t}_{0\le t\le 1}$，其中$p_0=p$为某个已知的初始分布，$p_1=q$为目标分布。
更加具体地说，我们训练的是一个能够描述样本瞬时速度的神经网络。这个网络后面被用于沿着概率路径${p_t}$将初始分布转为目标分布。
在训练完成后，我们 (1) 从$p$中采样新样本 (2) 求解向量场所描述的ODE，来得到分布$q$中的一个新样本。
在形式上，向量场是一个依赖于时间的函数$u:[0,1]\times\mathbb{R}^d\mapsto\mathbb{R}^d$。这个向量场用公式$(1)$中的ODE来定义一个流$\psi:[0,1]\times\mathbb{R}^d\mapsto\mathbb{R}^d$：
\[\begin{equation} \frac{\mathrm{d}}{\mathrm{d}t}\psi(t,x)=u(t,\psi(t,x)) \end{equation}\]其中$\psi(0,x)=x$。
我们称向量场$u$能够产生一条概率路径${p_t}$，当且仅当对应的流$\psi$满足：
\[\begin{equation} X_t:=\psi(t,X_0)\sim p_t \end{equation}\]其中$X_0\sim p_0$。
三、FM的步骤
- 第一步：从目标分布$q$中收集一些样本作为训练集。
- 第二步：设计一条连续时间的概率路径${p_t}$，满足$p_0=p$且$p_1=q$。
- 第三步：使用回归 (Regression) 的方式，训练一个参数为$\theta$的向量场$u^\theta$。
- 第四步：从$p$中采样一个新样本$X_0\sim p$，根据向量场所定义的概率路径，得到$q$上的一个样本$X_1=u^\theta(X_0)\sim q$。
四、实际场景下
假设初始分布为$p=\mathcal{N}(0,1)$，则概率路径${p_t}$由公式$(3)$定义：
\[\begin{equation} \begin{aligned} p_t(x)&=\int p_{t\mid 1}(x\mid x_1)q(x_1)\mathrm{d}x_1\\ p_{t\mid 1}(x\mid x_1)&=\mathcal{N}(tx_1,(1-t)^2) \end{aligned} \end{equation}\]这个路径也被称为条件最优运输 (conditional optimal transport)。
我们可以通过线性插值，定义从$X_0\sim p$到$X_1\sim q$的变换过程中任意时刻$t$下的样本$X_t\sim p_t$：
\[\begin{equation} X_t=tX_1+(1-t)X_0 \end{equation}\]在训练时，我们的目标是去最小化当前的向量场$u^\theta$与目标向量场$u$之间的距离，即：
\[\begin{equation} \mathcal{L}_{FM}=\mathbb{E}_{t,X_t}\left\| u^\theta(t,X_t)-u(t,X_t) \right\|^2 \end{equation}\]其中$t\sim\mathcal{U}[0,1]$。
然而这个loss非常难实现，因为目标向量场$u$是定义在两个随机向量的联合分布上。 但当我们只关注某一个目标样本$x_1$时，问题就从联合概率变为一个边缘概率，也就变得可解了。
具体来说，我们需要求解下面这个ODE：
\[\begin{equation} \frac{\mathrm{d}}{\mathrm{d}t}X_{t\mid1}=u(t,X_{t\mid1}\mid x_1) \end{equation}\]其中，
\[\begin{equation} X_{t\mid1}=tx_1+(1-t)X_0\sim \mathcal{N}(tx_1,(1-t)^2) \end{equation}\]这个ODE的解为如下的条件向量场（证明见附录）：
\[\begin{equation} u(t,X_{t\mid1}\mid x_1)=\frac{x_1-X_{t\mid1}}{1-t} \end{equation}\]我们定义如下的条件FM损失：
\[\begin{equation} \mathcal{L}_{CFM}=\mathbb{E}_{t,X_t,X_1}\left\| u^\theta(t,X_t)-u(t,X_t\mid X_1) \right\|^2 \end{equation}\]将ODE的解代入公式$(9)$，就得到了FM一个最简单的实现：
\[\begin{equation} \mathcal{L}_{CFM}^{OT,Gauss}=\mathbb{E}_{t,X_t,X_1}\left\| u^\theta(t,X_t)-(X_1-X_0) \right\|^2 \end{equation}\]其中$X_0\sim\mathcal{N}(0,1)$。
附录
公式(8)的证明
从公式$(6)$出发：
\[\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}X_{t\mid1} &=\frac{\mathrm{d}}{\mathrm{d}t}\left(tx_1+(1-t)X_0\right)\\ &=x_1-X_0\\ &=u(t,X_{t\mid1}\mid x_1) \end{aligned} \end{equation}\]由公式$(7)$做等价变形，我们有
\[\begin{equation} X_0=\frac{X_{t\mid 1}-tx_1}{1-t} \end{equation}\]代入上式得：
\[\begin{equation} \begin{aligned} u(t,X_{t\mid1}\mid x_1) &=x_1-\frac{X_{t\mid 1}-tx_1}{1-t}\\ &=\frac{x_1-X_{t\mid 1}}{1-t} \end{aligned} \end{equation}\]Q.E.D.
