概率路径

概率路径与条件概率路径

流匹配 Flow Matching 理论的起点是 概率路径 Probability Paths (图5第二行)。直观上,概率路径描述了一个随时间变化的概率分布序列 (pt)0t1{(p_t)}_{0 \le t \le 1}。它在单位时间内将简单的先验噪声分布平滑地转换为复杂的真实数据分布 。

在实际训练中,我们无法观测全局分布的演化——我们不知道pdatap_{data}的分布,更不知道怎么在pdatap_{data}pinitp_{init}中插值出一条路径。

但是我们可以对每一个具体的样本 zpdataz \sim p_{data} 来决定其从噪声演化到这一个取样出来的数据点的轨迹,这就是 条件概率路径 Conditional Probability Path (图5第一行)pt(xz)p_t(x|z) 。该路径必须满足特定的边界条件:

  1. t=0t=0 时,分布代表初始的噪声状态,通常设定为标准高斯分布 p0(z)=pinit=N(0,Id)p_0(\cdot|z) = p_{init} = \mathcal{N}(0, I_d)
  2. t=1t=1 时,分布必须塌缩到特定的数据点 zz 上,即 p1(z)=δzp_1(\cdot|z) = \delta_z,其中 δz\delta_z 是以 zz 为中心的 Dirac delta 分布,反映在图中就是一个确定的点。

如此,我们定义了从噪声到有序数据的演化起点和终点。

边缘概率路径

现在不考虑这条路径是如何构造的,假设我们已经通过某种手段“插值”出了这些条件路径 pt(xz)p_t(x |z),数学上很容易通过积分将其边缘化得到一个 边缘概率路径 Marginal Probability Path pt(x)p_t(x)

pt(x)=pt(xz)pdata(z)dzp_t(x) = \int p_t(x|z) p_{data}(z) dz

这个边缘分布代表的是在时间 tt 时刻,所有数据点 zz 条件分布的加权平均,其中权重由数据分布 pdata(z)p_{data}(z) 决定。放在神经网络中理解的话——我们可以通过不断学习每个数据点的条件路径来间接地学习整个边缘路径的演化规律。只要学习的数据足够多,就能近似得出整个数据的分布演化过程。

在任何时刻 tt,边缘分布 ptp_t 都是由各个数据点对应的条件分布组成的混合模型 。由于 pt(xz)p_t(x|z) 在边界上的定义,边缘路径自然地满足 p0=pinitp_0 = p_{init}p1=pdatap_1 = p_{data}。这意味着,只要我们能够找到一个驱动粒子沿着 ptp_t 演化的机制,就能实现从噪声生成数据的目标。

高斯条件概率路径

回到之前的问题——已经有了起点和终点,如何得到中间的路径呢?深度学习中最常用的路径形式是高斯路径:

pt(xz)=N(x;αtz,βt2Id)p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d)

其中 αt\alpha_tβt\beta_t 是时间 tt噪声调度器 Noise Schedulers 函数,满足边界条件 α0=0,β0=1\alpha_0 = 0, \beta_0 = 1α1=1,β1=0\alpha_1 = 1, \beta_1 = 0。这个高斯路径的意义在于,它通过线性变换和添加噪声的方式将数据点 zz 转换为一个在空间中逐渐收敛到 zz 的分布。下面列出几个常见的噪声调度器方案:

路径类型 αt\alpha_t βt\beta_t 几何特征
最优传输路径 tt 1t1-t 严格直线演化,速度恒定
方差保持路径 1σt2\sqrt{1-\sigma_t^2} σt\sigma_t 曲线演化,沿球面移动
方差爆炸路径 1 σ(t)\sigma(t) 均值不变,仅方差扩张

通过调整 αt\alpha_tβt\beta_t 的具体形式,流匹配可以灵活地模拟从直线到曲线的各种演化轨迹。其中,最优传输路径(相当于线性插值)因其能够诱导最直接、最平滑的生成轨迹而备受青睐 。

那么从条件概率路径 ptp_t 中采样的过程可以记作:

zpdataϵpinit=N(0,Id)    x=αtz+βtϵptz \sim p_{data} \quad \epsilon \sim p_{init} = \mathcal{N}(0, I_d) \implies x=\alpha_tz+\beta_t\epsilon \sim p_t

向量场

条件向量场

边缘概率路径定义了从初始分布到数据分布的演化过程,给出了XtX_t的变化过程。现在要做的是找到推动这种演化进行的向量场——或者说速度,也就是之前定义的uu

对每一个数据点zRz \in \mathbb{R} ,我们有一个 条件向量场 Conditional Vector Field。根据之前向量场的定义,很明显:

X0pinitddtXt=uttarget(Xtz)    Xtpt(z)(0t1)X_0 \sim p_{init} \quad \frac{d}{dt} X_t = u_t^{target}(X_t|z) \implies X_t \sim p_t(\cdot | z) \quad (0 \le t \le 1)

其中的uttarget(Xtz)u_t^{target}(X_t|z)就是我们要找的条件向量场。

X0pinitX_0 \sim p_{init} 是从噪声分布中采样的起点。对于给定的初始点,该方程定义的流 ψt(x0)\psi_t(x_0) 唯一确定了粒子在 tt 时刻的位置 。向量场、ODE 和流是描述同一个动力学系统的三种不同视角:向量场定义了规则,ODE 描述了局部变化,而流则给出了宏观的轨迹结果。

连续性方程:为什么向量场能够改变概率分布?这背后的基础理论是物理学中的连续性方程(Continuity Equation)。该方程描述了概率质量在向量场驱动下的守恒演化过程:

tpt(x)+div(ptut)(x)=0\frac{\partial}{\partial t}p_t(x) + \text{div}(p_t u_t)(x) = 0

这里的 div(ptut)\text{div}(p_t u_t) 表示概率流的散度。直观地理解,如果我们将概率密度想象成一片“云雾”,连续性方程规定了云雾密度的变化必须完全由其流动引起 。公式的左侧 tpt(x)\partial_t p_t(x) 是密度的瞬时变化率,右侧的负散度项则代表了该点概率质量的净流入。这意味着概率质量既不会凭空消失,也不会无中生有,而是被向量场 utu_t 连续地从一个区域“搬运”到另一个区域 。在生成式建模中,只要我们找到一个满足连续性方程的向量场 utu_t,那么随着 tt 从 0 演化到 1,初始的简单噪声分布 p0p_0 就会被精确地“塑形”为目标数据分布 p1p_1

对于高斯条件路径 pt(xz)=N(x;αtz,βt2Id)p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d),我们可以利用流图 ψt(x0z)=αtz+βtx0\psi_t(x_0|z) = \alpha_t z + \beta_t x_0 来显式地求出对应的条件向量场。根据定义:

ddtψt(x0z)=α˙tz+β˙tx0=uttarget(ψt(x0z)z)\frac{d}{dt} \psi_t(x_0|z) = \dot{\alpha}_t z + \dot{\beta}_t x_0 = u_t^{target}(\psi_t(x_0|z) | z)

通过代数变换,将 x0x_0 替换为 (xαtz)/βt(x - \alpha_t z) / \beta_t,我们得到条件向量场的闭式解:

uttarget(xz)=(α˙tβ˙tβtαt)z+β˙tβtxu_t^{target}(x|z) = (\dot{\alpha}_t - \frac{\dot{\beta}_t}{\beta_t} \alpha_t) z + \frac{\dot{\beta}_t}{\beta_t} x

该公式清晰地展示了速度矢量的构成:一部分是受数据点 zz 吸引产生的引力项,另一部分是由于分布尺度变化产生的漂移项 。在特定的路径设计下,如 CondOT 路径(αt=t,βt=1t\alpha_t=t, \beta_t=1-t),该公式可以进一步简化为极其直观的形式:ut(xz)=zϵu_t(x|z) = z - \epsilon,其中 ϵ\epsilon 是初始噪声 。

边缘向量场

流匹配引入了条件向量场 uttarget(xz)u_t^{target}(x|z) 的概念。对于每一个数据点 zz,我们可以轻松构造一个与之对应的条件向量场,使得其产生的轨迹服从条件路径 pt(xz)p_t(x|z)。如果我们对这些条件向量场进行某种形式的加权平均,就能得到正确的边缘向量场:

uttarget(x)=uttarget(xz)pt(xz)pdata(z)pt(x)dzu_t^{target}(x) = \int u_t^{target}(x|z) \frac{p_t(x|z) p_{data}(z)}{p_t(x)} dz


权重项 pt(xz)pdata(z)pt(x)\frac{p_t(x|z) p_{data}(z)}{p_t(x)} 实际上就是根据贝叶斯公式计算出的后验概率 pt(zx)p_t(z|x)。换句话说,边缘向量场在位置 xx 的速度,等于所有可能将粒子引向不同数据点 zz 的条件速度的期望值:

uttarget(x)=Ezpdata(x)[uttarget(xz)]u_t^{target}(x) = \mathbb{E}_{z \sim p_{data}(\cdot|x)} [u_t^{target}(x|z)]

这样理解还是不够直观。现在积分符号中有两个部分:一个uttarget(xz)u_t^{target}(x|z)——这是我们很容易通过概率路径得出来的条件向量场,也就是确定终点的时候,我们在这个时间点该向哪里走;而另一个部分pt(zx)p_t(z|x)则代表了在当前位置xx时,我们最终会被吸引到哪一个数据点zz的概率分布。

那么将这两项的乘积在dzdz上进行积分,就是对于当前的所在位置xx,前往不同目标zz在此时所需要的的向量场的加权平均,换言之就是给定位置xx时运动方向的期望——而所有的xx给出的uttargetu_t^{target},则构成了我们需要的边缘向量场,告诉一个数据点在任意一点应该往哪里“走”。

流匹配损失

定义概率路径和向量场后,接下来要考虑的是如何学习边缘向量场——也就是数据总体概率分布的演化“速度”。

要学习边缘向量场,首先要定义损失函数,即 流匹配损失 FM Loss 。理想的流匹配损失旨在使神经网络预测的向量场 vθ(x,t)v_\theta(x, t) 逼近真实的边缘向量场 uttarget(x)u_t^{target}(x)

LFM(θ)=EtUnif,xpt(x)[vθ(x,t)uttarget(x)2]\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t \sim \text{Unif}, x \sim p_t(x)} [\|v_\theta(x, t) - u_t^{target}(x)\|^2]

然而,由于 uttarget(x)u_t^{target}(x) 是不可知的(参考上一小节的uttarget(x)u_t^{target}(x)计算公式,里面需要用到一个pdata(z)p_{data}(z),我们当然无从得知),这一损失无法直接计算。所以将 条件流匹配损失 CFM Loss 作为其可计算的替代方案:

LCFM(θ)=Et,zpdata,xpt(xz)[vθ(x,t)uttarget(xz)2]\mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t, z \sim p_{data}, x \sim p_t(x|z)} [\|v_\theta(x, t) - u_t^{target}(x|z)\|^2]

这个式子中的条件向量场uttarget(xz)u_t^{target}(x|z)和条件路径pt(xz)p_t(x|z)是可以计算的。而数学上可以证明这两个损失函数对于参数 θ\theta 的梯度是完全等价的,即:

θLFM(θ)=θLCFM(θ)\nabla_\theta \mathcal{L}_{FM}(\theta) = \nabla_\theta \mathcal{L}_{CFM}(\theta)

这意味着,即便我们只是在强迫网络去匹配每一个孤立的条件速度矢量,在大数定律的作用下,网络最终会自发地学到所有这些矢量在边缘分布下的加权平均值,即真实的边缘向量场 。

FM/CFM 等价性证明

  1. 展开 FM 损失:

LFM(θ)=Et,x\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t, x}

其中第三项 uttarget2\|u_t^{target}\|^2 不包含 θ\theta,在优化时视为常数。
2. 处理交叉项:关键在于利用边缘化公式 uttarget(x)=uttarget(xz)pt(xz)pdata(z)pt(x)dzu_t^{target}(x) = \int u_t^{target}(x|z) \frac{p_t(x|z) p_{data}(z)}{p_t(x)} dz

Expt=pt(x)vθ(x,t)T(uttarget(xz)pt(xz)pdata(z)pt(x)dz)dx\mathbb{E}_{x \sim p_t} = \int p_t(x) v_\theta(x, t)^T \left( \int u_t^{target}(x|z) \frac{p_t(x|z) p_{data}(z)}{p_t(x)} dz \right) dx

=vθ(x,t)Tuttarget(xz)pt(xz)pdata(z)dxdz\dots = \int \int v_\theta(x, t)^T u_t^{target}(x|z) p_t(x|z) p_{data}(z) dx dz

这正好等价于 Ezpdata,xpt(xz)\mathbb{E}_{z \sim p_{data}, x \sim p_t(x|z)}
3. 合并回归目标: 同样的逻辑可以证明第一项 vθ2\|v_\theta\|^2 的期望在两组分布下也是相等的。因此,两个损失函数在关于 θ\theta 的一阶导数上具有完美的一致性 。

确定损失函数后,即可给出训练方法: