前言

学习 MIT S.184 课程,本篇文章记录课程第三讲 A 部分:分数函数和分数匹配,记录下个人学习笔记,和大家一起分享交流😄

Website:https://diffusion.csail.mit.edu/

Course Notes:An Introduction to Flow Matching and Diffusion Models

1. Reminder

今天我们将讨论两个重要主题:Score MatchingGuidance

在正式进入新内容之前,我们先快速回顾上一节课的核心结论。

首先,我们定义了 条件概率路径(conditional probability path)

p t ( ⋅ ∣ z ) p_t(\cdot\mid z) pt(z)

它描述的是:给定一个数据点 z z z,如何从初始噪声分布逐渐过渡到这个数据点。最常见的例子是高斯概率路径:

p t ( ⋅ ∣ z ) = N ( α t z , β t 2 I d ) p_t(\cdot\mid z) = \mathcal{N}(\alpha_t z,\beta_t^2I_d) pt(z)=N(αtz,βt2Id)

其中, α t \alpha_t αt β t \beta_t βt 是我们选择的调度器。 α t \alpha_t αt 控制数据成分, β t \beta_t βt 控制噪声成分。

同时,我们还定义了对应的 条件向量场(conditional vector field)

u t t a r g e t ( x ∣ z ) u_t^{\mathrm{target}}(x\mid z) uttarget(xz)

它的作用是让 ODE 的演化遵循条件概率路径。对于高斯概率路径,这个条件向量场可以写成 z z z x x x 的线性组合,其系数由 α t , β t \alpha_t,\beta_t αt,βt 以及它们的导数决定。

接着,我们通过对数据点 z z z 进行边缘化,得到了 边缘概率路径(marginal probability path)

p t ( x ) = ∫ p t ( x ∣ z ) p d a t a ( z ) d z p_t(x)=\int p_t(x\mid z)p_{\mathrm{data}}(z)dz pt(x)=pt(xz)pdata(z)dz

边缘概率路径描述的是从初始噪声分布到整个数据分布之间的插值。它满足:

p 0 = p i n i t , p 1 = p d a t a p_0=p_{\mathrm{init}}, \qquad p_1=p_{\mathrm{data}} p0=pinit,p1=pdata

然后,我们进一步构造了 边缘向量场(marginal vector field)

u t t a r g e t ( x ) = ∫ u t t a r g e t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z u_t^{\mathrm{target}}(x) = \int u_t^{\mathrm{target}}(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz uttarget(x)=uttarget(xz)pt(x)pt(xz)pdata(z)dz

这个公式可以理解为:在当前位置 x x x 处,根据后验分布对所有可能数据点对应的条件向量场做加权平均。

上一节课的关键结论是:如果我们从初始分布出发,沿着这个边缘向量场模拟 ODE,那么样本分布会遵循边缘概率路径,因此最终会到达数据分布:

X 0 ∼ p i n i t , d d t X t = u t t a r g e t ( X t ) ⇒ X 1 ∼ p d a t a X_0\sim p_{\mathrm{init}}, \qquad \frac{d}{dt}X_t=u_t^{\mathrm{target}}(X_t) \quad\Rightarrow\quad X_1\sim p_{\mathrm{data}} X0pinit,dtdXt=uttarget(Xt)X1pdata

因此,边缘向量场就是我们真正希望神经网络学习的对象。

上一节课还推导了 Flow Matching 的训练算法。

训练时,我们并不直接回归不可计算的边缘向量场,而是回归可计算的条件向量场。具体来说,每次训练时:

  1. 从数据集中采样一个数据点 z z z
  2. 随机采样一个时间 t t t
  3. 从条件概率路径中采样一个带噪样本 x x x
  4. 让神经网络输出 u t θ ( x ) u_t^\theta(x) utθ(x)
  5. 用条件向量场 u t t a r g e t ( x ∣ z ) u_t^{\mathrm{target}}(x\mid z) uttarget(xz) 作为监督目标进行回归。

对应损失为:

L C F M ( θ ) = E t , z , x [ ∥ u t θ ( x ) − u t t a r g e t ( x ∣ z ) ∥ 2 ] \mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t,z,x} \left[ \left \| u_t^\theta(x) - u_t^{\mathrm{target}}(x\mid z) \right \|^2 \right] LCFM(θ)=Et,z,x[ utθ(x)uttarget(xz) 2]

虽然训练时回归的是条件向量场,但由于 Conditional Flow Matching Loss 和真正的 Flow Matching Loss 只相差一个与参数无关的常数,所以二者具有相同的梯度和最小化器。因此,最小化条件流匹配损失就能学习到边缘向量场。

训练完成后,采样过程就回到 Lecture 1 中讲过的 ODE 模拟。

我们从初始噪声分布中采样:

X 0 ∼ p i n i t X_0\sim p_{\mathrm{init}} X0pinit

然后使用训练好的神经网络向量场模拟 ODE:

d d t X t = u t θ ( X t ) \frac{d}{dt}X_t=u_t^\theta(X_t) dtdXt=utθ(Xt)

最终返回终点 X 1 X_1 X1 。如果模型训练得足够好,那么这个终点就应该是来自数据分布的样本。

到这里为止,我们已经知道如何训练并使用一个流模型生成样本。

2. Score Function

接下来,我们从一个新的角度来理解这些模型:score 函数

Score Matching 可以看作是对前面 Flow Matching 框架的另一种视角。它不仅能帮助我们重新理解已经学过的内容,还能进一步把模型扩展到 SDE,也就是扩散模型。

在 Lecture 1 中,我们讨论过扩散模型。它们不是模拟确定性的 ODE,而是模拟带随机性的 SDE。今天引入的 score 函数,将帮助我们理解如何在 SDE 框架下进行生成,并且也会为本节课后半部分的 Guidance 做准备。Guidance 关心的是如何根据 prompt 等条件来控制生成结果。

那么,什么是 score 函数?

假设我们有一个概率分布 q ( x ) q(x) q(x) ,它的对数似然为:

log ⁡ q ( x ) \log q(x) logq(x)

Score 函数定义为对数似然关于 x x x 的梯度:

∇ x log ⁡ q ( x ) \nabla_x \log q(x) xlogq(x)

直观上,score 函数告诉我们:在当前位置 x x x ,如果想让对数似然增加得最快,应该朝哪个方向移动。也就是说,score 函数指向概率密度上升最快的方向。

左图展示的是某个分布的对数似然 log ⁡ q ( x ) \log q(x) logq(x) ,右图中的黑色箭头则表示对应的 score 函数 ∇ x log ⁡ q ( x ) \nabla_x\log q(x) xlogq(x) 。在概率密度较低的位置,箭头通常会指向更高密度区域;在高密度区域附近,箭头会指向局部峰值附近。

目前这只是一个概率论中的概念,还没有直接涉及机器学习。接下来我们会将它应用到前面定义的概率路径上。

在我们的设定中,有两类概率路径:

  • 条件概率路径:

p t ( x ∣ z ) p_t(x\mid z) pt(xz)

  • 边缘概率路径:

p t ( x ) p_t(x) pt(x)

因此,也可以定义两类 score 函数。

第一类是 条件 score 函数(conditional score function)

∇ x log ⁡ p t ( x ∣ z ) \nabla_x \log p_t(x\mid z) xlogpt(xz)

这里需要注意,梯度是对 x x x 求的。因为 p t ( x ∣ z ) p_t(x\mid z) pt(xz) 同时涉及 x x x z z z ,但 z z z 是给定的数据点,而我们关心的是当前位置 x x x 的变化方向。

第二类是 边缘 score 函数(marginal score function)

∇ x log ⁡ p t ( x ) \nabla_x \log p_t(x) xlogpt(x)

它描述的是边缘概率路径在当前位置 x x x 处的对数密度梯度。

接下来我们关心一个问题:

边缘 score 和条件 score 之间有什么关系?

答案和上一节课中边缘向量场与条件向量场之间的关系非常类似。边缘 score 可以写成条件 score 关于后验分布的加权平均:

∇ x log ⁡ p t ( x ) = ∫ ∇ x log ⁡ p t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z \nabla_x \log p_t(x) = \int \nabla_x \log p_t(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz xlogpt(x)=xlogpt(xz)pt(x)pt(xz)pdata(z)dz

其中:

p t ( x ∣ z ) p d a t a ( z ) p t ( x ) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } pt(x)pt(xz)pdata(z)

正是后验分布:

p t ( z ∣ x ) p_t(z\mid x) pt(zx)

因此,上式也可以写成:

∇ x log ⁡ p t ( x ) = E z ∼ p t ( z ∣ x ) [ ∇ x log ⁡ p t ( x ∣ z ) ] \nabla_x \log p_t(x) = \mathbb{E}_{z\sim p_t(z\mid x)} \left[ \nabla_x \log p_t(x\mid z) \right] xlogpt(x)=Ezpt(zx)[xlogpt(xz)]

直观地说,给定当前位置 x x x ,它可能来自许多不同的数据点 z z z 。边缘 score 就是在这些可能的来源数据点上,对条件 score 做后验加权平均。

这个结构与边缘向量场非常类似:

u t t a r g e t ( x ) = E z ∼ p t ( z ∣ x ) [ u t t a r g e t ( x ∣ z ) ] u_t^{\mathrm{target}}(x) = \mathbb{E}_{z\sim p_t(z\mid x)} \left[ u_t^{\mathrm{target}}(x\mid z) \right] uttarget(x)=Ezpt(zx)[uttarget(xz)]

区别只是:一个平均的是条件向量场,另一个平均的是条件 score。

下面简单证明这个关系。

从边缘 score 的定义开始:

∇ x log ⁡ p t ( x ) = ∇ x p t ( x ) p t ( x ) \nabla_x \log p_t(x) = \frac{\nabla_x p_t(x)}{p_t(x)} xlogpt(x)=pt(x)xpt(x)

而边缘概率路径的密度为:

p t ( x ) = ∫ p t ( x ∣ z ) p d a t a ( z ) d z p_t(x) = \int p_t(x\mid z)p_{\mathrm{data}}(z)dz pt(x)=pt(xz)pdata(z)dz

因此:

∇ x log ⁡ p t ( x ) = ∇ x p t ( x ) p t ( x ) = ∇ x ∫ p t ( x ∣ z ) p d a t a ( z ) d z p t ( x ) = ∫ ∇ x p t ( x ∣ z ) p d a t a ( z ) d z p t ( x ) \nabla_x \log p_t(x) = \frac{ \nabla_x p_t(x) }{ p_t(x) } = \frac{ \nabla_x \int p_t(x\mid z)p_{\mathrm{data}}(z)dz }{ p_t(x) } = \frac{ \int \nabla_x p_t(x\mid z)p_{\mathrm{data}}(z)dz }{ p_t(x) } xlogpt(x)=pt(x)xpt(x)=pt(x)xpt(xz)pdata(z)dz=pt(x)xpt(xz)pdata(z)dz

接着使用恒等式:

∇ x p t ( x ∣ z ) = p t ( x ∣ z ) ∇ x log ⁡ p t ( x ∣ z ) \nabla_x p_t(x\mid z) = p_t(x\mid z)\nabla_x\log p_t(x\mid z) xpt(xz)=pt(xz)xlogpt(xz)

于是得到:

∇ x log ⁡ p t ( x ) = ∫ p t ( x ∣ z ) ∇ x log ⁡ p t ( x ∣ z ) p d a t a ( z ) d z p t ( x ) = ∫ ∇ x log ⁡ p t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z \nabla_x \log p_t(x) = \frac{ \int p_t(x\mid z) \nabla_x\log p_t(x\mid z) p_{\mathrm{data}}(z)dz }{ p_t(x) } = \int \nabla_x\log p_t(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz xlogpt(x)=pt(x)pt(xz)xlogpt(xz)pdata(z)dz=xlogpt(xz)pt(x)pt(xz)pdata(z)dz

这就证明了边缘 score 是条件 score 的后验加权平均。

下面看一个具体例子:高斯概率路径的条件 score。

回顾高斯条件概率路径:

p t ( x ∣ z ) = N ( x ; α t z , β t 2 I d ) p_t(x\mid z) = \mathcal{N}(x;\alpha_t z,\beta_t^2 I_d) pt(xz)=N(x;αtz,βt2Id)

它的密度可以写成:

p t ( x ∣ z ) = 1 ( 2 π ) d / 2 β t d exp ⁡ ( − 1 2 β t 2 ∥ x − α t z ∥ 2 ) p_t(x\mid z) = \frac{1}{(2\pi)^{d/2}\beta_t^d} \exp \left( -\frac{1}{2\beta_t^2} \|x-\alpha_t z\|^2 \right) pt(xz)=(2π)d/2βtd1exp(2βt21xαtz2)

取对数得到:

log ⁡ p t ( x ∣ z ) = − d 2 log ⁡ ( 2 π ) − d log ⁡ β t − 1 2 β t 2 ∥ x − α t z ∥ 2 \log p_t(x\mid z) = -\frac{d}{2}\log(2\pi) -d\log\beta_t -\frac{1}{2\beta_t^2} \|x-\alpha_t z\|^2 logpt(xz)=2dlog(2π)dlogβt2βt21xαtz2

x x x 求梯度:

∇ x log ⁡ p t ( x ∣ z ) = − x − α t z β t 2 \nabla_x \log p_t(x\mid z) = -\frac{x-\alpha_t z}{\beta_t^2} xlogpt(xz)=βt2xαtz

也就是:

∇ x log ⁡ p t ( x ∣ z ) = − 1 β t 2 x + α t β t 2 z \nabla_x \log p_t(x\mid z) = -\frac{1}{\beta_t^2}x + \frac{\alpha_t}{\beta_t^2}z xlogpt(xz)=βt21x+βt2αtz

这个公式说明,高斯条件 score 同样是 x x x z z z 的线性组合。它指向的是让当前点 x x x 更接近条件高斯分布中心 α t z \alpha_tz αtz 的方向。

如果 x x x 离均值 α t z \alpha_t z αtz 很远,那么 score 的大小也会更大;如果 x x x 正好接近均值,那么 score 会接近零。

现在我们可以把条件概率路径、条件向量场和条件 score 放在一起比较。

条件概率路径:

p t ( ⋅ ∣ z ) = N ( α t z , β t 2 I d ) p_t(\cdot\mid z) = \mathcal{N}(\alpha_tz,\beta_t^2I_d) pt(z)=N(αtz,βt2Id)

描述的是从噪声分布到单个数据点 z z z 的插值。

条件向量场:

u t t a r g e t ( x ∣ z ) u_t^{\mathrm{target}}(x\mid z) uttarget(xz)

描述的是 ODE 应该如何运动,才能让样本分布遵循这条条件概率路径。

条件 score:

∇ x log ⁡ p t ( x ∣ z ) \nabla_x\log p_t(x\mid z) xlogpt(xz)

描述的是在当前时间 t t t 和位置 x x x 处,条件概率密度上升最快的方向。

对于高斯概率路径,条件向量场和条件 score 都是 x x x z z z 的线性组合,只是系数不同。

类似地,边缘概率路径、边缘向量场和边缘 score 之间也有对应关系。

边缘概率路径:

p t ( x ) = ∫ p t ( x ∣ z ) p d a t a ( z ) d z p_t(x) = \int p_t(x\mid z)p_{\mathrm{data}}(z)dz pt(x)=pt(xz)pdata(z)dz

描述的是从初始噪声分布到整个数据分布之间的插值。

边缘向量场:

u t t a r g e t ( x ) = ∫ u t t a r g e t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z u_t^{\mathrm{target}}(x) = \int u_t^{\mathrm{target}}(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz uttarget(x)=uttarget(xz)pt(x)pt(xz)pdata(z)dz

描述的是让 ODE 遵循边缘概率路径所需的速度场。

边缘 score:

∇ x log ⁡ p t ( x ) = ∫ ∇ x log ⁡ p t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z \nabla_x\log p_t(x) = \int \nabla_x\log p_t(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz xlogpt(x)=xlogpt(xz)pt(x)pt(xz)pdata(z)dz

描述的是边缘概率路径在当前位置 x x x 的对数密度梯度。

可以看到,边缘向量场和边缘 score 的公式结构非常相似:二者都是对条件对象进行后验加权平均。

接下来是一个非常重要的观察:对于高斯概率路径,向量场和 score 函数之间可以相互重参数化。

具体来说,我们可以找到两个只依赖于时间的系数 a t a_t at b t b_t bt ,使得条件向量场可以写成条件 score 的形式:

u t t a r g e t ( x ∣ z ) = a t ∇ x log ⁡ p t ( x ∣ z ) + b t x u_t^{\mathrm{target}}(x\mid z) = a_t\nabla_x\log p_t(x\mid z) + b_t x uttarget(xz)=atxlogpt(xz)+btx

其中:

a t = ( β t 2 α ˙ t α t − β ˙ t β t ) , b t = α ˙ t α t a_t= \left( \beta_t^2\frac{\dot{\alpha}_t}{\alpha_t} - \dot{\beta}_t\beta_t \right), \qquad b_t= \frac{\dot{\alpha}_t}{\alpha_t} at=(βt2αtα˙tβ˙tβt),bt=αtα˙t

同样地,边缘向量场也可以写成边缘 score 的形式:

u t t a r g e t ( x ) = a t ∇ x log ⁡ p t ( x ) + b t x u_t^{\mathrm{target}}(x) = a_t\nabla_x\log p_t(x) + b_t x uttarget(x)=atxlogpt(x)+btx

这个结论的证明本质上只是代数运算。

原因是,对于高斯概率路径,条件向量场和条件 score 都是 x x x z z z 的线性组合。因此,我们可以通过重新组合系数,把一个对象写成另一个对象的函数。

这个结果非常重要,因为它说明:

学习速度场和学习 score 函数,在高斯概率路径下本质上是等价的。

也就是说,我们可以选择训练神经网络预测向量场:

u t t a r g e t ( x ) u_t^{\mathrm{target}}(x) uttarget(x)

也可以选择训练神经网络预测 score:

∇ x log ⁡ p t ( x ) \nabla_x\log p_t(x) xlogpt(x)

只要知道调度器 α t , β t \alpha_t,\beta_t αt,βt ,就可以在二者之间进行转换。

这也是早期扩散模型的基本思想之一:它们通常不直接学习速度场,而是学习 score 函数。训练完成后,再通过公式把 score 转换成需要的向量场或采样动态。

因此,很多论文并不使用 “vector field” 这个语言,而是使用 “score function” 这个语言。但从这里可以看到,在高斯概率路径下,它们并不是完全不同的东西,而只是同一生成动力学的不同参数化方式。

3. Score Matching

现在我们正式讨论 Score Matching

在上一部分中,我们已经看到,对于高斯概率路径,速度场和 score 函数之间可以相互重参数化。因此,除了直接学习边缘向量场之外,我们也可以换一种方式:学习边缘 score 函数。

也就是说,我们不再训练一个向量场网络:

u t θ ( x ) u_t^\theta(x) utθ(x)

而是训练一个 score 网络:

s t θ ( x ) ∈ R d s_t^\theta(x)\in \mathbb{R}^d stθ(x)Rd

它的目标是近似边缘 score:

s t θ ( x ) ≈ ∇ x log ⁡ p t ( x ) s_t^\theta(x)\approx \nabla_x\log p_t(x) stθ(x)xlogpt(x)

这其实是在用 score 函数重新表述上一节课中已经做过的事情。

如果想直接学习边缘 score,最自然的做法是定义 Score Matching Loss

L S M ( θ ) = E t ∼ U n i f , z ∼ p d a t a , x ∼ p t ( ⋅ ∣ z ) [ ∥ s t θ ( x ) − ∇ x log ⁡ p t ( x ) ∥ 2 ] \mathcal{L}_{\mathrm{SM}}(\theta) = \mathbb{E}_{t\sim\mathrm{Unif},z\sim p_{\mathrm{data}},x\sim p_t(\cdot\mid z)} \left[ \left \| s_t^\theta(x) - \nabla_x\log p_t(x) \right \|^2 \right] LSM(θ)=EtUnif,zpdata,xpt(z)[ stθ(x)xlogpt(x) 2]

这里的采样过程与 Flow Matching 很类似:我们随机采样时间 t t t ,采样数据点 z z z ,再从条件概率路径中采样带噪样本 x x x

这个损失的目标很直接:让 score 网络 s t θ ( x ) s_t^\theta(x) stθ(x) 回归边缘 score:

∇ x log ⁡ p t ( x ) \nabla_x\log p_t(x) xlogpt(x)

不过问题在于,边缘 score 和边缘向量场一样,通常是不可直接计算的。因为 p t ( x ) p_t(x) pt(x) 是边缘概率路径,它涉及对未知数据分布的积分:

p t ( x ) = ∫ p t ( x ∣ z ) p d a t a ( z ) d z p_t(x)=\int p_t(x\mid z)p_{\mathrm{data}}(z)dz pt(x)=pt(xz)pdata(z)dz

所以我们无法直接得到:

∇ x log ⁡ p t ( x ) \nabla_x\log p_t(x) xlogpt(x)

因此,直接优化 L S M \mathcal{L}_{\mathrm{SM}} LSM 在实践中不可行。

这时可以采用与 Flow Matching 完全类似的思路:既然边缘 score 不可计算,那么我们改为回归条件 score。

于是定义 Denoising Score Matching Loss

L D S M ( θ ) = E t ∼ U n i f , z ∼ p d a t a , x ∼ p t ( ⋅ ∣ z ) [ ∥ s t θ ( x ) − ∇ x log ⁡ p t ( x ∣ z ) ∥ 2 ] \mathcal{L}_{\mathrm{DSM}}(\theta) = \mathbb{E}_{t\sim\mathrm{Unif},z\sim p_{\mathrm{data}},x\sim p_t(\cdot\mid z)} \left[ \left \| s_t^\theta(x) - \nabla_x\log p_t(x\mid z) \right \|^2 \right] LDSM(θ)=EtUnif,zpdata,xpt(z)[ stθ(x)xlogpt(xz) 2]

这里的目标变成了条件 score:

∇ x log ⁡ p t ( x ∣ z ) \nabla_x\log p_t(x\mid z) xlogpt(xz)

而对于高斯概率路径,条件 score 是有解析公式的,因此它是可计算的。

这与上一节课中的关系完全平行:

  • Flow Matching 想回归边缘向量场,但不可计算;
  • Conditional Flow Matching 改为回归条件向量场,并且可计算;
  • Score Matching 想回归边缘 score,但不可计算;
  • Denoising Score Matching 改为回归条件 score,并且可计算。

更重要的是,和 Flow Matching 一样,下面的结论同样成立:

L S M ( θ ) = L D S M ( θ ) + C \mathcal{L}_{\mathrm{SM}}(\theta) = \mathcal{L}_{\mathrm{DSM}}(\theta) + C LSM(θ)=LDSM(θ)+C

其中 C C C 是一个与参数 θ \theta θ 无关的常数。

因此,Score Matching Loss 和 Denoising Score Matching Loss 具有相同的梯度和相同的最小化器。换句话说,虽然我们无法直接优化边缘 score 的回归损失,但可以通过优化去噪 score matching 损失来间接学习边缘 score。

为什么这个结论成立?

原因与上一节课中 Flow Matching 的证明完全类似。

上一节课我们证明过,边缘向量场可以写成条件向量场关于后验分布的加权平均:

u t t a r g e t ( x ) = ∫ u t t a r g e t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z u_t^{\mathrm{target}}(x) = \int u_t^{\mathrm{target}}(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz uttarget(x)=uttarget(xz)pt(x)pt(xz)pdata(z)dz

而在上一部分中,我们也证明了边缘 score 可以写成条件 score 关于同一个后验分布的加权平均:

∇ x log ⁡ p t ( x ) = ∫ ∇ x log ⁡ p t ( x ∣ z ) p t ( x ∣ z ) p d a t a ( z ) p t ( x ) d z \nabla_x\log p_t(x) = \int \nabla_x\log p_t(x\mid z) \frac{ p_t(x\mid z)p_{\mathrm{data}}(z) }{ p_t(x) } dz xlogpt(x)=xlogpt(xz)pt(x)pt(xz)pdata(z)dz

因此,只要把 Flow Matching 证明中的 “向量场” 换成 “score 函数”,就可以得到同样的结论。

也就是说,去噪 score matching 并不是一个全新的技巧,而是与 conditional flow matching 完全平行的一种训练方式。


Q & A

Q:边缘 score 是否可处理?它和边缘向量场一样不可计算吗?

A:是的,边缘 score 通常也是不可直接计算的。

我们真正想学习的是:

∇ x log ⁡ p t ( x ) \nabla_x\log p_t(x) xlogpt(x)

但由于 p t ( x ) p_t(x) pt(x) 是对所有数据点边缘化得到的分布,而我们不知道 p d a t a p_{\mathrm{data}} pdata 的显式密度,因此无法直接计算边缘 score。

去噪 score matching 的关键就在于:虽然边缘 score 不可计算,但条件 score:

∇ x log ⁡ p t ( x ∣ z ) \nabla_x\log p_t(x\mid z) xlogpt(xz)

是可计算的。只要我们能够从数据集中采样 z z z ,并从条件路径中采样 x x x ,就可以用条件 score 作为监督信号训练网络。

这正是早期扩散模型训练中的关键突破之一。Denoising Score Matching 的思想本身早于现代扩散模型,在去噪自编码器等研究中就已经出现,后来被应用到扩散模型中,成为训练扩散生成模型的重要基础。


接下来,我们把上面的目标写成具体算法。

一般形式的 score matching 训练过程如下:

  1. 从数据集中采样一个数据点:

z ∼ p d a t a z\sim p_{\mathrm{data}} zpdata

  1. 随机采样时间:

t ∼ U n i f ( [ 0 , 1 ] ) t\sim \mathrm{Unif}([0,1]) tUnif([0,1])

  1. 从条件概率路径中采样带噪样本:

x ∼ p t ( ⋅ ∣ z ) x\sim p_t(\cdot\mid z) xpt(z)

  1. 计算条件 score:

∇ x log ⁡ p t ( x ∣ z ) \nabla_x\log p_t(x\mid z) xlogpt(xz)

  1. 训练 score 网络回归这个条件 score:

L ( θ ) = ∥ s t θ ( x ) − ∇ x log ⁡ p t ( x ∣ z ) ∥ 2 \mathcal{L}(\theta) = \left \| s_t^\theta(x) - \nabla_x\log p_t(x\mid z) \right \|^2 L(θ)= stθ(x)xlogpt(xz) 2

这就是去噪 score matching 的基本过程。

对于高斯概率路径,我们可以进一步具体化。

回顾高斯条件概率路径:

p t ( x ∣ z ) = N ( α t z , β t 2 I d ) p_t(x\mid z)=\mathcal{N}(\alpha_t z,\beta_t^2I_d) pt(xz)=N(αtz,βt2Id)

从这个分布采样可以写成:

ϵ ∼ N ( 0 , I d ) , x = α t z + β t ϵ \epsilon\sim\mathcal{N}(0,I_d), \qquad x=\alpha_t z+\beta_t\epsilon ϵN(0,Id),x=αtz+βtϵ

而条件 score 为:

∇ x log ⁡ p t ( x ∣ z ) = − x − α t z β t 2 \nabla_x\log p_t(x\mid z) = -\frac{x-\alpha_t z}{\beta_t^2} xlogpt(xz)=βt2xαtz

由于:

x − α t z = β t ϵ x-\alpha_t z=\beta_t\epsilon xαtz=βtϵ

所以:

∇ x log ⁡ p t ( x ∣ z ) = − ϵ β t \nabla_x\log p_t(x\mid z) = -\frac{\epsilon}{\beta_t} xlogpt(xz)=βtϵ

因此,Denoising Score Matching Loss 可以写成:

L D S M ( θ ) = E t ∼ U n i f , z ∼ p d a t a , ϵ ∼ N ( 0 , I d ) [ ∥ s t θ ( α t z + β t ϵ ) + ϵ β t ∥ 2 ] \mathcal{L}_{\mathrm{DSM}}(\theta) = \mathbb{E}_{t\sim\mathrm{Unif},z\sim p_{\mathrm{data}},\epsilon\sim\mathcal{N}(0,I_d)} \left[ \left \| s_t^\theta(\alpha_t z+\beta_t\epsilon) + \frac{\epsilon}{\beta_t} \right \|^2 \right] LDSM(θ)=EtUnif,zpdata,ϵN(0,Id)[ stθ(αtz+βtϵ)+βtϵ 2]

这说明,在高斯路径下,训练 score 网络等价于让网络预测加入数据中的噪声方向。

更具体地说,网络输入的是带噪样本:

x t = α t z + β t ϵ x_t=\alpha_t z+\beta_t\epsilon xt=αtz+βtϵ

而监督目标是:

− ϵ β t -\frac{\epsilon}{\beta_t} βtϵ

如果我们稍微改变参数化,也可以让模型直接预测噪声 ϵ \epsilon ϵ 。这就是许多扩散模型中常见的 噪声预测(noise prediction) 形式。

高斯概率路径下的训练算法可以概括为:

  1. 从数据集中采样数据样本:

z ∼ p d a t a z\sim p_{\mathrm{data}} zpdata

  1. 采样时间:

t ∼ U n i f ( [ 0 , 1 ] ) t\sim \mathrm{Unif}([0,1]) tUnif([0,1])

  1. 采样高斯噪声:

ϵ ∼ N ( 0 , I d ) \epsilon\sim \mathcal{N}(0,I_d) ϵN(0,Id)

  1. 构造带噪样本:

x t = α t z + β t ϵ x_t=\alpha_t z+\beta_t\epsilon xt=αtz+βtϵ

  1. 计算损失:

L ( θ ) = ∥ s t θ ( x t ) + ϵ β t ∥ 2 \mathcal{L}(\theta) = \left \| s_t^\theta(x_t) + \frac{\epsilon}{\beta_t} \right \|^2 L(θ)= stθ(xt)+βtϵ 2

  1. 通过梯度下降更新参数。

这里需要注意一个实际问题:当 β t \beta_t βt 很小时, ϵ β t \frac{\epsilon}{\beta_t} βtϵ 项会变得很大,从而可能带来数值不稳定。因此,在真实扩散模型训练中,常常会使用不同的参数化方式、权重设计或噪声调度来缓解这个问题。

4. Sampling with SDEs

到目前为止,我们把 Flow Matching 重新表述成了 Score Matching 的视角。接下来要利用 score 函数来做一件新的事情:把 ODE 采样扩展为 SDE 采样。

回顾上一节课中的结论,如果我们从初始分布出发,并沿着边缘目标向量场演化 ODE:

X 0 ∼ p i n i t , d X t = u t t a r g e t ( X t ) d t X_0\sim p_{\mathrm{init}}, \qquad dX_t=u_t^{\mathrm{target}}(X_t)dt X0pinit,dXt=uttarget(Xt)dt

那么样本分布会沿着边缘概率路径演化:

X t ∼ p t X_t\sim p_t Xtpt

这说明 ODE 采样可以生成数据分布。

现在我们想问:能否在这个确定性 ODE 中加入随机噪声,使其变成 SDE,同时仍然让样本的边缘分布保持为同一条概率路径 p t p_t pt

也就是说,我们希望构造一个 SDE,使得:

X t ∼ p t X_t\sim p_t Xtpt

仍然成立。

为此,我们先任意选择一个扩散系数:

σ t ≥ 0 \sigma_t\geq 0 σt0

它控制我们在时间 t t t 注入多少随机噪声。然后构造下面这个 SDE:

X 0 ∼ p i n i t , d X t = [ u t t a r g e t ( X t ) + σ t 2 2 ∇ x log ⁡ p t ( X t ) ] d t + σ t d W t X_0\sim p_{\mathrm{init}}, \qquad dX_t = \left[ u_t^{\mathrm{target}}(X_t) + \frac{\sigma_t^2}{2} \nabla_x\log p_t(X_t) \right]dt + \sigma_t dW_t X0pinit,dXt=[uttarget(Xt)+2σt2xlogpt(Xt)]dt+σtdWt

核心结论是:

X t ∼ p t X_t\sim p_t Xtpt

也就是说,即使我们把原来的 ODE 改成带随机噪声的 SDE,只要在漂移项中额外加入:

σ t 2 2 ∇ x log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla_x\log p_t(X_t) 2σt2xlogpt(Xt)

就仍然可以让样本分布沿着同一条概率路径 p t p_t pt 演化。

这就是 score 函数真正发挥作用的地方。

为什么需要额外加上这个 score 项?

如果我们只是简单地把 ODE:

d X t = u t t a r g e t ( X t ) d t dX_t=u_t^{\mathrm{target}}(X_t)dt dXt=uttarget(Xt)dt

改成:

d X t = u t t a r g e t ( X t ) d t + σ t d W t dX_t=u_t^{\mathrm{target}}(X_t)dt+\sigma_t dW_t dXt=uttarget(Xt)dt+σtdWt

那么注入的随机噪声会让样本向各个方向扩散。这样一来,分布就不再严格沿着原来的概率路径 p t p_t pt 演化。

因此,我们需要一个额外的校正项来抵消噪声带来的扩散效应。

这个校正项就是:

σ t 2 2 ∇ x log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla_x\log p_t(X_t) 2σt2xlogpt(Xt)

它会把样本推向当前概率分布的高密度区域。直观上,噪声项 σ t d W t \sigma_t dW_t σtdWt 会把样本扩散开,而 score 项则会根据当前分布的密度梯度,把样本拉回到合理区域。二者配合后,整体分布仍然保持在目标概率路径 p t p_t pt 上。

因此,我们可以把这个 SDE 理解为:

原来的确定性流动

加上随机噪声

再加上一个由 score 函数提供的校正漂移项。

这个结论非常重要,因为它说明:同一条概率路径不仅可以由一个 ODE 实现,也可以由一整族 SDE 实现。

不同的 σ t \sigma_t σt 会得到不同的随机动力学:

d X t = [ u t t a r g e t ( X t ) + σ t 2 2 ∇ x log ⁡ p t ( X t ) ] d t + σ t d W t dX_t = \left[ u_t^{\mathrm{target}}(X_t) + \frac{\sigma_t^2}{2} \nabla_x\log p_t(X_t) \right]dt + \sigma_t dW_t dXt=[uttarget(Xt)+2σt2xlogpt(Xt)]dt+σtdWt

但在理论上,只要 score 和向量场是准确的,它们都会产生相同的边缘分布路径 p t p_t pt ,最终也都会到达数据分布。

这意味着训练完成后,我们不仅有一个确定性的 ODE 采样器,还可以构造出一整族带随机性的 SDE 采样器。理论上它们都能生成同一个目标分布,但在实际中,它们可能在采样质量、多样性、稳定性和计算成本上表现不同。


Q & A

Q:除了 σ t ≥ 0 \sigma_t\geq 0 σt0 之外,对扩散系数还有其他限制吗?

A:理论上, σ t \sigma_t σt 主要需要满足一些基本的正则性条件,例如足够连续或足够平滑,不要过于不规则。除此之外,它可以有很大自由度。

也就是说, σ t \sigma_t σt 是一个可以人为选择的函数。不同选择会对应不同的 SDE 采样动态。理论上它们都可以保持相同的边缘概率路径,但实际效果可能不同。

Q:为什么要引入噪声?既然理论上 ODE 已经可以生成数据,为什么还需要 SDE?

A:引入噪声给了我们更多采样自由度。

ODE 采样是确定性的:给定初始噪声后,轨迹就是固定的。而 SDE 采样在演化过程中不断注入随机性,因此即使从相同或相近的初始状态出发,也可能探索到不同轨迹。

理论上,不同的 σ t \sigma_t σt 对应的一整族模型都能沿着相同的概率路径生成数据。但在实践中,由于神经网络近似误差、数值离散误差和采样步数有限,不同 SDE 采样器可能表现不同。这样我们就可以通过实验选择效果更好的采样方式。

Q σ t d W t \sigma_t dW_t σtdWt σ t 2 2 ∇ x log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla_x\log p_t(X_t) 2σt2xlogpt(Xt) 这两项分别有什么直觉?

A σ t d W t \sigma_t dW_t σtdWt 是随机噪声项,它会让样本在各个方向上扩散。

而:

σ t 2 2 ∇ x log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla_x\log p_t(X_t) 2σt2xlogpt(Xt)

是 score 校正项。它利用当前分布 p t p_t pt 的 score,把样本推向更高概率密度的区域。

直观地说,噪声项会把概率质量扩散开,而 score 项会把它拉回到合理的分布形状中。二者共同作用,使得虽然单条轨迹变得随机,但整体样本分布仍然保持在同一条概率路径上。


前面我们已经引入了一个重要结论:如果我们已经有一个能够遵循概率路径 p t p_t pt 的边缘向量场:

u t t a r g e t ( x ) u_t^{\mathrm{target}}(x) uttarget(x)

那么可以在 ODE 的基础上加入随机噪声,并通过 score 函数进行修正,从而得到一个 SDE。这个 SDE 虽然单条轨迹是随机的,但整体分布仍然会沿着同一条概率路径 p t p_t pt 演化。

现在我们来证明这个结论。

证明中需要用到 SDE 对应的概率密度演化方程,也就是 Fokker-Planck 方程

假设随机过程满足 SDE:

X 0 ∼ p i n i t , d X t = u t ( X t ) d t + σ t d W t X_0\sim p_{\mathrm{init}}, \qquad dX_t=u_t(X_t)dt+\sigma_t dW_t X0pinit,dXt=ut(Xt)dt+σtdWt

其中, u t u_t ut 是漂移向量场, σ t \sigma_t σt 是扩散系数。若 X t X_t Xt 的边缘分布为:

X t ∼ p t X_t\sim p_t Xtpt

那么 p t ( x ) p_t(x) pt(x) 需要满足 Fokker-Planck 方程:

∂ t p t ( x ) = − div ⁡ ( p t u t ) ( x ) + σ t 2 2 Δ p t ( x ) , x ∈ R d ,   0 ≤ t ≤ 1 \partial_t p_t(x) = -\operatorname{div}(p_tu_t)(x) + \frac{\sigma_t^2}{2}\Delta p_t(x), \qquad x\in\mathbb{R}^d,\ 0\leq t\leq 1 tpt(x)=div(ptut)(x)+2σt2Δpt(x),xRd, 0t1

这个方程可以看作连续性方程在 SDE 情况下的推广。

回顾 ODE 情况下的连续性方程:

∂ t p t ( x ) = − div ⁡ ( p t u t ) ( x ) \partial_t p_t(x)=-\operatorname{div}(p_tu_t)(x) tpt(x)=div(ptut)(x)

它只描述概率质量如何被向量场搬运。而在 SDE 中,由于额外加入了布朗噪声,概率质量不仅会被向量场推动,还会发生随机扩散。因此,Fokker-Planck 方程中多了一项:

σ t 2 2 Δ p t ( x ) \frac{\sigma_t^2}{2}\Delta p_t(x) 2σt2Δpt(x)

这一项可以理解为类似热方程中的扩散项,用来描述噪声导致的概率质量扩散。

这里的 Δ \Delta ΔLaplacian 算子,定义为:

Δ w t ( x ) = ∑ i = 1 d ∂ 2 ∂ x i 2 w t ( x ) = div ⁡ ( ∇ w t ) ( x ) \Delta w_t(x) = \sum_{i=1}^{d} \frac{\partial^2}{\partial x_i^2}w_t(x) = \operatorname{div}(\nabla w_t)(x) Δwt(x)=i=1dxi22wt(x)=div(wt)(x)

也就是说,Laplacian 可以看作梯度的散度,或者 Hessian 矩阵的迹。

下面从直观角度理解 Fokker-Planck 方程:

在某个位置 x x x 附近,概率质量的变化来自两部分。

第一部分是由向量场 u t ( x ) u_t(x) ut(x) 引起的定向流动。它类似于 ODE 中的情况:概率质量会沿着向量场方向移动,因此某个区域的概率密度会因为流入或流出而变化。

第二部分是由扩散系数 σ t \sigma_t σt 引起的随机扩散。布朗运动会让样本向各个方向随机扩散,因此概率质量会像热量一样向周围散开。

因此,Fokker-Planck 方程表达的是:

概率密度的变化 = 向量场搬运造成的变化 * 随机扩散造成的变化。

也就是说:

∂ t p t ( x ) = − div ⁡ ( p t u t ) ( x ) + σ t 2 2 Δ p t ( x ) \partial_t p_t(x) = -\operatorname{div}(p_tu_t)(x) + \frac{\sigma_t^2}{2}\Delta p_t(x) tpt(x)=div(ptut)(x)+2σt2Δpt(x)

其中,第一项是连续性方程中的流动项,第二项是热扩散项。

现在我们使用 Fokker-Planck 方程证明前面的 SDE extension trick。

已知边缘向量场 u t t a r g e t u_t^{\mathrm{target}} uttarget 会让 ODE 遵循概率路径 p t p_t pt

d X t = u t t a r g e t ( X t ) d t ⇒ X t ∼ p t dX_t=u_t^{\mathrm{target}}(X_t)dt \quad\Rightarrow\quad X_t\sim p_t dXt=uttarget(Xt)dtXtpt

因此, p t p_t pt u t t a r g e t u_t^{\mathrm{target}} uttarget 满足连续性方程:

∂ t p t ( x ) = − div ⁡ ( p t u t t a r g e t ) ( x ) \partial_t p_t(x) = -\operatorname{div}(p_tu_t^{\mathrm{target}})(x) tpt(x)=div(ptuttarget)(x)

现在我们希望构造一个 SDE:

d X t = [ u t t a r g e t ( X t ) + σ t 2 2 ∇ log ⁡ p t ( X t ) ] d t + σ t d W t dX_t = \left[ u_t^{\mathrm{target}}(X_t) + \frac{\sigma_t^2}{2}\nabla\log p_t(X_t) \right]dt + \sigma_t dW_t dXt=[uttarget(Xt)+2σt2logpt(Xt)]dt+σtdWt

并证明它仍然满足:

X t ∼ p t X_t\sim p_t Xtpt

也就是说,我们要证明 p t p_t pt 满足该 SDE 对应的 Fokker-Planck 方程。

从连续性方程开始:

∂ t p t ( x ) = − div ⁡ ( p t u t t a r g e t ) ( x ) = − div ⁡ ( p t u t t a r g e t ) ( x ) − σ t 2 2 Δ p t ( x ) + σ t 2 2 Δ p t ( x ) \begin{aligned} \partial_t p_t(x) &= -\operatorname{div} \left( p_tu_t^{\mathrm{target}} \right)(x) \\[8pt] &= -\operatorname{div} \left( p_tu_t^{\mathrm{target}} \right)(x) - \frac{\sigma_t^2}{2}\Delta p_t(x) + \frac{\sigma_t^2}{2}\Delta p_t(x) \end{aligned} tpt(x)=div(ptuttarget)(x)=div(ptuttarget)(x)2σt2Δpt(x)+2σt2Δpt(x)

第二步只是加减了同一个项:

σ t 2 2 Δ p t ( x ) \frac{\sigma_t^2}{2}\Delta p_t(x) 2σt2Δpt(x)

因此等式没有改变。

接下来,利用 Laplacian 的定义:

Δ p t ( x ) = div ⁡ ( ∇ p t ) ( x ) \Delta p_t(x) = \operatorname{div}(\nabla p_t)(x) Δpt(x)=div(pt)(x)

所以:

− σ t 2 2 Δ p t ( x ) = − div ⁡ ( σ t 2 2 ∇ p t ) ( x ) -\frac{\sigma_t^2}{2}\Delta p_t(x) = -\operatorname{div} \left( \frac{\sigma_t^2}{2}\nabla p_t \right)(x) 2σt2Δpt(x)=div(2σt2pt)(x)

于是:

∂ t p t ( x ) = − div ⁡ ( p t u t t a r g e t ) ( x ) − div ⁡ ( σ t 2 2 ∇ p t ) ( x ) + σ t 2 2 Δ p t ( x ) \begin{aligned} \partial_t p_t(x) &= -\operatorname{div} \left( p_tu_t^{\mathrm{target}} \right)(x) - \operatorname{div} \left( \frac{\sigma_t^2}{2}\nabla p_t \right)(x) + \frac{\sigma_t^2}{2}\Delta p_t(x) \end{aligned} tpt(x)=div(ptuttarget)(x)div(2σt2pt)(x)+2σt2Δpt(x)

接下来使用 score 函数的恒等式:

∇ log ⁡ p t ( x ) = ∇ p t ( x ) p t ( x ) \nabla\log p_t(x)=\frac{\nabla p_t(x)}{p_t(x)} logpt(x)=pt(x)pt(x)

因此:

∇ p t ( x ) = p t ( x ) ∇ log ⁡ p t ( x ) \nabla p_t(x)=p_t(x)\nabla\log p_t(x) pt(x)=pt(x)logpt(x)

代入可得:

∂ t p t ( x ) = − div ⁡ ( p t u t t a r g e t ) ( x ) − div ⁡ ( p t σ t 2 2 ∇ log ⁡ p t ) ( x ) + σ t 2 2 Δ p t ( x ) \begin{aligned} \partial_t p_t(x) &= -\operatorname{div} \left( p_tu_t^{\mathrm{target}} \right)(x) - \operatorname{div} \left( p_t \frac{\sigma_t^2}{2} \nabla\log p_t \right)(x) + \frac{\sigma_t^2}{2}\Delta p_t(x) \end{aligned} tpt(x)=div(ptuttarget)(x)div(pt2σt2logpt)(x)+2σt2Δpt(x)

再利用散度算子的线性性,将前两项合并:

∂ t p t ( x ) = − div ⁡ ( p t [ u t t a r g e t + σ t 2 2 ∇ log ⁡ p t ] ) ( x ) + σ t 2 2 Δ p t ( x ) \begin{aligned} \partial_t p_t(x) &= -\operatorname{div} \left( p_t \left[ u_t^{\mathrm{target}} + \frac{\sigma_t^2}{2}\nabla\log p_t \right] \right)(x) + \frac{\sigma_t^2}{2}\Delta p_t(x) \end{aligned} tpt(x)=div(pt[uttarget+2σt2logpt])(x)+2σt2Δpt(x)

这正是下面这个 SDE 对应的 Fokker-Planck 方程:

d X t = [ u t t a r g e t ( X t ) + σ t 2 2 ∇ log ⁡ p t ( X t ) ] d t + σ t d W t dX_t = \left[ u_t^{\mathrm{target}}(X_t) + \frac{\sigma_t^2}{2}\nabla\log p_t(X_t) \right]dt + \sigma_t dW_t dXt=[uttarget(Xt)+2σt2logpt(Xt)]dt+σtdWt

因此,该 SDE 的边缘分布确实满足:

X t ∼ p t X_t\sim p_t Xtpt

这就证明了:只要我们在加入随机噪声 σ t d W t \sigma_t dW_t σtdWt 的同时,额外加入 score 校正项:

σ t 2 2 ∇ log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla\log p_t(X_t) 2σt2logpt(Xt)

就可以保持原来的概率路径不变。


Q & A

Q:边缘向量场和 score 函数是否都需要分别学习?

A:在高斯概率路径下,边缘向量场和 score 函数可以相互重参数化。因此,通常不需要分别学习两个网络。

如果训练的是向量场网络,可以通过前面推导的转换公式得到 score;如果训练的是 score 网络,也可以转换得到对应的向量场。

不过,这种 “向量场与 score 可互相转换” 的结论依赖于高斯概率路径。而 SDE extension trick 本身更一般,它说明只要有边缘向量场和对应 score,就可以构造一族遵循同一概率路径的 SDE。


现在我们可以回到采样问题。

前面已经知道,如果训练得到了边缘向量场,就可以通过 ODE 采样生成数据:

d X t = u t t a r g e t ( X t ) d t dX_t=u_t^{\mathrm{target}}(X_t)dt dXt=uttarget(Xt)dt

而通过 SDE extension trick,我们还可以选择任意扩散系数 σ t \sigma_t σt ,构造 SDE 采样器:

d X t = [ u t t a r g e t ( X t ) + σ t 2 2 ∇ log ⁡ p t ( X t ) ] d t + σ t d W t dX_t = \left[ u_t^{\mathrm{target}}(X_t) + \frac{\sigma_t^2}{2}\nabla\log p_t(X_t) \right]dt + \sigma_t dW_t dXt=[uttarget(Xt)+2σt2logpt(Xt)]dt+σtdWt

对于高斯概率路径,边缘向量场可以写成 score 的形式:

u t t a r g e t ( x ) = a t ∇ log ⁡ p t ( x ) + b t x u_t^{\mathrm{target}}(x) = a_t\nabla\log p_t(x) + b_tx uttarget(x)=atlogpt(x)+btx

因此,将它代入上面的 SDE,可以得到完全用 score 表达的随机采样过程:

d X t = [ ( a t + σ t 2 2 ) ∇ log ⁡ p t ( X t ) + b t X t ] d t + σ t d W t dX_t = \left[ \left( a_t+\frac{\sigma_t^2}{2} \right) \nabla\log p_t(X_t) + b_tX_t \right]dt + \sigma_t dW_t dXt=[(at+2σt2)logpt(Xt)+btXt]dt+σtdWt

如果我们训练了一个 score 网络:

s t θ ( x ) ≈ ∇ log ⁡ p t ( x ) s_t^\theta(x)\approx \nabla\log p_t(x) stθ(x)logpt(x)

那么采样时可以直接将 score 网络代入:

d X t = [ ( a t + σ t 2 2 ) s t θ ( X t ) + b t X t ] d t + σ t d W t dX_t = \left[ \left( a_t+\frac{\sigma_t^2}{2} \right) s_t^\theta(X_t) + b_tX_t \right]dt + \sigma_t dW_t dXt=[(at+2σt2)stθ(Xt)+btXt]dt+σtdWt

这就是扩散模型中的 SDE 采样形式。

它说明:训练完成后,我们可以不仅仅使用确定性的 ODE 采样,还可以选择不同的扩散系数 σ t \sigma_t σt ,得到一整族随机采样过程。

理论上,如果 score 和向量场完全准确,并且连续时间 SDE 可以精确模拟,那么所有这些采样过程都会沿着同一条概率路径 p t p_t pt ,最终生成同样的数据分布。

下面是 ODE 采样和 SDE 采样的直观对比:

左边是确定性采样,右边是随机采样。可以看到,二者最终都在生成蛋白质结构,但轨迹形态明显不同。

ODE 采样中,给定初始噪声后,轨迹是确定的;而 SDE 采样中,由于每一步都会注入随机噪声,轨迹会更加随机和粗糙。

另一个蛋白质结构生成的例子如下:

这个例子展示了通过 SDE 采样将噪声逐步转化为蛋白质结构的过程。与 ODE 相比,SDE 采样的动态中包含随机扰动,因此生成路径更加随机。

既然 ODE 采样已经可以生成数据,为什么还要引入 SDE 采样?

从理论上看,如果模型是完美的,并且采样过程也能精确模拟,那么所有扩散系数 σ t \sigma_t σt 对应的 SDE 都应该得到相同的结果。因为它们都遵循同一条概率路径,最终都应该采样到数据分布:

X 1 ∼ p d a t a X_1\sim p_{\mathrm{data}} X1pdata

但在实践中,情况并不完全一样。主要原因有两个。

首先是 训练误差(training error)

神经网络不可能完美学习边缘向量场或 score 函数。无论使用 Flow Matching 还是 Score Matching,训练出来的网络都只是目标函数的近似。因此,不同采样动态对这种近似误差的敏感程度可能不同。

其次是 模拟误差(simulation error)

无论是 ODE 还是 SDE,实际采样时都不可能精确求解连续时间方程,只能使用有限步数的数值方法进行离散化。离散化会带来误差,而不同的 σ t \sigma_t σt 选择会改变这个误差的表现。

因此,虽然理论上不同 SDE 采样器等价,但在实践中它们可能产生不同效果。我们可以把 σ t \sigma_t σt 看作一个可调节的采样超参数,通过实验找到更好的采样动态。

不过,一个重要经验是:ODE 采样通常已经能取得很好的结果

因此,SDE 采样更应该被看作一种可选工具,而不是必须使用的步骤。对于很多任务,只使用确定性 ODE 采样就已经足够有效。

当然,在某些下游应用中,随机演化可能更有用。例如:

  • 微调或后训练;
  • 推理时优化;
  • 需要在采样过程中进行探索的任务;
  • 某些需要多样性的生成场景。

这些情况下,SDE 采样提供的随机性可能会带来额外优势。

最后补充一个有趣的联系:Langevin dynamics 与分子动力学模拟。

在化学和物理中,分子动力学模拟常常需要从某个平衡分布中采样,例如玻尔兹曼分布:

p B o l t z m a n n ( x ) = 1 Z exp ⁡ ( − U ( x ) ) p_{\mathrm{Boltzmann}}(x) = \frac{1}{Z}\exp(-U(x)) pBoltzmann(x)=Z1exp(U(x))

其中, U ( x ) U(x) U(x) 是势能函数, Z Z Z 是归一化常数。

如果我们考虑一个恒定不变的概率路径:

p t ( x ) = p B o l t z m a n n ( x ) p_t(x)=p_{\mathrm{Boltzmann}}(x) pt(x)=pBoltzmann(x)

并令边缘向量场为零:

u t t a r g e t ( x ) = 0 u_t^{\mathrm{target}}(x)=0 uttarget(x)=0

那么 SDE extension trick 给出的 SDE 就变成:

d X t = σ t 2 2 ∇ log ⁡ p t ( X t ) d t + σ t d W t dX_t = \frac{\sigma_t^2}{2} \nabla\log p_t(X_t)dt + \sigma_t dW_t dXt=2σt2logpt(Xt)dt+σtdWt

这就是 Langevin dynamics 的形式。

由于:

p t ( x ) = p B o l t z m a n n ( x ) = 1 Z exp ⁡ ( − U ( x ) ) p_t(x)=p_{\mathrm{Boltzmann}}(x) = \frac{1}{Z}\exp(-U(x)) pt(x)=pBoltzmann(x)=Z1exp(U(x))

因此:

∇ log ⁡ p t ( x ) = − ∇ U ( x ) \nabla\log p_t(x) = -\nabla U(x) logpt(x)=U(x)

也就是说,动力学会受到势能梯度的影响,同时又会受到随机噪声扰动。

这个过程有一个重要性质:

X 0 ∼ p B o l t z m a n n ⇒ X t ∼ p B o l t z m a n n X_0\sim p_{\mathrm{Boltzmann}} \quad\Rightarrow\quad X_t\sim p_{\mathrm{Boltzmann}} X0pBoltzmannXtpBoltzmann

也就是说,如果初始分布是玻尔兹曼分布,那么经过 Langevin dynamics 演化后,分布仍然保持为玻尔兹曼分布。这就是平衡分布的含义。

这个例子说明,本课程中讨论的 SDE、score 和概率路径,不仅出现在生成模型中,也和物理、化学中的分子动力学模拟有深刻联系。


Q & A

Q:这里的 U ( x ) U(x) U(x) 是什么?

A U ( x ) U(x) U(x) 可以理解为能量或势能函数。

在统计物理中,玻尔兹曼分布会给低能量状态更高概率,给高能量状态更低概率。分子动力学模拟通常希望从这样的分布中采样,从而描述系统在平衡状态下可能出现的构型。

这里不需要深入统计物理细节,只需要知道:Langevin dynamics 是一种通过 “梯度项 + 随机噪声” 来保持或采样目标分布的动力学形式。这和我们在扩散模型中看到的 SDE 结构非常相似。


5. Key takeaway

最后回顾一下这一部分的核心内容:

第一,对于高斯概率路径,学习边缘向量场和学习 score 函数是等价的

因为二者之间存在转换公式:

u t t a r g e t ( x ) = a t ∇ log ⁡ p t ( x ) + b t x u_t^{\mathrm{target}}(x) = a_t\nabla\log p_t(x)+b_tx uttarget(x)=atlogpt(x)+btx

所以我们既可以训练网络预测速度场,也可以训练网络预测 score。

第二,Denoising Score Matching 提供了一种学习边缘 score 的简单方式

边缘 score:

∇ log ⁡ p t ( x ) \nabla\log p_t(x) logpt(x)

不可直接计算,但条件 score:

∇ log ⁡ p t ( x ∣ z ) \nabla\log p_t(x\mid z) logpt(xz)

是可计算的。通过回归条件 score,我们可以间接学习边缘 score。这与 Conditional Flow Matching 学习边缘向量场的逻辑完全平行。

第三,score 函数使 SDE 采样成为可能

在 ODE 采样中加入噪声后,必须加入 score 校正项:

σ t 2 2 ∇ log ⁡ p t ( X t ) \frac{\sigma_t^2}{2}\nabla\log p_t(X_t) 2σt2logpt(Xt)

这样才能保持相同的概率路径。因此,训练好的 score 模型不仅可以用于确定性 ODE 采样,也可以用于随机 SDE 采样。

总体来说,这一部分告诉我们:score 函数不仅是 Flow Matching 的另一种参数化方式,也是连接扩散模型、随机采样和物理中 Langevin dynamics 的关键工具。

下一部分将进入本节课的另一个核心主题:Guidance,尤其是无分类器引导。Score 函数会在理解条件生成和引导机制时继续发挥关键作用。

结语

本节课围绕 Score Function 与 Score Matching 展开,从一个全新的角度重新审视了 Flow Matching 与扩散模型之间的内在联系 。

首先,我们从 score 的定义出发,将其理解为对数密度的梯度: ∇ x log ⁡ p t ( x ) \nabla_x \log p_t(x) xlogpt(x) 。它刻画了分布在局部空间中 “上升最快的方向”。在此基础上,我们分别定义了条件 score 与边缘 score,并证明了一个关键结论:边缘 score 可以表示为条件 score 在后验分布下的加权平均。这一结构与 Flow Matching 中 “边缘向量场 = 条件向量场的后验平均” 完全一致。

随后,我们进一步揭示了一个重要事实:在高斯概率路径下,Score Matching 与 Flow Matching 本质上是同一类问题的不同参数化形式。具体而言,向量场与 score 之间可以通过时间相关系数相互转换。这意味着生成模型的学习目标,既可以是 “预测速度场”,也可以是 “预测 score 函数”,二者在理论上等价。

在训练方法上,我们推导了 Denoising Score Matching,说明虽然边缘 score 不可直接计算,但可以通过条件 score 构造可训练的监督信号。这一思想与 Flow Matching 中的 conditional objective 完全平行,共同构成了现代扩散模型训练的核心基础。

在采样部分,我们将 score 引入 SDE 框架,得到了关键的 SDE extension 结果:通过在 ODE 采样中加入噪声项,并引入 score 校正漂移项,可以构造一族保持同一概率路径的随机动力系统。这一结果揭示了一个统一视角:ODE 采样与 SDE 采样本质上是同一生成过程的两种实现方式。

此外,我们还将 score matching 与 Langevin dynamics 联系起来,说明其不仅是生成模型中的工具,也深刻连接了统计物理中的平衡分布采样问题。

下一讲我们将进入另一个核心主题 Guidance,敬请期待🤗

参考

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐