第3章 得分匹配与去噪得分匹配
3.1 得分函数:概率密度的梯度
做分子生成,最直接的想法是——先学一个概率密度 p(x),然后从里面采样。但这条路走不通。为啥?因为归一化常数 Z 算不出来。你想想看,一个分子构型空间动辄上千维,配分函数根本没法求。
所以换了个思路:不学密度本身,改学密度对输入 x 的梯度。这个梯度就是 得分函数(Score Function):
s_θ(x) = ∇_x log p(x)
直观上,得分函数指向数据密度增加最快的方向。你沿着这个方向走,就能找到更「像真实分子」的构型。我记得第一次看到这个定义时,觉得这招挺妙的——绕开了归一化常数,直接去抓分布的「形状」。
我在项目中遇到过一个问题:直接用 p(x) 建模,训练时 loss 震荡得厉害。换成得分函数之后,训练稳定多了。原因是梯度操作天然去掉了常数项,模型不需要去拟合那个巨大的 Z 值。
3.2 Langevin Dynamics:从得分走到样本
学好了得分函数,怎么用它生成分子?Langevin Dynamics 就是那个采样引擎。它的更新公式很简单:
x_{t+1} = x_t + (ε/2) · s_θ(x_t) + √ε · z_t, z_t ~ N(0, I)
每一步做两件事:
- 沿得分方向走 —— 向高概率区域移动
- 加噪声探索 —— 防止陷入局部模式
说白了,就是「朝对的方向走,但别走太死」。步长 ε 控制移动速度,噪声项保证你能跳出局部陷阱。
我个人的习惯是,在分子生成任务里,初始 x_0 从标准高斯采样,然后迭代 200–500 步。步长 ε 设成 0.01 左右,再用一个指数衰减 schedule。嗯,这里要注意:步长太大,采样会发散;步长太小,收敛太慢。
3.3 得分匹配:让估计靠近真实
有了目标(得分函数)和采样方法(Langevin Dynamics),接下来就是怎么训练 s_θ(x)。得分匹配(Score Matching)给出了一个直接的损失:
J_SM(θ) = ½ ∫ p(x) ‖ s_θ(x) - ∇_x log p(x) ‖² dx
但这个损失里有 ∇_x log p(x) —— 又是真实密度,还是算不了。得分匹配通过分部积分,把它变成了一个等价形式:
J_SM(θ) = E_{p(x)} [ Tr(∇_x s_θ(x)) + ½ ‖ s_θ(x) ‖² ] + const
这下好了,只需要 s_θ(x) 的 Jacobian 迹,不需要真实密度。但是……高维数据下,算这个迹的代价非常高。每个样本都要算一次 Hessian 的迹,分子体系动辄几百个原子,根本扛不住。
我曾经在这个坑里趴了一周。训练一个 50 维的简单分子系统,显存直接爆了。后来换了去噪得分匹配,才把问题解决。
3.4 去噪得分匹配:更稳定的训练
去噪得分匹配(Denoising Score Matching, DSM)换了一个角度:对数据加噪声,然后估计噪声后的得分。加噪声后的数据分布 q_σ(x̃|x) 是一个已知的高斯分布:
q_σ(x̃|x) = N(x̃; x, σ²I)
它的得分函数有闭式解:
∇_x̃ log q_σ(x̃|x) = -(x̃ - x) / σ²
DSM 的目标函数就变成:
J_DSM(θ) = ½ E_{p(x)} E_{q_σ(x̃|x)} [ ‖ s_θ(x̃) + (x̃ - x)/σ² ‖² ]
你看,这里面没有 Hessian 迹,没有复杂的高维计算。只需要让 s_θ(x̃) 去拟合一个简单的解析表达式。训练稳定,计算高效。
我建议你把 DSM 当作默认方案。在分子生成任务里,噪声尺度 σ 的选择会影响生成质量。太小了去噪效果不够,太大了分子结构被破坏。我个人习惯用 σ = 0.1 起步,然后做多尺度训练。
下面这个流程总结了得分匹配和去噪得分匹配在训练与采样阶段的完整关系:
3.5 两种方法的对比
我把得分匹配和去噪得分匹配的核心差异整理成了表格,方便你对照:
| 对比维度 | 得分匹配 (SM) | 去噪得分匹配 (DSM) |
|---|---|---|
| 目标函数 | Tr(∇_x s_θ(x)) + ½‖s_θ(x)‖² | ½‖s_θ(x̃) + (x̃−x)/σ²‖² |
| 计算成本 | 高(Hessian 迹) | 低(解析梯度) |
| 高维适用性 | 差(维度平方增长) | 好(线性复杂度) |
| 训练稳定性 | 容易梯度爆炸 | 稳定 |
| 低密度区表现 | 差(无监督信号) | 好(噪声覆盖) |
| 分子生成推荐 | 不推荐 | 推荐 |
3.6 Langevin Dynamics 采样步骤
这里给出一个 Langevin Dynamics 在分子生成中的标准采样流程。我在项目里基本按这个模板跑:
- 初始化:从标准高斯分布采样初始构型
x_0 ~ N(0, I) - 设置 schedule:定义步长序列 {ε_t},通常从 0.1 指数衰减到 0.001
- 迭代更新:对 t = 0 到 T-1,执行 Langevin 更新
- 得分计算:用训练好的
s_θ(x_t)计算梯度方向 - 完成:T 步后得到生成样本
x_T
对应的 Python 伪代码:
def langevin_sampling(s_theta, x_init, eps_list, T):
x = x_init
for t in range(T):
eps = eps_list[t]
z = torch.randn_like(x)
x = x + (eps / 2) * s_theta(x) + (eps ** 0.5) * z
return x
嗯,这里要提醒一下:s_theta 返回的是梯度方向,不是概率值。很多人刚开始容易搞混,以为得分函数输出的是一个概率,其实它输出的是一个向量——指向高概率区的方向。
correct_step(Metropolis-Hastings 校正)可以保证采样严格收敛到目标分布。不过生成任务中通常不校,效率优先。
3.7 避坑指南
最后把我踩过的坑整理一下,希望你绕开:
- 步长选择:ε 太大 → 采样发散;ε 太小 → 收敛慢。我一般先用 ε=0.01 试跑 50 步,观察样本的范数变化,再调整。
- 噪声尺度 σ:DSM 中 σ 太小 → 去噪效果差;σ 太大 → 分子结构被破坏。分子生成任务建议从 σ=0.05 到 0.5 做多尺度训练。
- 网络输出范围:得分函数的输出可能非常大(尤其在低密度区)。我建议在输出层加一个
tanh或clip,防止梯度爆炸。 - 初始化方式:不要从全零初始化。我记得有一次偷懒从零初始化 Langevin Dynamics,结果跑了 300 步还在原点附近打转。换成高斯初始化后立刻正常了。
好了,这一章的核心就这些。得分函数是你理解扩散模型的钥匙,Langevin Dynamics 是生成引擎,DSM 是训练基石。把这些搞透,后面的章节你会走得很顺。