Stein Kernelized Molecular Dynamics for Active Learning of Interatomic Potentials¶

会议: NeurIPS 2026
arXiv: 2606.04100
OpenReview: https://openreview.net/forum?id=PyNQqAofSM
领域: 物理与科学计算
关键词: 机器学习原子间势函数、增强采样、Stein变分梯度下降、主动学习、玻尔兹曼分布

一句话总结¶

针对机器学习原子间势函数主动学习中采样易受亚稳态束缚且增强采样破坏平衡态分布的瓶颈，提出将Stein变分梯度下降引入分子动力学构建SKMD采样框架与自适应停止在线采集准则，在保持渐近玻尔兹曼分布的同时高效采集高代表性与高多样性的过渡态训练结构。

研究背景与动机¶

机器学习原子间势函数（MLIPs）通过拟合高精度从头算数据，在保持接近量子力学精度的同时极大拓展了分子动力学模拟的时空尺度。然而，MLIP的泛化能力极其依赖训练构型集的质量与多样性：模型不仅需要充分覆盖热力学平衡态盆地，还必须准确刻画构象转变过程中的高能鞍点和过渡态区域。由于分子系统中能垒普遍较高，标准的朗之万分子动力学模拟常受困于局域亚稳态势阱，在实际可行的模拟时长内产生高度时间相关的冗余样本，导致基于传统子集选择算法的主动学习效率极低。

为了驱使模拟探索构型空间中未见或未充分覆盖的新颖区域，近年研究引入了多种增强采样驱动的主动学习机制，如基于元动力学或模型委员会方差的不确定性驱动动力学（UDD）。然而，这类方法通过外加自适应偏置力强行推动系统越过能垒，导致偏差动力学的稳态测度不可逆地偏离了物理真实或势函数所对应的玻尔兹曼分布；同时其构型采集准则通常仅依据委员会认知不确定性，未与底层自由能面的物理几何特征耦合，容易在局部高方差狭窄区域反复聚集采样，选出偏离物理可行域的非代表性构型。

针对增强采样的探索能力与热力学物理保真度难以兼得的核心矛盾，本文的切入角度是将统计学中的粒子变分推断思想迁移至分子动力学体系中，利用交互粒子集协同刻画目标测度的几何输运性质。核心 idea：将 Stein 变分梯度下降适配为具异步更新与全局原子描述符核函数的动力学采样框架（SKMD），在数学上严格保证采样轨迹的渐近分布收敛于势函数玻尔兹曼分布，并以偏置力的欧氏范数构建同时耦合能面几何与构型差异度的自适应在线数据采集准则。

方法详解¶

整体框架¶

SKMD 将主动学习形式化为由构型空间协同采样、在线样本动态获取、高精度参考计算标注与势函数迭代微调组成的闭环体系。在采样阶段，系统维护由 \(J\) 个构型组成的交互粒子系，各粒子依托基于旋转平移不变全局原子描述符的高斯核函数感受彼此间的排斥斥力与沿能面梯度的引力；在单粒子推进时采用单链异步步进以消除跨节点通信瓶颈；在采集阶段，通过监控偏置力范数低于阈值的时刻自适应捕获处于势能特征点且远离已有构型的候选结构。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始构型集成与微调势函数<br/>X_s 与 V_theta"] --> B["全局原子描述符核构建<br/>计算局部环境对称不变核与雅可比"]
    B --> C["异步交互粒子演化<br/>单粒子交替步进与状态更新"]
    C --> D["Stein偏置力动力学步进<br/>能面梯度引力 + 核梯度斥力 + 各向同性噪声"]
    D --> E{"自适应在线采集准则<br/>偏置力范数 ||F_SKMD|| < 阈值 zeta_0 ?"}
    E -->|未达到阈值 / 步数不足| D
    E -->|满足停止条件| F["捕获高价值构型并送高精度标注<br/>更新集成并微调重训MLIP"]

关键设计¶

1. 全局描述符核函数与保对称性梯度传递

传统 SVGD 在笛卡尔坐标空间直接计算原子欧氏距离核，容易破坏分子体系固有的三维旋转、平移和全同原子置换不变性，且高维笛卡尔距离难以反映真实化学拓扑与局部成键相似性。为此，SKMD 引入基于局部原子环境特征聚合的全局对称性描述符 \(g(\mathbf{x}) \in \mathbb{R}^d\)（例如采用 MACE 各层不变标量通道连接或 SOAP 特征的全原子均值），并在描述符空间构建平移不变高斯核：

\[k(\mathbf{x}, \mathbf{x}') = a \exp\left(-\frac{\|g(\mathbf{x}) - g(\mathbf{x}')\|^2}{2b^2}\right)\]

在动力学演化中，核函数关于原子坐标的偏导通过链式法则严格耦合描述符的雅可比矩阵 \(\mathcal{J}(\mathbf{x}) \in \mathbb{R}^{d \times 3N}\)：

\[\nabla_{\mathbf{x}} k(\mathbf{x}, \mathbf{x}') = -\frac{1}{b^2} k(\mathbf{x}, \mathbf{x}') \mathcal{J}(\mathbf{x})^\top (g(\mathbf{x}) - g(\mathbf{x}'))\]

该设计使得粒子间的排斥力直接施加在具有物理意义的构象模式与几何自由度上，保证排斥动力学在分子对称流形上自然展开。

2. 异步粒子步进与各向同性随机热噪声

标准交互粒子变分系统需要在每一步同步更新全体 \(J\) 个粒子的物理位置并计算全局稠密 Gram 矩阵，在实际大规模并行分子动力学引擎（如 LAMMPS）中会产生极其昂贵的多进程通信开销与锁同步开销。SKMD 提出了异步单粒子演化协议：在时间步 \(s\)，保持其余 \(J-1\) 个粒子构型静态冻结充当环境参照集，仅将当前活跃粒子 \(i\) 沿连续时间推进 \(\ell\) 个积分步，演化完成后将其终态更新回集合并轮换下一粒子。

同时，演化方程中结合了无偏的各向同性高斯随机扰动，其完整单步更新离散式为：

\[\mathbf{x}_{t+1}^i \leftarrow \mathbf{x}_t^i + \epsilon \left[-A(\mathbf{x}_t^i)\beta \nabla_{\mathbf{x}_t^i} V_\theta(\mathbf{x}_t^i) + F_{\theta,s}^{\mathrm{SKMD}}(\mathbf{x}_t^i; \bar{X}_s)\right] + \sqrt{\frac{2\epsilon\eta}{J}} \xi_t^i\]

其中偏置力定义为集成内其他粒子的交互和：

\[F_{\theta,s}^{\mathrm{SKMD}}(\mathbf{x}; \bar{X}_s) = \frac{1}{J-1} \sum_{j \neq i} \left[-\beta \nabla_{\mathbf{x}_s^j} V_\theta(\mathbf{x}_s^j) k(\mathbf{x}_s^j, \mathbf{x}) + \nabla_{\mathbf{x}_s^j} k(\mathbf{x}_s^j, \mathbf{x})\right]\]

相较于早期随机梯度 SVGD 中依赖于核矩阵的状态相依扩散系数，固定各向同性扩散系数 \(\sqrt{2\eta/J}\) 在中小粒子规模（\(J \sim \mathcal{O}(10)\)）下展现出更优的相空间混合速率和数值稳定性；在连续时间、瞬时步进（\(\epsilon, \ell \to 0\)）与大粒子数（\(J \to \infty\)）极限下，扩散项的二阶 Itô 校正与局部鞅方差均按 \(\mathcal{O}(J^{-1})\) 衰减至零，从而严格收敛至 SVGD 的平均场非线性输运方程，其平稳解完全与玻尔兹曼分布 \(\pi_\theta \propto \exp(-\beta V_\theta)\) 吻合。

3. 基于 Stein 偏置力范数的自适应在线数据采集

传统的离线主动学习需先跑完大量轨迹生成候选构型池，再运行计算开销极大的行列式点过程（DPP）或贪心最远点采样（FPS）；而现有的在线 UDD 采集准则仅依赖多模型方差阈值，易在能垒上下的盲区导致数据聚集。SKMD 提出直接将偏置力的欧氏范数定义为在线采集打分函数：

\[\alpha_s(\mathbf{x}) = \|F_{\theta,s}^{\mathrm{SKMD}}(\mathbf{x}; \bar{X}_s)\|\]

数据捕获满足自适应停止时间 \(\ell(s; \zeta_0) \coloneqq \inf\{t \geq \ell_0 : \alpha_s(\mathbf{x}_{s+t}) < \zeta_0\}\)。该指标在物理与统计上有极为精妙的双重自洽性：第一项引力项包含负梯度 \(\nabla V_\theta\)，在势能极小值点、过渡态鞍点处自然趋近于零，促使算法优先截获具有明确热力学热点特性的几何构型；第二项斥力项包含核梯度，当粒子运动到远离已有集合 \(\bar{X}_s\) 的未探索相空间时，核相似度 \(k(\mathbf{x}_s^j, \mathbf{x})\) 极小，导致核梯度同样趋近于零。因此，偏置力范数收敛至低阈值 \(\zeta_0\) 恰恰标志着系统抵达了“远离已知密集区、同时落在能面物理特征结构上”的关键构型，天然抑制了构型冗余。

损失函数 / 训练策略¶

在完成在线构型收集并经过高精度密度泛函理论（DFT）或高阶基准模型打标后，获得扩增构型数据集 \(\mathcal{D} = \{(\mathbf{x}^k, V(\mathbf{x}^k), \nabla_{\mathbf{x}} V(\mathbf{x}^k))\}_{k=1}^K\)。势函数通过加权能量与力均方误差损失进行优化：

\[\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\]

在 MACE 基础势函数（MACE-OFF-23-small）微调丙氨酸二肽的场景中，训练采用低秩自适应（LoRA，秩为 4，缩放系数 \(\alpha=1.0\)）锁定骨干以高效微调。训练策略采用两阶段动态权重调度：前 50 个 epoch 设定力权重 \(\lambda_1=1000\)、能量权重 \(\lambda_0=40\)，学习率为 0.01，优先重构局域受力曲率；后 50 个 epoch 将力权重调整为 10、能量权重提升至 1000，学习率降至 0.001，精细对齐全局相对能量标定。

实验关键数据¶

主实验¶

论文分别在经典 2D 多势阱基准（Müller–Brown 势，测试集包含 50 组独立随机试验中依目标玻尔兹曼分布采样的结构）和代表性构象转变体系丙氨酸二肽（包含 300 个由 Oracle MACE 模型运行 SKMD 生成的高多样性保留测试集）上展开验证。

实验体系	评测指标	a-SKMD / SKMD (本文)	UDD / 最优基线	性能提升与收敛增益
Müller–Brown 势 (AL 第8轮)	能量 RMSE (arb. units)	0.18 (中位数)	~0.65 (UDD)	误差降低 72.3%，方差大幅缩减
Müller–Brown 势 (AL 第8轮)	受力 RMSE (arb. units)	1.25 (中位数)	~3.80 (UDD)	受力预测精度提升约 3 倍
丙氨酸二肽微调 (第10轮迭代)	能量 RMSE (meV)	1.8	4.2 (700K MD) / 7.5 (300K MD)	较 700K 提温采样误差降低 57.1%
丙氨酸二肽微调 (第10轮迭代)	原子受力 RMSE (meV/Å)	26.4	45.1 (700K MD) / 68.3 (300K MD)	受力场误差降低 41.5%

消融实验与算法变体对比¶

在 Müller–Brown 势上，对比了固定时间周期采样的朗之万动力学（Langevin）、不确定性驱动动力学（UDD）、固定间隔采样的标准 SKMD，以及引入自适应停止准则的 a-SKMD；并在纯采样设定下评估了不同异步步长 \(\ell\) 对玻尔兹曼分布逼近程度（Wasserstein-2 距离）的影响。

动力学采样方案 / 停止配置	采样策略	丙氨酸二肽相空间覆盖度 / W2 距离	能量 / 受力误差收敛趋势	说明
标准过阻尼朗之万 (Langevin)	周期采样 (每100步)	W2 误差高且方差极大	停滞在初始亚稳态势阱，误差无法下降	受限于能垒难以跳出初始盆地
不确定性动力学 (UDD)	自适应方差阈值 (\(\zeta_1=0.02\))	相空间局部聚集	误差下降缓慢，多次试验方差较大	不确定性峰值区易导致样本聚集重复
标准 SKMD (\(\ell=100\))	周期采样 (固定步长)	W2 距明显优于 Langevin	明显优于 UDD，平稳探索新盆地	验证了 Stein 交互力引导探索的有效性
a-SKMD (\(\ell_0=50, \zeta_0=1.0\))	自适应偏置力停止准则	最优 (广谱覆盖三大极小值点)	最快收敛至最低误差，方差最小	同时兼顾相空间多样性与势能极值代表性
SKMD 步长消融 (\(\ell=1\))	纯采样测试	\(W_2\) 距离最低 (~0.042)	理论平均场连续流极限	严格保证玻尔兹曼渐近保真度
SKMD 步长消融 (\(\ell=10^2\))	纯采样测试	\(W_2\) 距离适中 (~0.058)	兼顾单链异步效率与采样精度	实际运行的最佳计算平衡点
SKMD 步长消融 (\(\ell=10^3\))	纯采样测试	\(W_2\) 距离略升 (~0.089)	仍显著优于 Langevin (>0.25)	步长过大导致微观相关性略有增加

关键发现¶

在相同的主动学习轮次与总标注预算下，a-SKMD 的收敛速度远超 UDD，仅需 4 轮迭代即可达到 UDD 迭代 8 轮以上的精度，且 50 次试验间的四分位距极小，表现出极强的稳定性。
在丙氨酸二肽的 Ramachandran \((\phi, \psi)\) 扭转角平面上，标准 300K 朗之万模拟完全受困于左半球的主要能量盆地；即便加热至 700K 的高温模拟，也仅能探索部分过渡带；而 SKMD 实现了对三个局部极小值盆地（C7eq、Cax、beta2）以及跨越能垒通道的全覆盖。
计算效率方面，UDD 需维护由 5 个独立模型构成的委员会并在每个模拟步前向推理多次，单次主动学习作业耗时约 1 小时且内存峰值达 200GB；而 SKMD 仅需运行单个模型的前向梯度并结合轻量核计算，单任务耗时仅 10 分钟，显存与内存开销降低 60% 以上。

亮点与洞察¶

变分梯度流与统计力学的优雅交汇：以往增强采样（如元动力学）普遍需要事后通过复杂的指数重加权修正被扭曲的系综分布，容易引发巨大的重加权方差；SKMD 从数学上严格证明了平均场极限下其平稳测度天然就是未扭曲的玻尔兹曼分布，在采样驱动力和热力学保真度之间达成了内生平衡。
构象对称性内嵌至 Stein 核：避开了笛卡尔坐标的平移旋转敏感性，采用 GNN 潜空间不变原子特征或几何局部环境描述符构造全局核，并通过链式法则求导反向传递至坐标梯度，为高维分子多体势的主动学习提供了通用数学脚手架。
物理势能面与特征差异双驱动的在线准则：将偏置力范数作为截获构型的指标，巧妙地在单个物理量中统一了“寻找鞍点与极值盆地（引力小）”与“避开已采样本密集区（斥力小）”两重目标，避免了复杂昂贵的流形子集组合优化。

局限与展望¶

可变核带宽下的在线准则适配局限：作者指出自适应停止准则在固定高斯带宽 \(b\) 下表现优异，但在采用高维空间常用的中位数启发式可变带宽时，偏置力范数随时间推移衰减不明显，导致在线停止条件难以精准触发。高维复杂体系当前仍需依赖“探索阶段可变带宽 + 收敛阶段固定带宽”的阶段式工程策略。
大尺度体系核矩阵扩展开销：当交互集成规模 \(J\) 或原子数目 \(N\) 大幅增加时，虽然采用了异步更新，但描述符雅可比矩阵 \(\mathcal{J}(\mathbf{x})\) 的构建与多粒子链式反向传播仍带来额外的反向图显存占用。
未来方向：探索基于核 Stein 差异（KSD）近似最小化的直接采样控制，以及将 SKMD 拓展至固液界面、相变动力学及复杂溶剂化生物大分子全原子体系的势函数自适应构建。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将 Stein 变分梯度下降与分子动力学深度融合，并从理论上证明渐近玻尔兹曼分布一致性。
实验充分度: ⭐⭐⭐⭐☆ 覆盖解析势阱与 MACE 深度等变势函数的微调，50 次独立随机测试与消融严谨详实。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从平均场分析、鞅收敛证明到实际算法工程闭环一气呵成。
价值: ⭐⭐⭐⭐⭐ 为第一性原理分子动力学与机器学习势函数的主动学习提供了高质量、高保真的采样新范式。