手把手推导：从Score Function到Langevin采样，彻底搞懂SGM扩散模型的数学原理

手把手推导从Score Function到Langevin采样彻底搞懂SGM扩散模型的数学原理在生成模型领域Score-Based Generative ModelingSGM正以其独特的数学美感和理论深度吸引着越来越多的研究者。与常见的生成对抗网络GAN或变分自编码器VAE不同SGM通过直接建模概率密度的梯度场即Score Function来实现数据生成这一思路不仅优雅而且在实践中展现出了惊人的效果。本文将带您深入SGM的核心数学原理从Score Function的定义出发逐步推导到Langevin采样算法的实现细节让您真正理解这一强大模型背后的理论基础。1. Score Function概率密度的导航图1.1 Stein Score Function的定义与几何意义在概率论中给定一个概率密度函数p(x)其Stein Score Function定义为s(x) \nabla_x \log p(x)这个看似简单的定义蕴含着深刻的几何意义。Score Function实际上描述了概率密度函数在x点处的最陡上升方向。想象你站在一个概率密度构成的山丘上Score Function就是告诉你应该朝哪个方向走才能最快到达山顶即概率密度最大的区域。对于高斯分布N(μ, σ²I)我们可以具体计算出其Score Function\begin{aligned} p(x) \frac{1}{(2\pi)^{d/2}\sigma^d}\exp\left(-\frac{\|x-\mu\|^2}{2\sigma^2}\right) \\ \log p(x) -\frac{\|x-\mu\|^2}{2\sigma^2} C \\ s(x) \nabla_x \log p(x) -\frac{x-\mu}{\sigma^2} \end{aligned}这个结果清晰地展示了Score Function的方向性它总是指向分布的中心均值μ且随着距离中心越远其强度越大。1.2 加噪过程中的Score Function变化在SGM中我们通过逐步加噪的过程来构造一系列分布。设原始数据分布为q(x₀)在加噪步骤t时数据分布变为q(x_t|x_0) \mathcal{N}(x_0, \sigma_t^2 I)对应的边缘分布为q(x_t) \int q(x_t|x_0)q(x_0)dx_0根据1.1节的推导条件分布的Score Function为\nabla_{x_t} \log q(x_t|x_0) -\frac{x_t - x_0}{\sigma_t^2}然而我们真正需要的是边缘分布的Score Functions_t(x_t) \nabla_{x_t} \log q(x_t)这个无法直接计算因为q(x_t)的表达式涉及积分。这正是我们需要神经网络来近似Score Function的原因。2. Score Matching训练Score Network2.1 目标函数的推导我们的目标是训练一个网络sθ(xₜ,t)来近似真实的Score Function sₜ(xₜ)。直接的想法是最小化以下L₂损失\mathcal{L} \mathbb{E}_{q(x_t)} \left[ \| s_\theta(x_t,t) - \nabla_{x_t} \log q(x_t) \|^2 \right]但这个目标不可行因为我们无法直接计算∇ₓlog q(xₜ)。幸运的是我们可以利用Score Matching的技巧通过以下等价形式来避免直接计算\mathcal{L} \mathbb{E}_{q(x_t|x_0)q(x_0)} \left[ \| s_\theta(x_t,t) - \nabla_{x_t} \log q(x_t|x_0) \|^2 \right] C其中C是与θ无关的常数。根据1.2节的结果我们可以将目标函数具体化为\mathcal{L} \mathbb{E} \left[ \left\| s_\theta(x_t,t) \frac{x_t - x_0}{\sigma_t^2} \right\|^2 \right]2.2 与DDPM的联系有趣的是SGM的目标函数可以与DDPMDenoising Diffusion Probabilistic Models建立联系。考虑加权损失\mathcal{L}^* \sigma_t^4 \mathcal{L} \mathbb{E} \left[ \left\| \frac{x_t - x_0}{\sigma_t} \sigma_t s_\theta(x_t,t) \right\|^2 \right]令z (xₜ - x₀)/σₜ这是标准正态噪声zθ -σₜsθ(xₜ,t)则损失变为\mathcal{L}^* \mathbb{E} \left[ \| z - z_\theta \|^2 \right]这正是DDPM中预测噪声的目标函数这一发现揭示了两种看似不同的生成模型背后深刻的统一性。3. Langevin动力学与采样过程3.1 Langevin动力学的基本原理Langevin动力学提供了一种通过Score Function进行采样的方法。其离散形式为x_{i1} x_i \epsilon \nabla_x \log p(x_i) \sqrt{2\epsilon} z_i其中zᵢ ∼ N(0,I)ϵ是步长。这一更新公式可以理解为漂移项ϵ∇ₓlog p(xᵢ)推动样本向高概率区域移动扩散项√(2ϵ)zᵢ保证探索性避免陷入局部极值在适当条件下当ϵ→0且迭代次数→∞时xᵢ的分布会收敛到p(x)。3.2 SGM中的退火Langevin采样在SGM中我们采用退火(annealing)策略从高噪声级别逐步过渡到低噪声级别。具体算法如下初始化xₜ⁰ ∼ N(0,I)tT对每个t执行N步Langevin更新x_t^{i1} x_t^i \frac{\alpha_t}{2} s_\theta(x_t^i,t) \sqrt{\alpha_t} z_i将最终样本作为下一噪声级别的初始值xₜ₋₁⁰ xₜᴺt t-1重复步骤2-3直到t0其中αₜ是噪声级别相关的步长。这种退火策略有效地克服了传统Langevin采样在多模态分布中的混合问题。4. 实践中的技巧与优化4.1 噪声调度与步长选择在实践中噪声级别{σₜ}的选择至关重要。常见的策略包括调度类型公式特点线性调度σₜ σₘₐₓ - (σₘₐₓ-σₘᵢₙ)(t/T)简单直接指数调度σₜ σₘₐₓ(σₘᵢₙ/σₘₐₓ)^(t/T)更关注低噪声区域余弦调度σₜ cos(πt/2T)平滑过渡步长αₜ通常与σₜ²成正比以确保不同噪声级别的更新幅度相对一致。4.2 网络架构设计Score Network sθ(xₜ,t)通常采用类似UNet的结构但有以下特殊考虑时间嵌入将噪声级别t通过正弦位置编码嵌入网络噪声条件在残差块中加入与σₜ相关的缩放因子归一化使用Group Normalization而非Batch Normalization一个典型的实现片段可能如下class ScoreNetwork(nn.Module): def __init__(self): super().__init__() self.time_embed nn.Sequential( nn.Linear(1, 128), nn.SiLU(), nn.Linear(128, 256) ) self.conv1 nn.Conv2d(3, 64, 3, padding1) self.downblocks nn.ModuleList([DownBlock(64*(2**i)) for i in range(4)]) self.upblocks nn.ModuleList([UpBlock(64*(2**(3-i))) for i in range(4)]) def forward(self, x, t): t_emb self.time_embed(t.view(-1,1)) h self.conv1(x) skips [] for block in self.downblocks: h block(h, t_emb) skips.append(h) for i, block in enumerate(self.upblocks): h block(h, skips[-(i1)], t_emb) return h4.3 采样加速技巧为了提高采样效率可以应用以下技巧子序列采样不必遍历所有T个噪声级别而是选择子序列{t₁,...,tₖ}预测校正交替使用Langevin步骤和Score Matching校正高阶方法使用二阶或更高阶的数值积分方法在实践中通常需要在生成质量和计算效率之间进行权衡。例如使用50-100个噪声级别配合100-200步Langevin更新可以在合理时间内获得高质量样本。