得分匹配的朗之万动力学——Score-Matching Langevin Dynamics (SMLD)

摘要

本次周报主要学习了得分匹配的朗之万动力学（Score-Matching Langevin Dynamics, SMLD）的基本思想、理论基础及其相关技术。SMLD通过学习数据分布的梯度（分数函数），结合朗之万等式生成样本。在学习过程中，首先介绍了从分布中采样的理论基础，包括朗之万等式、福克-普朗克方程以及随机梯度朗之万动力学（SGLD）的概念；接着详细分析了斯坦因分数函数的定义及其几何意义，并通过显式得分匹配、隐式得分匹配和去噪得分匹配三种技术讨论了如何估计分数函数。最后，结合噪声条件分数网络（NCSN）的工作，展示了得分匹配的实际应用方法及生成图像的过程。本次学习为理解扩散模型及其变体提供了理论支持。

Abstract

This report focuses on the study of Score-Matching Langevin Dynamics (SMLD), its fundamental principles, theoretical foundations, and related techniques. SMLD generates samples by learning the gradient of the data distribution (score function) and employing the Langevin equation. The report begins with an introduction to the theoretical foundation of sampling from distributions, including the Langevin equation, Fokker-Planck equation, and Stochastic Gradient Langevin Dynamics (SGLD). Subsequently, the Stein score function is defined and its geometric interpretation is analyzed. Three techniques for estimating the score function—explicit score matching, implicit score matching, and denoising score matching—are discussed in detail. Finally, real-world applications of score matching are demonstrated, particularly through the Noise Conditional Score Network (NCSN), which explains its practical use in generating images. This report provides essential theoretical support for understanding diffusion models and their variants.

周报内容

0. 上期补充

这一张主要是介绍${\alpha }_t$不作为提前设定好的超参数时，应当如何神经网络进行学习。主要提出利用单调的神经网络来计算信噪比(signal-to-noise ratio, SNR)，进而参与优化原始的生成目标。

这一张主要是介绍从3种不同的视角，提出了3种等价的优化目标。第1种优化目标，是预测原图$x_0$，这是上一次周报扩散模型复习——Diffusion Models Review(Understanding Diffusion Models: A Unified Perspective论文公式推导)中已经介绍过的。接着，利用重参数化的技巧，写出原图$x_0$的式子，带入第1种优化目标的式子中，重新化简，从而得到新的优化目标：预测原始噪声${ϵ}_0$。

这一张，是从样本梯度和分数函数的角度来解释扩散模型的优化目标。通过引入Tweedie公式，利用后验均值的等价关系，将原图$x_0$的式子写为带有分数函数的式子。

最后一张详细地解释了扩散模型从分数函数的角度出发推导得到的优化函数，即建模一个预测分数函数的神经网络$s_{\theta }(x_t,t)$。在论述的最后，解释了预测原始噪声${ϵ}_0$与预测分数函数${\nabla }_{x_t}logp(x_t)$的等价关系，两者的等价关系仅由一个随时间变化的${\bar{\alpha }}_t$所控制。

1. 本期的基本思想

本次周报的主要目标是学习得分匹配的朗之万动力学(Score-Matching Langevin Dynamics , SMLD)。SMLD先于DDPM提出，DDPM是受其启发，在SMLD的基础上进一步改善了SMLD的生成质量和效率。SMLD的主要思想是通过学习数据分布的梯度（分数函数），然后利用朗之万动力学等式生成图像（样本）。SMLD主要包含3个部分：朗之万等式(the Langevin equation)，(斯坦因)分数函数(the (Stein) score function)和得分匹配的损失(the score-matching loss)。SMLD的主要目标就是从指定分布$p(x)$中生成数据：

1. 当$p(x)$已知时，将${\nabla }_x\log p(x_t)$代入Langevin equation，生成样本$x_t$。这一部分对应SMLD的第1部分。
2. 当$p(x)$未知时，需训练一个神经网络$s_{\theta }(x_t)$拟合${\nabla }_x\log p(x_t)$，然后再转入1的步骤，代入到Langevin equation中生成样本$x_t$。如何设定$s_{\theta }(x_t)$是SMLD第2部分要讨论的问题，而如何训练$s_{\theta }(x_t)$进行拟合是SMLD第3部分讨论的问题。

本次周报的学习主要基于论文：Tutorial on Diffusion Models for Imaging and Vision

2. 从一个分布中采样（Sampling from a Distribution）

想象给定了一个分布$p(x)$，我们希望从这个分布中采样。如果$x$是一维变量，那么我们可以计算累积分布函数（cumulative distribution function, CDF）的逆函数来实现，累积分布函数其实就是我们一般称为的分布函数。即：

1. 计算逆函数：对于目标分布的累积分布函数$F(x)$，计算其逆函数$F^{-1}(u)$。此逆函数将均匀分布变量$U\in [0,1]$映射到目标分布。
2. 应用逆函数：生成一个均匀分布的随机数$u\sim \mathrm{Uniform}[0,1]$，然后通过$x=F^{-1}(u)$得到服从目标分布的随机变量$X$。

对于高维的分布，上述的做法则不再适用。然而，“给高概率样本更高权重”的直觉想法仍然是有效的。例如，如果我们想从橙色数据集中采样出样本，几乎可以肯定的是，我们想要球形的橙子而不是立方体的橙子。

“给高概率样本更高权重”的直觉，这句话的意思是：在概率论与统计学中，样本的权重通常反映其出现的概率。高概率的样本更可能被频繁选中，这一原则是许多采样方法的基础。例如：

• 概率分布采样：若样本服从某个分布，高概率区域（如正态分布的均值附近）的样本会被更频繁地生成。
• 重要性采样：通过调整权重，使低概率但重要的样本也能被有效利用（如估计罕见事件的期望值）。

回到如何从分布中采样的问题上，如果我们的目标是生成水果的图片，那么我们应该希望与那些”奇怪的“水果图像相比，潜在分布$p(x)$中“正常的”水果图像应该具有更高的值。因此要从一个分布中采样，从较高的分布位置中选择样本是更自然的。我们可以将这个想法转化为一个优化问题：

$$x^{\ast }=\underset{x}{\mathrm{argmax}}\log p(x)$$

其目的是最大化分布$p(x)$的对数似然。但是这种方法没有告诉我们如何采样低概率的样本，这个问题将由朗之万等式(Langevin equation)解决。

这里我们先要介绍这里的最大化对数似然与极大似然估计的区别。在极大似然估计中，数据样本$x$是固定的，但是模型的参数是变化的。这里的最大化优化问题，模型的参数是固定的，但是数据样本$x$是变化的，也就是模型现在是确定的，我们的目标是从该模型中采样出最可能的样本。下表展示了采样问题和极大似然估计的区别：

现在，我们继续讨论最大化对数似然的优化问题。如果$p(x)$是一个简单的参数模型，则最大化将具有解析解。然而，在一般的高维空间中，优化问题常会受到局部极小值的影响。所以，没有一个算法能够解决所有的问题。最合理优先考虑的方法就是梯度下降算法，针对我们这里的优化目标，梯度下降法可以写为：

$$x_{t+1}=x_t+\tau {\nabla }_x\log p(x_t)$$

其中，${\nabla }_x\log p(x_t)$是$\log p(x)$是在$x_t$处取值的梯度值，$\tau$是步长。这里作者使用"+“符号，而不是”-"符号，是因为这里在解最大化的问题。

2.1 朗之万等式（Langevin equation）

对于离散时间$t=1,\dots ,T$，朗之万等式(Langevin equation)用于从一个已知的分布$p(x)$迭代地进行采样。公式写为：

$$x_{t+1}=x_t+\tau {\nabla }_x\log p(x_t)+\sqrt{2\tau }z,z\sim \mathcal{N}(0,I)$$

其中$\tau$是步长，是用户控制的；$x_0$是白噪声，即$x_0$服从标准正态分布。

下面展示一个例子，$p(x)$是一个参数为$\mu$和$\sigma$的高斯分布，利用Langevin equation从这个高斯分布采样的过程为：

如果我们忽视噪声项$\sqrt{2\tau }z$，那么Langevin equation就是纯粹的梯度下降法。但在生成模型的场景下：

$$Langevinequation=gradientdescent+noise$$

为什么我们需要“梯度下降+噪声项”的形式，而不是纯粹的“梯度下降”？

一种解释是，我们这里关注的重点不是解一个优化问题，而是如何从一个分布中采样的问题。通过将随机噪声引入梯度下降的步骤中，可以随机选取一个样本，沿着目标函数的轨迹移动，而不是停留在原地。如果接近峰值，就会左右轻微移动；如果远离峰值，梯度方向会将其拉向峰值。如果峰值周围的曲率较陡，采样点就会接近稳态点$x_T$。如果曲率较平缓，采样点就会较分散。这样采样的效果会比较好，能够很好地描绘出我们指定的分布$p(x)$​。

举个例子来说： 纯粹的梯度下降像一颗石子从山顶滚到最近的山谷（找到局部极小值）；而梯度下降+噪声像一群蜜蜂在山间飞行，蜜蜂受重力（梯度）影响飞向山谷，但翅膀振动（噪声）让它们偶尔飞离谷底，最终蜜蜂的分布密度反映山谷的深浅（概率密度高低）。

总结来说，有以下几点理由：

• 逃离局部极小：噪声提供随机扰动，避免陷入局部极小值（优化中需要避免的“陷阱”，但采样中需要覆盖整个分布）。
• 动态平衡：
  • 靠近峰值（概率密度高）：梯度方向指向峰值，但噪声使样本在附近轻微震荡（类似“围绕峰值徘徊”）。
  • 远离峰值（概率密度低）：梯度方向将样本拉向高概率区域，噪声影响较小。
  • 曲率影响：
    • 若峰值处曲率大（分布尖锐），样本集中在狭窄区域。
    • 若曲率小（分布平坦），样本分散在更广区域。
• 通过反复应用“梯度方向+噪声”的更新规则，最终样本会收敛到目标分布$p(x)$。

列出第2个例子：考虑一个高斯混合分布$p(x)={\pi }_1\mathcal{N}(x\mid {\mu }_1,{\sigma }_1^2)+{\pi }_2\mathcal{N}(x\mid {\mu }_2,{\sigma }_2^2)$。通过数学计算，可以得到${\nabla }_x\log p(x)$。为了说明，我们选择${\pi }_1=0.6,{\mu }_1=2,{\sigma }_1=0.5,{\pi }_2=0.4$, ${\mu }_2=-2,{\sigma }_2=0.2$，初始化$x_0=0$，并令$\tau =0.05$。我们将上述的梯度下降迭代运行$T=500$次，并画出相应的$p(x_t)$的轨迹。如上图可见，序列 { x 1 , x 2 , … , x T } \{x_1,x_2,\ldots,x_T\} {x1​,x2​,…,xT​}随着高斯分布的形状，爬到了一个峰值，这是纯**“梯度下降”**的采样方式。
加上噪声项后的“**梯度下降+噪声”**模式就更有意思了，序列$x_t$没有停留在极值，而是绕过了极值并停留在极值的旁边一个地方。越接近峰值，停在那里的概率就越高。

上图展示了样本轨迹的一个有趣的现象：从任意一个位置开始，$x_t$会按朗之万动力学运动。这个运动方向并不完全是任意的，而是存在一定量的预定义漂移，同时在每一步都存在一定程度的随机性。漂移由${\nabla }_x\log p(x)$决定，而随机性则来自$z$。作者将同个高斯混合分布的两种模式（纯“梯度下降”和“梯度下降+噪声”）用不同的颜色标出, 以便更好地可视化。红色的为纯“梯度下降”，蓝色的为“梯度下降+噪声”，我们发现蓝色比红色在收敛时更为集中。

给出第3个例子：按照第2个例子，继续考虑一个高斯混合分布$p(x)={\pi }_1\mathcal{N}(x\mid {\mu }_1,{\sigma }_1^2)+{\pi }_2\mathcal{N}(x\mid {\mu }_2,{\sigma }_2^2)$。为了说明，我们选择${\pi }_1=0.6,{\mu }_1=2,{\sigma }_1=0.5,{\pi }_2=0.4$, ${\mu }_2=-2,{\sigma }_2=0.2$，假设从$x_0\sim Uniform[-3,3]$中采集$M=10000$个均匀分布的样本。使用朗之万采样法更新$t=100$次，生成的样本的直方图如上图所示，最后很好地描绘了真实的分布。

2.2 福克-普朗克方程（Fokker-Planck equation）

Fokker-Planck equation是一种稍微正式的方式来证明Langevin equation的合理性。Fokker-Planck equation是随机过程的基本结果。对于任何的马尔可夫过程（例如，维纳过程和布朗运动），解$x_t$的动态由随机微分方程（即Langevin equation）描述。但是，由于$x_t$无论在任何时间$t$都是随机变量，因此存在与每个$x_t$相关的潜在概率分布$p(x,t)$。Fokker-Planck equation提供了关于这个潜在分布的数学描述，即该分布必须满足一个偏微分方程。

在我们的问题背景下，Langevin equation的解$x_t$在满足Fokker-Planck equation的时间$t$时，将具有概率分布$p(x,t)$（对于Langevin equation，其对应的Fokker-Planck equation为）：

∂ t p ( x , t ) = − ∂ x { [ ∂ x ( log ⁡ p ( x ) ) ] p ( x , t ) } + ∂ x 2 p ( x , t ) {\partial _t}p(x,t) = - {\partial _x}\left\{ {[{\partial _x}(\log p(x))]p(x,t)} \right\} + \partial _x^2p(x,t) ∂t​p(x,t)=−∂x​{[∂x​(logp(x))]p(x,t)}+∂x2​p(x,t)

其中$\log p(x)$是真实分布的对数似然。

推导出Fokker-Planck equation比较复杂，但验证是否满足Fokker-Planck equation并不难：

2.3 随机梯度的朗之万动力学（Stochastic Gradient Langevin Dynamics）

由Langevin equation控制$x_t$的动力学行为通常称为朗之万动力学(Langevin Dynamics)。Langevin动力学通过在梯度下降的过程中添加噪声，使得最终能够从目标分布中采样，而不是仅仅寻找一个最优解。而SGD是优化算法，通过使用小批量数据（minibatch）来估计梯度，从而加速训练过程，其随机性主要来自于小批量的随机选择。而SGD本身不包含噪声，工作Bayesian learning via stochastic gradient Langevin dynamics将SGD与Langevin动力学结合形成的SGLD方法，并在经典的最大后验估计(Maximum-A-Posteriori, MAP)背景下进行了比较。

详细的参数解释可见上述引用的工作。

3. 斯坦因分数函数（(Stein’s) Score Function）

3.1 分数函数的定义

对数似然函数的梯度${\nabla }_x\log p(x)$是朗之万等式迭代过程中的关键，其中${\nabla }_x\log p(x)$正式的名称是斯坦因分数函数(Stein’s score function)：

$$s_{\theta }(x)\stackrel{def}{=}{\nabla }_x\log p_{\theta }(x)$$

需要注意的是，应当将Stein’s score function与一般的分数函数区分开，一般的分数函数定义如下：

$$s_x(\theta )\stackrel{def}{=}{\nabla }_{\theta }\log p_{\theta }(x)$$

一般的分数函数是关于$\theta$的对数似然的梯度，而Stein’s score function是关于数据点$x$的。因此极大似然估计使用的是一般的分数函数，而Langevin dynamics使用的是Stein’s score function。但是，大部分扩散模型的文献中都将Stein’s score function直接称为score function。因此，在后面的论述中score function其实都是Stein’s score function。

Example 3.4 展示了一维高斯分布的score function

Example 3.5 展示了混合高斯分布的score function

上图用两个例子展示了概率密度函数和相关的score function的图形

3.2 分数函数的几何意义

理解分数函数的方法是铭记，它是相当于数据$x$的梯度。对于任何的高维分布$p(x)$，其梯度会形成一个向量场(vector field)：

• $\log p(x)$变化大的地方，梯度（分数）大，$\left|s(x)\right|$(分数值的绝对值)远离零。当$\log p(x)$接近峰值时，$\left|s(x)\right|$接近零。
• 向量场指示数据点如何在等高线图中移动。在下图中，展示了高斯混合模型（包含两个高斯分布）的等高线图，作者画出了箭头以指示向量场。现在我们想象一个数据点在向量场中，朗之万等式将数据点沿着向量场的方向移动到局部的极值点区域（高斯分布的峰值区域）。
• 在物理学中，分数函数等效于“漂移”。此名称表明，扩散颗粒应如何流到最低的能量状态。

4. 得分匹配技术（Score Matching Techniques）

最困难的问题是，在分布$p(x)$未知的情况下应该如何得到${\nabla }_x\log p(x)$。在这一部分，我们先学习一些已知的技术。

4.1 显式的得分匹配（Explicit Score Matching）

假设给定一份数据集$\chi =\left\{x^{(1)},\dots ,x^{(M)}\right\}$，人们提出利用经典的核密度估计来定义分布：

$$q_h(x)=\frac{1}{M}\sum _{m=1}^M\frac{1}{h}K\left(\frac{x-x^{(m)}}{h}\right)$$

其中$h$是核函数$K(\cdot )$的超参数。下图展示了核密度估计的基本想法，左侧的图展示了以不同数据点$x^{(m)}$为中心的多个核函数，这些单独核的总和就是总体的核密度估计出的分布$q(x)$。右侧图展示的是真实数据的直方图和相关的核密度估计。

因为分布$q(x)$是未知分布$p(x)$的近似，我们可以通过学习基于$q(x)$形式的分数函数$s_{\theta }(x)$。下面定义损失函数以训练网络，显式得分匹配的损失为：

通过代入核密度估计的式子，我们可以发现损失为：

到此，我们已经推导完用以训练网络的损失函数了。一旦我们训练完网络$s_{\theta }$，我们就可以将其代入朗之万等式中，递归地进行采样：

$$x_{t+1}=x_t+\tau s_{\theta }(x_t)+\sqrt{2\tau }z$$

显式得分匹配的问题是，核密度估计是对真实分布的非参数估计。特别地，当我们的样本数量有限并且样本处于高维空间中时，核密度估计的性能可能很差。

4.2 隐式的得分匹配（Implicit Score Matching）

在隐式得分匹配中，显式的得分匹配损失会被隐含的分数替换：

其中${\nabla }_x{s}_{\theta }(x)$是$s_{\theta }(x)$的雅可比矩阵，隐式的得分匹配损失可以通过蒙特卡洛近似：

其中${\partial }_i{s}_{\theta }(x^{(m)})=\frac{\partial }{\partial x_i}[s_{\theta }(x){]}_i=\frac{{\partial }^2}{\partial x_i^2}\log p(x)$。如果用深度神经网络实现分数函数的模型，则跟踪操作可能难以计算，因此隐式得分匹配不可扩展。

4.3 去噪的得分匹配（Denoising Score Matching）

鉴于显式和隐式得分匹配的缺点，现在引入一种更常用的得分匹配技术：去噪得分匹配(Denoising Score Matching, DSM)。在DSM中，损失函数定义为：

这里主要的不同就是用一个条件分布$q(x|x^{\prime })$代替了分布$q(x)$，后者需要一个估计（例如，核密度估计），但是前者并不需要。下面给出一个例子：

特殊地，我们令 $q(x|x^{\prime })=\mathcal{N}(x\mid x^{\prime },{\sigma }^2)$，可以让 $x=x^{\prime }+\sigma z$，这样可以得到：

因此,去噪声分数匹配的损失函数就变成了：

如果我们用$x$替换虚变量$x^{\prime }$，在给定训练数据集的前提下，$q(x)$的采样可以被$p(x)$的采样代替。总结得到:

上式的美妙之处就在于它具有高度的可解释性。$x+\sigma z$是在原图像$x$中有效地添加噪声$\sigma z$。分数函数$s_{\theta }$应该取这个噪声图像，并预测噪声$\frac{x}{{\sigma }^2}$。预测噪声等价于去噪声，因为任何去噪声后的图像加上预测噪声都会得到噪声图像。因此，上式是去嗓声步骤。

工作 A Connection Between Score Matching and Denoising Autoencoders说明了显式得分匹配ESM和去噪得分匹配DSM的等价关系，正是这种等价关系允许DSM能够估计分数函数。在上述工作中，证明了$J_{DSM}(\theta )$与$J_{ESM}(\theta )$仅差一个常数项的关系。

得分匹配模型中的训练过程通常是通过最大程度地减少去噪的得分匹配损失函数来完成的，训练步骤可以写为给定一个训练集 { x ( m ) } m = 1 M \{ {x^{(m)}}\} _{m = 1}^M {x(m)}m=1M​，我们要训练一个网络$\theta$，目标是:

下图说明了分数函数$s_{\theta }(x)$的训练过程：

上述训练过程假定固定噪声水平$\sigma$。将其推广到多个噪声水平并不困难，宋博士的工作噪声条件分数网络（NCSN）讨论可以用如下的损失来代替前面的目标：

其中，单个的损失函数是根据噪声水平${\sigma }_1,\dots ,{\sigma }_L$定义的：

系数函数$\lambda ({\sigma }_i)$通常根据经验选择为$\lambda (\sigma )={\sigma }^2$。同时，噪声水平序列通常满足条件：$\frac{{\sigma }_1}{{\sigma }_2}=\dots =\frac{{\sigma }_{L-1}}{{\sigma }_L}>1$

在推理阶段，假设我们已经训练了得分估计器$s_{\theta }$。为了生成图像，我们使用Langevin equation迭代地采样来去噪图像。在NCSN的情况下，可以通过退火重要性采样实现相应Langevin equation：

其中，${\alpha }_i={\sigma }_i^2/{\sigma }_L^2$是步长，$s_{\theta }(x_t,{\sigma }_i)$是噪声${\sigma }_i$的得分匹配函数。对于迭代次数$t$，每一步都会重复从$i=1$到$L$的${\sigma }_i$的计算，详细工作可以参照宋博士的工作噪声条件分数网络（NCSN） 。

总结

本次周报通过对得分匹配的朗之万动力学（SMLD）的学习，深入理解了模型从数据分布中采样的核心思想及其实现过程。核心要点包括朗之万等式的采样机制、斯坦因分数函数的定义及其在优化中的作用，以及得分匹配技术在不同场景下的实现方式。显式得分匹配和隐式得分匹配虽然提供了理论工具，但存在一定的局限性，而去噪得分匹配（DSM）通过引入噪声模型有效解决了这些问题，并在噪声条件分数网络（NCSN）中得到了具体应用。通过结合朗之万等式和分数匹配技术，SMLD为生成模型中的高效采样和生成提供了强大的理论支持。本次学习对于理解扩散模型的优化目标和生成过程具有重要意义，也为后续研究奠定了坚实的理论基础。