用于原子间势主动学习的Stein核化分子动力学

# 基于Stein核化分子动力学的机器学习力场主动学习

来源：https://arxiv.org/html/2606.04100

Joanna Zou¹、Fraser Birks²、Dallas Foster³、Youssef Marzouk¹
¹MIT Schwarzman计算学院计算科学与工程中心
²华威大学工程学院华威预测建模中心
³NVIDIA

###### 摘要

机器学习原子间势（MLIP）能够实现高效精确的原子尺度模拟，但其性能在很大程度上取决于训练数据的质量与多样性。我们提出Stein核化分子动力学（SKMD）——一种增强采样方法，利用相互作用粒子动力学为MLIP的主动学习和微调获取信息丰富的训练构型。SKMD对应Stein变分梯度下降的随机变体，通过引入异步粒子更新和全局原子描述符核函数，使其适用于分子动力学场景；该核函数提供了一种具有对称性感知能力的构型相似度度量。与分子动力学中其他增强采样方法不同，SKMD将Boltzmann分布保留为动力学的渐近分布。这一性质在探索多样构型与吸引至能量landscape高概率区域之间实现了平衡。我们进一步提出一种高效在线数据获取方法，采用自适应停止准则，在模拟过程中选取非冗余训练数据。我们在Müller–Brown势的神经网络模型主动学习以及丙氨酸二肽的MACE原子间势微调两个任务上对SKMD进行了验证。与主动学习基线方法相比，在相同训练样本数量下，我们的方法以更少的训练迭代次数取得了更高的模型精度。

## 1 引言

原子尺度化学现象模拟领域的诸多现代进展，源于机器学习原子间势（MLIP）的应用——这类数据驱动的原子力场代理模型能够以远超*从头算*方法可行范围的系统规模和时间尺度开展分子动力学（MD）模拟。MLIP的精度在很大程度上取决于训练数据的质量：训练构型必须能够代表系统的关键热力学态以及连接这些状态的过渡态，MLIP才能正确刻画化学性质。这些过渡态或未观测热力学态的训练数据难以获取，原因在于模拟过程中转变事件发生频率极低。此外，利用量子力学参考计算标注数据的高昂代价限制了可以加入训练集的样本数量。

在MLIP的**主动学习**中，我们通过交替进行训练数据收集和基于增广数据集重新训练模型，逐步提升模型精度。现有大量MLIP主动学习方法利用子集选择方法——例如基于D-最优设计\[47\]、CUR分解\[10\]、MaxVol算法\[38, 37\]、高斯过程方差\[28, 58, 59, 62, 61\]、熵最大化\[29, 52\]、委员会查询\[3\]以及行列式点过程\[66\]——从模拟MD轨迹中筛选信息丰富的子集加入训练集。然而，标准MD轨迹可能陷于能量盆地，无论采用何种子集选择技术，都会产生高度相关的数据，从而限制模型的改进空间。

为此，近期的主动学习方法在分子动力学中引入**增强采样**来促进对构型空间新区域的探索，包括元动力学\[26\]、不确定性驱动动力学\[31\]以及超主动学习\[57\]。这些方法引入自适应偏置力，驱动动力学走向构型空间中代表性不足的区域，并定义采集准则以在模拟过程中选取非冗余训练数据。然而，偏置动力学不能保持与MLIP相关联的Boltzmann分布的保真性，且采集准则通常不考虑底层能量landscape。因此，所选训练构型可能无法代表物理上有意义的构型或热力学态的真实分布。

我们提出Stein核化分子动力学（SKMD）来解决上述问题，这是一种用于MLIP主动学习的新型增强采样方法。我们的核心思想是将贝叶斯推断和统计学中的变分推断方法迁移至分子动力学中的采样问题。SKMD源自Stein变分梯度下降（SVGD）\[34\]——一种基于粒子的变分推断算法，利用相互作用粒子集来近似目标分布。我们的方法优于其他增强采样方法之处在于，将MLIP的Boltzmann分布保留为动力学的渐近分布。此外，SKMD偏置力提供了一种手段，可定义采集准则，在选取多样构型与能量landscape指导的构型之间取得平衡。

主动学习与模型微调密切相关，后者通过奖励函数所指定的数据空间区域来提升模型输出精度。基于流的生成方法虽已被用于Boltzmann采样和微调\[45, 46, 44\]，但这些方法需要在目标增强采样区域中已存在训练数据，且在数据覆盖不足的区域往往难以充分采样。我们认为，增强采样框架更适合主动学习任务，因为局部粒子变换能够发现现有训练数据未曾覆盖的热力学态。

我们的主要贡献总结如下：

- 我们提出SKMD，作为SVGD的随机变体，通过异步粒子更新和全局原子描述符核函数实现，从而使算法适用于分子动力学场景。
- 我们证明SKMD动力学的渐近分布是系统的Boltzmann分布。在命题 1 中，我们证明其平均场极限与SVGD的平均场极限一致，后者在适当条件下收敛至Boltzmann分布。
- 我们以自适应停止准则的形式开发了一种在线数据获取方法，详见附录C.3节和附录D.1节。
- 我们证明SKMD在数据生成和MLIP主动学习方面优于其他采样技术，并在Müller-Brown势神经网络模型学习和有机分子MACE基础模型对丙氨酸二肽的微调两个问题上进行了验证。

## 2 背景

### 2.1 机器学习原子力场

**原子间势**将由 $N$ 个原子组成的构型的总势能建模为原子位置 $\mathbf{x}=(x_{1},\ldots,x_{N})\in\mathbb{R}^{3N}$ 的函数，其中 $x_{n}\in\mathbb{R}^{3}$。经典经验势\[53, 50, 14, 6\]是具有简单参数形式的解析函数，而机器学习原子间势（MLIP）是从密度泛函理论（DFT）计算等高保真参考数据中学习得到的灵活函数近似。设 $V:\mathbb{R}^{3N}\to\mathbb{R}$ 和 $-\nabla_{\mathbf{x}}V:\mathbb{R}^{3N}\to\mathbb{R}^{3N}$ 分别对应参考DFT计算的势能和力。MLIP $V_{\theta}:\mathbb{R}^{3N}\to\mathbb{R}$ 通常通过加权最小二乘目标函数学习得到，其中模型参数 $\theta\in\Theta$ 是损失函数 $\mathcal{L}$ 的最小化解：

$$\mathcal{L}(\theta)=\frac{\lambda_{0}}{K}\sum_{k=1}^{K}\|V_{\theta}(\mathbf{x}^{k})-V(\mathbf{x}^{k})\|^{2}+\frac{\lambda_{1}}{K}\sum_{k=1}^{K}\|\nabla_{\mathbf{x}}V_{\theta}(\mathbf{x}^{k})-\nabla_{\mathbf{x}}V(\mathbf{x}^{k})\|^{2}\ ,\tag{1}$$

其中 $\lambda_{0},\lambda_{1}>0$，在训练集 $\mathcal{D}\coloneqq\{(\mathbf{x}^{k},V(\mathbf{x}^{k}),\nabla_{\mathbf{x}}V(\mathbf{x}^{k}))\}^{K}_{k=1}$ 上求值。

**原子描述符**是局部原子环境的特征表示，构成许多MLIP的基础。局部描述符 $\tilde{g}(\mathbf{x})\in\tilde{\Omega}\subseteq\mathbb{R}^{d}$ 可以从数据中学习，如来自NequIP\[9\]、Allegro\[43\]或MACE\[8, 7\]等GNN势的不变潜在表示；也可以从局部原子近邻的对称性适配基函数显式构造，以强制满足SO(3)下的不变性，如SOAP描述符\[5\]、双谱分量\[54\]和ACE基函数\[18\]。局部描述符是逐原子的表示，而**全局描述符** $g(\mathbf{x})\in\Omega\subseteq\mathbb{R}^{d}$ 是多原子构型的表示，通常是局部描述符的组合 $\tilde{g}(\mathbf{x})=[\tilde{g}^{1}(\mathbf{x}),\ldots,\tilde{g}^{N}(\mathbf{x})]\in\tilde{\Omega}^{N}\subseteq\mathbb{R}^{Nd}$。文献\[29\]中使用的全局描述符示例之一是局部描述符的均值 $g(\mathbf{x})=\frac{1}{N}\sum_{n=1}^{N}\tilde{g}^{n}(\mathbf{x})$。根据具体应用，更具信息量的全局描述符可以是部分原子（例如块体构型中的界面原子）的局部描述符均值。

**Boltzmann采样**在分子动力学中是指生成按原子系统构型Boltzmann分布分布的样本：

$$\pi(\mathbf{x})=\frac{1}{Z}\exp\big(-\beta V(\mathbf{x})\big)$$

其中 $Z$ 为归一化常数。给定一个比参考计算代价更低的MLIP，通常用MLIP的Boltzmann分布来近似 $\pi$：

$$\pi_{\theta}(\mathbf{x})=\frac{1}{Z_{\theta}}\exp\big(-\beta V_{\theta}(\mathbf{x})\big)$$

Boltzmann采样的常用方法是使用Langevin恒温器进行分子动力学模拟，对应欠阻尼Langevin动力学。若动力学满足遍历性，则位置的边缘不变分布与 $\pi_{\theta}$ 一致\[32\]。在Langevin动力学的过阻尼极限下，$\pi_{\theta}$ 仍然是位置的边缘不变分布，此时分子动力学方程为：

$$\textup{d}\mathbf{x}_{t}=-\nabla_{\mathbf{x}}V_{\theta}(\mathbf{x}_{t})\textup{d}t+\sqrt{2\beta^{-1}}\textup{d}W_{t}\tag{2}$$

可用于从Boltzmann分布采样。然而，在实际模拟时间尺度上，无论是欠阻尼还是过阻尼Langevin动力学的模拟，都不能保证从 $\pi_{\theta}$ 精确采样，因为动力学容易受到亚稳性的影响。由于分子系统的Boltzmann分布通常具有高度多模态性，轨迹可能在自由能landscape中被高势垒分隔的亚稳盆地中长时间受困，导致混合缓慢、全分布采样不足\[25\]。

### 2.2 基于粒子的变分推断

**Stein变分梯度下降（SVGD）**是一种基于粒子的变分推断算法，利用状态空间 $\mathcal{X}$ 上的相互作用粒子集 $\bar{X}_{t}=\{\mathbf{x}^{i}_{t}\}_{i=1}^{J}$ 的经验分布 $\hat{q}_{t}$ 来近似目标密度 $\pi$。对于 $i=1,\ldots,J$ 中每个粒子的演化，给定时间步长 $\epsilon>0$ 和对称半正定核函数 $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$，其更新方程为：

$$\mathbf{x}^{i}_{t+1}\leftarrow\mathbf{x}^{i}_{t}+\epsilon\hat{\phi}^{*}_{t}(\mathbf{x}^{i}_{t};\bar{X}_{t})\ ,\tag{3a}$$

$$\hat{\phi}^{*}_{t}(\cdot;\bar{X}_{t})=\frac{1}{J}\sum_{i=1}^{J}\Big[\nabla_{\mathbf{x}^{j}_{t}}\log\pi(\mathbf{x}^{j}_{t})k(\mathbf{x}^{j}_{t},\cdot)+\nabla_{\mathbf{x}^{j}_{t}}k(\mathbf{x}^{j}_{t},\cdot)\Big],\quad\mathbf{x}^{j}_{t}\in\bar{X}_{t}\ .\tag{3b}$$

可以证明，式(3a)对应描述粒子动力学的连续时间ODE的Euler离散化，式(3b)对应由Stein恒等式\[49\]得到的期望的Monte Carlo估计量：

$$\textup{d}\mathbf{x}^{i}_{t}=\phi^{*}_{t}(\mathbf{x}^{i}_{t})\textup{d}t\ ,\tag{4a}$$

$$\phi^{*}_{t}(\cdot)=\mathbb{E}_{\mathbf{x}\sim\hat{q}_{t}}[\nabla_{\mathbf{x}}\log\pi(\mathbf{x})k(\mathbf{x},\cdot)+\nabla_{x}k(\mathbf{x},\cdot)]\ .\tag{4b}$$

在时间 $t\to\infty$ 和粒子数 $J\to\infty$ 的极限下，经验分布 $\hat{q}_{t}$ 在Kullback–Leibler（KL）散度意义下弱收敛至 $\pi$\[35, 36\]。