概述


生成式模型本质是从已知分布采样并通过某种变换得到目标分布,核心是在已知分布的情况下,如何建模得到目标分布的过程。以扩散模型 DDPM 为例,其从已知的高斯分布 N(0,I)N(\boldsymbol{0},\boldsymbol{I}) 采样,通过逐渐的去噪得到目标分布。

流匹配模型(Flow Matching)由连续归一化流(Continuous Normalizing Flows, CNF)和Neural ODE演进而来,其被广泛应用于目前的生成式系统中,如Stable Diffusion 3、Meta MovieGen 等。Flow Matching可以看做是一种特殊的扩散模型,其数学原理也较为晦涩难懂,而讲解其原理这是本篇文章的核心解决目标之一。

在正式介绍 Flow Matching 模型之前,我们可参考下图,其介绍了目前几种不同的生成模型策略。可以看到,Flow Matching 模型类似于 VAE,通过构造一个可逆的变换 ff 来做生成和反推的工作,这个变换也被称为流(Flow)。

连续归一化流


预备知识

流,简单来说就是一种通过可逆变换将简单分布(比如高斯分布)映射到复杂数据分布的技术。而标准化流(Normalizing Flow)是流的一种重要类别,广泛应用于生成模型、密度估计等领域。

对于一个给定的简单分布变量 z∈Rnz\in\mathbb{R}^n,我们通过可逆变换 x=f(z)x=f(z) 得到新变量 x∈Rnx\in\mathbb{R}^n。根据概率论知识,可以求得 xx 的概率分布为

px(x)=pz(f−1(x))∣det⁡(∂f−1∂x)∣=pz(z)∣det⁡(∂f∂z)−1∣p_x(x)=p_z(f^{-1}(x))\left|\det\left(\frac{\partial f^{-1}}{\partial x}\right)\right|=p_z(z)\left|\det\left(\frac{\partial f}{\partial z}\right)^{-1}\right|

其中的 ∂f−1∂x=J\dfrac{\partial f^{-1}}{\partial x}=\boldsymbol{J} 是一个雅可比矩阵,且有 Jij=∂(f−1)i∂xjJ_{ij}=\dfrac{\partial (f^{-1})_i}{\partial x_j},当 det⁡(J)=1\det(\boldsymbol{J})=1 时,得到的全新分布将和初始分布 pz(z)p_z(z) 具有同样的规模。对两边取对数可以得到

log⁡px(x)=log⁡pz(f−1(x))+log⁡∣det⁡(J)∣\boxed{\log p_x(x)=\log p_z(f^{-1}(x))+\log|\det(\boldsymbol{J})|}

这一对数似然结果可以被看做两项

  • 变量 z=f−1(x)z=f^{-1}(x) 在先验下的对数概率
  • 变换在 xx 处的体积缩放对数,也就是 ff 对概率空间的影响

流的复合

单个可逆网络表达能力有限。归一化流的做法是将 ff 分解为 KK个简单可逆变换的复合。设 z0∼π(z)z_0\sim \pi(z) 为初始简单分布变量,经过 KK 步后到达数据空间

zK=fθK∘fθK−1∘⋯∘fθ1(z0)z_K=f_{\theta_K}\circ f_{\theta_{K-1}}\circ\cdots\circ f_{\theta_1}(z_0)

其中的每一步 fθkf_{\theta_k} 都是一个可逆且可微的神经网络 θk\theta_k,复合后有 x=zK=f(z0)x=z_K=f(z_0)。对于复合函数 f=fθK∘⋯∘fθ1f=f_{\theta_K}\circ\cdots\circ f_{\theta_1},其反函数的雅可比行列式由链式法则给出

∣det⁡(∂f−1∂x)∣=1∣det⁡(∂f∂z0)∣=∏k=1K1∣det⁡(∂fθk∂zk−1)∣\left|\det\left(\frac{\partial f^{-1}}{\partial x}\right)\right|=\dfrac{1}{\left|\det\left(\dfrac{\partial f}{\partial z_0}\right)\right|}=\prod_{k=1}^K\frac{1}{\left|\det\left(\dfrac{\partial f_{\theta_k}}{\partial z_{k-1}}\right)\right|}

将上述式子取对数同样可以得到

log⁡px(x)=log⁡π(z0)−∑k=1Klog⁡∣det⁡(∂fθk∂zk−1)∣\boxed{\log p_x(x)=\log\pi(z_0)-\sum_{k=1}^K\log\left|\det\left(\frac{\partial f_{\theta_k}}{\partial z_{k-1}}\right)\right|}

上述过程描述了一个离散形式的归一化流,通过构建神经网络 θk\theta_{k} 可以实现从简单分布 π(z0)\pi(z_0) 到复杂分布 px(x)p_x(x) 的演变。在归一化流中,可逆变换 fθkf_{\theta_k} 负责移动样本,雅可比行列式 det⁡(J)\det(\boldsymbol{J}) 负责计算概率分布密度的变化。

传统的归一化流通常是离散的变换,比如通过一系列精心设计的神经网络层(例如RealNVP或Glow)逐步完成从 z0z_0 到 xx 的映射。但是在每一步的求解中需要计算雅可比行列式,这带来了较大的计算复杂度,并且限制了流生成的概率分布多样性。

连续归一化流

连续归一化流(Continuous Normalizing Flow, CNF)将离散归一化流的逐层变换过程推广为连续时间上的演化。为理解这一推广,令 tk=kΔtt_k=k\Delta t,其中 Δt=1/K\Delta t=1/K,并考虑接近恒等映射的变换:

zk+1=zk+Δt vtk(zk;θ)z_{k+1}=z_k+\Delta t\,v_{t_k}(z_k;\theta)

其中的 vtk(zk;θ):Rn→Rnv_{t_k}(z_k;\theta):\mathbb R^n\to\mathbb R^n 是由神经网络参数化的时变向量场,也称为速度场。给定时刻 tt 和位置 zz,它输出样本在该位置的瞬时运动速度。这一神经网络可以由上式改写记为

vtk(zk;θ)=zk+1−zkΔtv_{t_k}(z_k;\theta)=\frac{z_{k+1}-z_k}{\Delta t}

在适当条件下,当 Δt→0\Delta t\to0 时,这一离散更新对应于下述常微分方程 ODE

dztdt=vt(zt;θ),z0∼π\frac{\mathrm{d}z_t}{\mathrm{d}t}=v_t(z_t;\theta),\quad z_0\sim\pi

因此,CNF 中的样本不再仅在有限个层之间跳转,而是沿着速度场确定的轨迹连续运动。这里的 t∈[0,1]t\in[0,1] 是生成过程的时间参数,也可以理解为连续化的网络深度。

流映射与速度场

为了区分样本当前位置和把初始位置送到当前位置的函数,定义流映射

ϕt:Rn→Rn\phi_t:\mathbb R^n\to\mathbb R^n

流映射函数满足 zt=ϕt(z0)z_t=\phi_t(z_0),代入 ODE 得到

ddtϕt(z0)=vt(ϕt(z0);θ),ϕ0(z0)=z0\frac{\mathrm{d}}{\mathrm{d}t}\phi_t(z_0)=v_t\bigl(\phi_t(z_0);\theta\bigr),\quad\phi_0(z_0)=z_0

这与前面关于 ztz_t 的 ODE 完全等价,但两者强调的对象不同,详见下表。

符号 含义
ztz_t 一个样本在时刻 tt 的位置
vt(x;θ)v_t(x;\theta) 时刻 tt、位置 xx 处的瞬时速度
ϕt(z0)\phi_t(z_0) 从初始位置 z0z_0 出发,运动到时刻 tt 后的位置

由于 t∈[0,1]t\in [0,1],因此完整映射 ff 在 CNF 中对应于终点流映射

f=ϕ1,x=ϕ1(z0)f=\phi_1,\quad x=\phi_1(z_0)

由于网络直接参数化的是 vtv_t,故 ϕt\phi_t 可以通过求解 ODE 得到隐式表达

ϕt(z0)=z0+∫0tvτ(ϕτ(z0);θ) dτ\phi_t(z_0)=z_0+\int_0^tv_{\tau}\bigl(\phi_{\tau}(z_0);\theta\bigr)\,\mathrm{d}\tau

由于积分中的速度依赖沿途位置,这通常不是一个可以直接计算的显式积分,而需要数值 ODE 求解器。但是可以证明:CNF 的可逆性来自 ODE 解 ϕt\phi_t 的存在性和唯一性,不要求神经网络 vt(x;θ)v_t(x;\theta) 是一个可逆函数。

样本的连续变化

初始样本服从先验分布 z0∼p0=πz_0\sim p_0=\pi,经过流映射后有

zt=ϕt(z0)∼ptz_t=\phi_t(z_0)\sim p_t

在整个流过程中,我们可以得到一条概率密度路径(一个映射)

p:[0,1]×Rn→R⩾0p:[0,1]\times\mathbb R^n\to\mathbb R_{\geqslant0}

上式表明,对于时变向量场 vt∈Rn(t∈[0,1])v_t\in\mathbb{R}^n(t\in[0,1]),映射 pp 把其映射为一个具有非负性的一维概率分布。根据概率的归一化性质,对于任意的 tt 都有

∫Rnpt(x) dx=1\int_{\mathbb R^n}p_t(x)\,\mathrm{d}x=1

流映射与密度路径之间的关系可以记为

pt=[ϕt]∗πp_t=[\phi_t]_*\pi

这表示从 p0=πp_0=\pi 采样,经过流映射 ϕt\phi_t 得到样本密度 ptp_t。再利用预备知识节的随机变量变换公式,可以得到

pt(zt)=p0(ϕt−1(zt))∣det⁡(∂ϕt−1∂zt)∣p_t(z_t)=p_0\bigl(\phi_t^{-1}(z_t)\bigr)\left|\det\left(\frac{\partial \phi_t^{-1}}{\partial z_t}\right)\right|

沿着轨迹 zt=ϕt(z0)z_t=\phi_t(z_0),由于 z0=ϕt−1(zt)z_0=\phi^{-1}_t(z_t),上式可以写成

pt(zt)=p0(z0)∣det⁡(∂ϕt∂z0)−1∣p_t(z_t)=p_0(z_0)\left|\det \left(\frac{\partial \phi_t}{\partial z_0}\right)^{-1}\right|

两边取对数有

log⁡pt(zt)=log⁡p0(z0)−log⁡∣det⁡(∂ϕt∂z0)∣\boxed{\log p_t(z_t)=\log p_0(z_0)-\log\left| \det \left(\frac{\partial \phi_t}{\partial z_0}\right)\right|}

似然估计

到这里可以看到,CNF 使用的仍然是普通归一化流的变量替换公式。现在考虑连续时间结构,我们进一步把雅可比行列式的变化写成微分方程 ODE。定义流映射对初始位置的雅可比矩阵

At=∂ϕt∂z0∈Rn×n\boldsymbol{A}_t=\frac{\partial \phi_t}{\partial z_0}\in\mathbb R^{n\times n}

因为 ϕ0(z0)=z0\phi_0(z_0)=z_0 是恒等映射,因此 A0=In\boldsymbol{A}_0=\boldsymbol{I}_n。雅可比矩阵对时间 tt 求导有

dAtdt=∂∂t[∂∂z0(z0+∫0tvτ(ϕτ(z0);θ) dτ)]=∂∂z0vt(ϕt(z0);θ)=∂vt(ϕt(z0);θ)∂ϕt(z0)⋅∂ϕt∂z0=∂vt∂ztAt=BAt\begin{aligned}\frac{\mathrm{d}A_t}{\mathrm{d}t}&=\frac{\partial}{\partial t}\left[\frac{\partial}{\partial z_0}\left(z_0+\int_0^tv_{\tau}\bigl(\phi_{\tau}(z_0);\theta\bigr)\,\mathrm{d}\tau\right)\right]\\&=\frac{\partial}{\partial z_0}v_t\bigl(\phi_t(z_0);\theta\bigl)\\&=\frac{\partial v_t\bigl(\phi_t(z_0);\theta\bigl)}{\partial \phi_t(z_0)}\cdot\frac{\partial\phi_t}{\partial z_0}\\&=\frac{\partial v_t}{\partial z_t}\boldsymbol{A}_t\\&=\boldsymbol{BA}_t\end{aligned}

其中的 B\boldsymbol{B} 也是一个雅可比矩阵。根据行列式的求导公式

ddtdet⁡(At)=det⁡(At)⋅tr(At−1dAtdt)\frac{\mathrm{d}}{\mathrm{d}t}\det(\boldsymbol{A}_t)=\det(\boldsymbol{A}_t)\cdot\mathrm{tr}\left(\boldsymbol{A}_t^{-1}\frac{\mathrm{d}\boldsymbol{A}_t}{\mathrm{d}t}\right)

可以得到下述求导结果

ddtlog⁡∣det⁡(At)∣=tr(At−1dAtdt)\frac{\mathrm{d}}{\mathrm{d}t}\log|\det(A_t)|=\mathrm{tr}\left(\boldsymbol{A}_t^{-1}\frac{\mathrm{d}\boldsymbol{A}_t}{\mathrm{d}t}\right)

代入 dAtdt\dfrac{\mathrm{d}\boldsymbol{A}_t}{\mathrm{d}t} 的化简结果有

ddtlog⁡∣det⁡At∣=tr⁡(At−1BAt)=tr⁡(B)=∑i=1n∂vti(zt;θ)∂zti=∇ztvt(zt;θ)\begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\log|\det \boldsymbol{A}_t|&=\operatorname{tr}\left(\boldsymbol{A}_t^{-1}\boldsymbol{BA}_t\right)\\&=\operatorname{tr}\left(\boldsymbol{B}\right)\\&=\sum_{i=1}^n\frac{\partial v_t^i(z_t;\theta)}{\partial z_t^i}\\&=\nabla_{z_t}v_t(z_t;\theta)\end{aligned}

其中的第二步使用了迹的循环不变性。该式表明:整个流映射的对数体积缩放率,等于当前位置处速度场雅可比矩阵的迹(速度场的散度)。再对上一小节对数概率似然公式求关于 tt 的导数,可以得到

ddtlog⁡pt(zt)=−∇zt⋅vt(zt;θ)\frac{\mathrm{d}}{\mathrm{d}t}\log p_t(z_t)=-\nabla_{z_t}\cdot v_t(z_t;\theta)

移项后可得重要结论

ddtlog⁡pt(zt)+∇ztvt(zt;θ)=0\boxed{\frac{\mathrm{d}}{\mathrm{d}t}\log p_t(z_t)+\nabla_{z_t}v_t(z_t;\theta)=0}

这就是 CNF 的瞬时变量替换等式,Neural ODE 论文将其作为连续流计算密度的核心结果。在直观解释该公式之前,先考虑进一步的推导,将上式对 t∈[0,1]t\in[0,1] 积分可得

log⁡p1(z1)=log⁡π(z0)−∫01∇ztvt(zt;θ) dt\log p_1(z_1)=\log\pi(z_0)- \int_0^1\nabla_{z_t} v_t(z_t;\theta)\,\mathrm{d}t

连续性方程

根据散度的定义,散度控制的是局部体积的瞬时变化,并且由瞬时变量替换等式有

∇ztvt>0⇒局部体积膨胀,沿轨迹密度下降\nabla_{z_t} v_t>0 \quad\Rightarrow\quad \text{局部体积膨胀,沿轨迹密度下降}

∇ztvt<0⇒局部体积收缩,沿轨迹密度上升\nabla_{z_t}v_t<0 \quad\Rightarrow\quad \text{局部体积收缩,沿轨迹密度上升}

而在瞬时变量替换等式中,我们对 log⁡pt(zt)\log p_t(z_t) 求了关于 tt 的全导数。事实上,我们之前提到过 zt=z(t)z_t=z(t),因此利用链式法则和偏导数可以得到下式

ddtlog⁡pt(zt)=[∂∂tlog⁡pt(zt)+∂∂ztlog⁡pt(zt)∂ϕt(z0)∂t]=[∂∂tlog⁡pt(zt)+∇ztlog⁡pt(zt)⋅vt(zt;θ)]\begin{aligned}\frac{\mathrm{d}}{\mathrm{d}t}\log p_t(z_t)&=\left[\frac{\partial}{\partial t}\log p_t(z_t)+\frac{\partial}{\partial z_t}\log p_t(z_t)\frac{\partial \phi_t(z_0)}{\partial t}\right]\\&=\left[\frac{\partial}{\partial t}\log p_t(z_t)+\nabla_{z_t}\log p_t(z_t)\cdot v_t(z_t;\theta)\right]\end{aligned}

代入瞬时变量替换公式,并在等式两边乘以 ptp_t,得到

∂pt(zt)∂t+vt(zt;θ)⋅∇ztpt(zt)+pt(zt)∇zt⋅vt(zt;θ)=0\frac{\partial p_t(z_t)}{\partial t}+v_t(z_t;\theta)\cdot\nabla_{z_t}p_t(z_t)+p_t(z_t)\nabla_{z_t}\cdot v_t(z_t;\theta)=0

由乘积求导公式,可以整理为一个更简洁的以变量 x,tx,t 为主的方程形式

∂pt(x)∂t+∇x(pt(x)vt(x;θ))=0\boxed{\frac{\partial p_t(x)}{\partial t}+\nabla_x\bigl(p_t(x)v_t(x;\theta)\bigr)=0}

这就是连续性方程,表达概率质量在速度场作用下的守恒。在适当的正则性及唯一性条件下,给定初始密度 p0=πp_0=\pi,这个方程刻画了速度场 vtv_t 所生成的密度路径。其物理含义为:流体中任意封闭体积内的质量变化率等于流入流出该空间的流体质量流量之差。

CNF 生成与训练


CNF 生成

现在加入参数标记,以区分模型产生的密度和之后人为指定的目标密度路径

dztdt=vt(zt;θ),ptθ=[ϕtθ]∗π\frac{\mathrm{d}z_t}{\mathrm{d}t}=v_t(z_t;\theta),\quad p_t^\theta=[\phi_t^\theta]_*\pi

设真实数据分布为 qq。生成建模的目标是 p1θ≈qp_1^\theta\approx q。现在考虑生成新样本,假设我们已经得到了一个训练好的神经网络 θ\theta,在生成样本时有

z0∼π,z1=ϕ1θ(z0)z_0\sim\pi,\quad z_1=\phi_1^\theta(z_0)

也就是从先验采样,我们不妨假设 π=N(z0;0,I)\pi=N(z_0;\boldsymbol{0},\boldsymbol{I}),之后通过求解带有初值的 ODE 方程可以得到输出 x=z1x=z_1。求解 ODE 的过程可以近似看做一个欧拉近似迭代过程,考虑

zt+Δt≈zt+Δtvt(zt;θ)z_{t+\Delta t}\approx z_t+\Delta tv_{t}(z_t;\theta)

每一步的迭代都发生了三件事

  1. 将当前位置 ztz_t 和时间 tt 输入网络 θ\theta
  2. 网络输出当前的速度 vtv_t
  3. 按照 zt+Δt≈zt+Δtvt(zt;θ)z_{t+\Delta t}\approx z_t+\Delta tv_{t}(z_t;\theta) 更新位置

CNF 似然训练

在生成之前,网络 θ\theta 还没有训练好。随机初始化的速度场虽然也能推动样本运动,但终点分布 p1θp_1^{\theta} 通常不像真实数据 x∼qx\sim q。因此我们需要一个损失函数来评价参数 θ\theta,然后通过梯度下降更新它。对于似然训练,我们想要最小化负对数似然估计

L=−Ex∼q[log⁡p1θ(x)]\mathcal{L}=-\mathbb{E}_{x\sim q}[\log p_1^{\theta}(x)]

根据前面的分析有

log⁡p1θ(x)=log⁡π(z0)−∫01∇ztvt(zt;θ) dt\log p^{\theta}_1(x)=\log\pi(z_0)- \int_0^1\nabla_{z_t} v_t(z_t;\theta)\,\mathrm{d}t

训练数据是大量来自先验分布 qq 的真实样本 x(i)x^{(i)},但是我们并不知道其来自于哪一个起点变量 z0z_0,因此需要使用流映射的可逆性求解该样本的起点位置 z0(i)z_0^{(i)}。反向求解起点时,我们只需要知道当前状态的速度场 vt(θ)v_{t}(\theta) 就可以通过对 ODE 的 t:1→0t:1\rightarrow 0 积分就可以得到 z0(i)z_0^{(i)}。

在这个反向过程中,我们能够通过对速度场的散度积分,同时得到 log⁡p1θ(x(i))\log p_1^{\theta}(x^{(i)}) 计算所需的两项。此时我们就完成了样本 x(i)x^{(i)} 对应的损失计算,即负对数概率估计

l(x(i);θ)=−log⁡p1θ(x(i))l(x^{(i)};\theta)=-\log p_1^{\theta}(x^{(i)})

L=∑x(i)∼ql(x(i);θ)\mathcal{L}=\sum_{x^{(i)}\sim q}l(x^{(i)};\theta)

接着就能通过梯度方法进行参数优化了。

流匹配


训练目标

Flow Matching 的核心是:构造一条从噪声分布到数据分布的目标路径,为路径上的位置提供应该怎样运动的速度监督,再训练 CNF 的速度场。简单来说,其并没有采用似然训练方法,但仍然遵循 CNF 框架。

假设已知构造的目标路径为

p0=π∼N(0,I),p1≈qp_0=\pi\sim N(\boldsymbol{0},\boldsymbol{I}),\quad p_1\approx q

其对应的速度场为 utu_t,其需要满足 CNF 的连续性方程

∂pt(xt)∂t+∇x(pt(xt)ut(xt))=0\frac{\partial p_t(x_t)}{\partial t}+\nabla_x\bigl(p_t(x_t)u_t(x_t)\bigr)=0

现在使用神经网络构造出速度场 vt(xt;θ)v_t(x_t;\theta) 以拟合 ut(xt)u_t(x_t),有训练损失函数

LFM=Et∼U[0,1],pt(xt)[∥vt(xt;θ)−ut(xt)∥2]\mathcal{L}_{\mathrm{FM}}=\mathbb{E}_{t\sim U[0,1],p_t(x_t)}\left[\|v_t(x_t;\theta)-u_t(x_t)\|^2\right]

因此,该损失函数 LFM\mathcal{L}_{\mathrm{FM}} 衡量了当前速度场和理想速度场的区别。

条件速度场

训练目标看起来很美好,但问题在于我们怎么求得 ut(xt)u_t(x_t) 呢?虽然我们不知道 pt(xt)p_t(x_t) 和 ut(xt)u_t(x_t) 的具体形式,但我们可以通过添加条件将其转换为条件概率路径和条件速度场。对于给定的已知样本 x1x_1,我们可以考虑构造 ut(xt∣x1)u_t(x_t|x_1),最后通过这个向量场生成对应的条件概率路径 pt(xt∣x1)p_t(x_t|x_1)。

设采样 x1∼q(x1)x_1\sim q(x_1),对于边缘概率路径有

pt(xt)=∫pt(xt∣x1)q(x1) dx1p_t(x_t)=\int p_t(x_t|x_1)q(x_1)\,\mathrm{d}x_1

两边对 tt 求导得到

ddtpt(xt)=∫ddtpt(xt∣x1)q(x1) dx1\frac{\mathrm{d}}{\mathrm{d}t}p_t(x_t)=\int \frac{\mathrm{d}}{\mathrm{d}t}p_t(x_t|x_1)q(x_1)\,\mathrm{d}x_1

左端应用利用瞬时变量替换等式,右端对 p(xt∣x1)p(x_t|x_1) 应用连续性方程并调换积分顺序有

−∇xt(ut(xt)pt(xt))=−∇xt(∫ut(xt∣x1)pt(xt∣x1)q(x1) dx1)-\nabla_{x_t}(u_t(x_t)p_t(x_t))=-\nabla_{x_t}\left(\int u_t(x_t|x_1)p_t(x_t|x_1)q(x_1)\,\mathrm{d}x_1\right)

因此两式的被求散度式相等,移项可以得到

ut(xt)=∫ut(xt∣x1)pt(xt∣x1)q(x1)pt(xt) dx1\boxed{u_t(x_t)=\int u_t(x_t|x_1)\frac{p_t(x_t|x_1)q(x_1)}{p_t(x_t)}\,\mathrm{d}x_1}

上式表明了 ut(xt)u_t(x_t) 和条件构造 ut(xt∣x1),pt(xt∣x1)u_t(x_t|x_1),p_t(x_t|x_1) 的关系,但这个积分也不容易求解。作者在原文中给出了另外一个损失函数,其衡量了条件概率路径下,当前速度场和理想条件速度场的区别。

LCFM=Et∼U[0,1],q(x1),pt(xt∣x1)[∥vt(xt;θ)−ut(xt∣x1)∥2]\mathcal{L}_{\mathrm{CFM}}=\mathbb{E}_{t\sim U[0,1],q(x_1),p_t(x_t|x_1)}\left[\|v_t(x_t;\theta)-u_t(x_t|x_1)\|^2\right]

设对于所有的 xt∈Rnx_t\in\mathbb{R}^n 和 t∈[0,1]t\in[0,1] 都有 pt(xt)>0p_t(x_t)>0,可以证明

∇θLCFM=∇θLFM\nabla_{\theta}\mathcal{L}_{\mathrm{CFM}}=\nabla_{\theta}\mathcal{L}_{\mathrm{FM}}

这样,我们的训练就不再依赖于一个抽象的边缘向量场 vt(xt;θ)v_{t}(x_t;\theta),而是依赖于条件向量场 vt(xt∣x1)v_t(x_t|x_1)。这样我们就可以利用一定的训练数据对模型进行训练。

条件概率路径

在求 ut(xt∣x1)u_t(x_t|x_1) 的式子中,我们仍然不知道具体的条件概率路径 pt(xt∣x1)p_t(x_t|x_1)。类似于 DDPM,我们需要定义具体的前向过程 q(xt∣xt−1)q(x_{t}|x_{t-1}) 才能基于该过程进行训练。这里,作者定义了一个条件概率路径的高斯分布形式

pt(xt∣x1)=N(xt;μt(x1),σt2(x1)I)p_t(x_t|x_1)=N(x_t;\mu_t(x_1),\sigma^2_t(x_1)\boldsymbol{I})

为了使这个条件概率路径有较为良好的性质,设边界条件

p0(xt=0∣x1)=N(x;0,I),p1(xt=1∣x1)=N(x;x1,ε2I)p_0(x_{t=0}|x_1)=N(x;\boldsymbol{0},\boldsymbol{I}),\quad p_1(x_{t=1}|x_1)=N(x;x_1,\varepsilon^2\boldsymbol{I})

对于一个标准高斯分布的噪音起点 x0∼N(0,I)x_0\sim N(\boldsymbol{0},\boldsymbol{I}),定义流映射

xt=ψt(x0∣x1)=σt(x1)x0+μt(x1)∼N(xt;μt(x1),σt2(x1)I)x_t=\psi_t(x_0|x_1)=\sigma_t(x_1)x_0+\mu_t(x_1)\sim N\bigr(x_t;\mu_t(x_1),\sigma_t^2(x_1)\boldsymbol{I}\bigr)

代入 CNF 的 ODE 方程有

ddtψt(x0∣x1)=ut(ψt(x0∣x1)∣x1)\frac{\mathrm{d}}{\mathrm{d}t}\psi_t(x_0|x_1)=u_t\bigr(\psi_t(x_0|x_1)|x_1\bigr)

此时固定 x0,x1x_0,x_1,对上述映射求关于 tt 的导数有

x˙t=ut(xt∣x1)=σ˙t(x1)x0+μ˙t(x1)\dot{x}_t=u_t(x_t|x_1)=\dot{\sigma}_t(x_1)x_0+\dot{\mu}_t(x_1)

上式告诉了 xtx_t 的下一步运动方向,亦即条件速度场 ut(xt∣x1)u_t(x_t|x_1),但速度场的输入应当是当前位置 xtx_t 而非初始位置 x0x_0。因此上式中的 x0x_0 应当被消去,根据流映射的定义和可逆性有 x0=xt−μt(x1)σt(x1)x_0=\dfrac{x_t-\mu_t(x_1)}{\sigma_t(x_1)},因此可以得到

ut(xt∣x1)=σ˙t(x1)σt(x1)(xt−μt(x1))+μ˙t(x1)u_t(x_t|x_1)=\frac{\dot{\sigma}_t(x_1)}{\sigma_t(x_1)}(x_t-\mu_t(x_1))+\dot{\mu}_t(x_1)

现在来确定 μt\mu_t 和 σt\sigma_t 的形式,选择线性插值形式

μt=tx1,σt=1−(1−ε)t,0<ε<1\mu_t=tx_1,\quad \sigma_t=1-(1-\varepsilon)t,\quad 0<\varepsilon<1

代入条件速度场公式有

ut(xt∣x1)=x1−(1−ε)xt1−(1−ε)t,t∈[0,1]\boxed{u_t(x_t|x_1)=\frac{x_1-(1-\varepsilon)x_t}{1-(1-\varepsilon)t},\quad t\in[0,1]}

采用带有 x0,x1x_0,x_1 的表达形式,有

xt=[1−(1−ε)t]x0+tx1,ut(xt∣x0,x1)=x1−(1−ε)x0x_t=[1-(1-\varepsilon)t]x_0+tx_1,\quad u_t(x_t|x_0,x_1)=x_1-(1-\varepsilon)x_0

这样,我们就使用一对 (x0,x1)(x_0,x_1) 构造了一个直线速度条件场。

虽然我们使用了线性插值,但是用大量噪声与数据之间的线性插值,使得我们容易计算条件速度标签;通过条件速度回归,间接训练神经网络拟合整体目标速度场。这个整体速度场通常是非线性的,生成轨迹也因此可以是曲线。

批量训练

现在让真实样本随机变化并取样,设 X0,X1X_0,X_1 相互独立且

X1∼q,X0∼N(0,I)=πX_1\sim q,\quad X_0\sim N(\boldsymbol{0},\boldsymbol{I})=\pi

假设批次取样有 B\mathcal{B} 个样本对 (x0(i),x1(i))(x_0^{(i)},x_1^{(i)}),对每一对分别采样时间 tit_i,构造中间位置

xti(i)=[1−(1−ε)ti]x0(i)+tix1(i)x_{t_i}^{(i)}=[1-(1-\varepsilon)t_i]x_0^{(i)}+t_ix_1^{(i)}

这个 xti(i)x_{t_i}^{(i)} 将作为网络的输入,而根据我们对流映射的定义可以计算出

xt=1(i)=x1(i)+εx0(t)x_{t=1}^{(i)}=x_1^{(i)}+\varepsilon x_0^{(t)}

这一对样本的条件速度标签由构造的条件速度场给出

uti(xti(i)∣x1(i))=x1(i)−(1−ε)x0(i)=yiu_{t_i}(x_{t_i}^{(i)}|x_1^{(i)})=x_1^{(i)}-(1-\varepsilon)x_0^{(i)}=y_i

这是我们期望得到的条件速度,而神经网络的输出为

y^i=vti(xti(i);θ)\hat{y}_i=v_{t_i}(x_{t_i}^{(i)};\theta)

根据损失函数的定义,可以得到

LCFM=1B∑i=1B∥y^i−yi∥2\mathcal{L}_{\mathrm{CFM}}=\frac{1}{\mathcal{B}}\sum_{i=1}^{\mathcal{B}}\|\hat{y}_i-y_i\|^2

然后就可以使用梯度更新法实现参数优化了。

总结

Flow Matching 的核心在于:构造从噪声到数据的简单路径,用路径的运动速度监督神经网络。

  • 训练时:随机采样噪声 x0x_0、真实数据 x1x_1 和时间 tt,通过线性插值直接计算中间位置 xtx_t 及速度标签 ut(xt∣x1)u_t(x_t|x_1),让网络根据时间和当前位置预测速度。大量样本的回归使网络学到推动整体分布变化的速度场,而无需在训练中求解 ODE。
  • 生成时:从新噪声 x0x_0 出发,沿学到的速度场求解 ODE,逐步得到数据样本。

没错,尽管 Flow Matching 涉及到大量的数学推导,但最终形式却异常简单(^▽^)。