针对原子间作用力的几何增强预训练
针对原子间作用力的几何增强预训练
引言
分子动力学(MD)模拟在物理、化学、生物和材料科学等领域扮演着重要角色,为原子水平的过程提供了洞见。MD模拟的精确度和效率取决于所选择描述分子体系中原子相互作用的相互原子作用力(interatomic potentials)函数。经典MD使用经验公式,参数需要拟合,计算代价低但精度不够。而第一性原理MD则通过求解薛定谔方程获得精确的相互作用力,但计算量极大。因此,机器学习相互原子作用力(MLIPs)通过使用机器学习模型拟合第一性原理计算的能量和力而达到接近ab initio的精度以及较高效率,成为一种有前景的替代方案。
MLIPs的性能和通用性受限于标记数据的稀缺性,因为获取标记数据需要耗费巨大的第一性原理计算代价。各种自监督学习方法已被探索过,以从大量未标记数据中学习通用表征,再在有限标记数据上进行微调提取任务特定信息。然而,现有方法对于MLIPs领域的预训练数据集获取和预训练任务设计仍存在局限性。
论文概况
本文提出了一种针对MLIPs的几何增强自监督学习框架GPIP。该框架包含两个核心组成部分:
几何结构生成:利用经验力场进行经典分子动力学模拟高效生成大规模分子几何结构作为无标注预训练数据。
几何增强预训练:设计掩码、去噪和对比学习三个互补的自监督预训练任务,从生成的无标注结构数据中同时捕获拓扑和空间结构信息。
通过GPIP的两个步骤,MLIPs在仅花费少量计算资源的情况下,就能显著提高性能。该方法不依赖任何已有数据集,只需在目标分子系统上生成MD轨迹,故而避免了现有数据集所覆盖系统有限的问题,具有极好的普适性。
文章评估了GPIP在从小分子到复杂周期性系统等广泛基准上的表现,证明了该方法的有效性和鲁棒性。此外,文中还开发了一个新的电解液数据集,包含更多元素种类和复杂配置,用于更全面评测MLIPs的能力。
研究工作流程
a) 流程概述
使用经典分子动力学模拟生成目标分子体系的大量几何构象作为无标注数据。
对生成的无标注构象应用三种几何增强的自监督学习任务:掩码去噪、带噪掩码还原和对比学习,预训练图神经网络(GNN)捕获构象的拓扑和空间结构信息。
在第一性原理计算得到的少量数据上对预训练的GNN进行微调,学习任务相关信息。
b) 具体细节
无标注数据生成
文中在MD17、ISO17、液体水和电解液四个不同复杂程度的体系上,使用经典分子动力学软件LAMMPS结合经验力场(如OPLS-AA、TIP3P等),在不同温度下模拟MD轨迹,从中采样出大量分子构象作为无标注预训练数据集。
自监督学习任务
掩码去噪还原(Masked denoising):随机掩码部分原子特征,加噪其坐标,利用GNN从可见原子插值推理出被掩码原子的特征。
带噪掩码还原(Denoising with masked atoms):随机掩码部分原子特征,加噪整体构象坐标,训练GNN预测加性噪声而非还原原始构象,以强制捕获空间结构信息。
3D网络对比学习(Contrastive learning with 3D networks):构建一个3D网络捕获分子全局3D结构信息,通过最大化GNN与3D网络输出的互信息,使GNN也学习全局3D信息。
微调
在上述三种任务的无监督预训练后,将预训练得到的GNN在具有第一性原理标签的少量数据上进行监督微调,以进一步学习与任务相关的能量和力等信息。
c) 研究结论
GPIP能极大提升MLIPs在各类基准上的精度和泛化能力。
GPIP计算代价极低,相比增加第一性原理标签数据,更具有成本优势。
三种自监督任务互补,单一任务效果有限,结合则能有效捕获构象的拓扑和空间信息。
GPIP具有极好的普适性,不依赖任何先验数据集,避免了现有数据集所覆盖系统有限的问题。
d) 研究意义
科学意义:提出了一种廉价、高效、通用的MLIPs预训练范式,解决了预训练数据和任务设计的局限。
应用价值:提高了MLIPs在各类分子体系上的模拟精度,推进了MD模拟在多个领域的应用。
e) 研究创新亮点
无标注MD构象数据预训练思路创新,避免了昂贵第一性原理计算。
结合掩码、去噪、对比的多任务自监督学习框架设计独特。
开发了新的电解液数据集,更全面评测MLIPs能力。
实验评估体系全面,涵盖广泛基准和分子复杂程度。
该研究为MLIPs模型的低成本、高性能化提供了一种行之有效的解决方案,体现了自监督学习和分子模拟两个领域的创新。