[论文解读] Sampling-free Inference for Ab-Initio Potential Energy Surface Networks

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|May 30, 2022

Machine Learning in Materials Science被引用 5

一句话总结

该论文提出PlaNet，一种无需采样的推理框架，通过在变分蒙特卡洛（VMC）优化过程中获得的噪声中间能量上训练代理图神经网络（GNN），实现了从头算势能面预测的加速。该代理模型可实现毫秒级能量推理——比VMC快7个数量级，同时保持化学精度。PESNet++通过引入受限、双占据轨道的归纳偏置改进了波函数建模，在N₂等复杂体系上将能量误差降低高达74%。

ABSTRACT

Recently, it has been shown that neural networks not only approximate the ground-state wave functions of a single molecular system well but can also generalize to multiple geometries. While such generalization significantly speeds up training, each energy evaluation still requires Monte Carlo integration which limits the evaluation to a few geometries. In this work, we address the inference shortcomings by proposing the Potential learning from ab-initio Networks (PlaNet) framework, in which we simultaneously train a surrogate model in addition to the neural wave function. At inference time, the surrogate avoids expensive Monte-Carlo integration by directly estimating the energy, accelerating the process from hours to milliseconds. In this way, we can accurately model high-resolution multi-dimensional energy surfaces for larger systems that previously were unobtainable via neural wave functions. Finally, we explore an additional inductive bias by introducing physically-motivated restricted neural wave function models. We implement such a function with several additional improvements in the new PESNet++ model. In our experimental evaluation, PlaNet accelerates inference by 7 orders of magnitude for larger molecules like ethanol while preserving accuracy. Compared to previous energy surface networks, PESNet++ reduces energy errors by up to 74%.

研究动机与目标

• 解决神经波函数方法中能量推理的高计算成本问题，该问题限制了其在大规模分子模拟中的应用。
• 通过用快速代理推理替代变分蒙特卡洛（VMC）中的蒙特卡洛积分，克服其计算瓶颈。
• 通过基于物理的架构约束，提升神经波函数在多种分子几何构型下的泛化能力和精度。
• 实现以往基于标准VMC的神经波函数无法实现的高分辨率、多维势能面探索。

提出的方法

• 利用VMC优化过程中获得的噪声中间能量估计值训练代理图神经网络（GNN），重新利用原本被丢弃的训练信号。
• 将代理训练建模为在线学习问题，以应对VMC优化过程中能量分布变化和标签噪声的问题。
• 实施时间序列数据聚合，以降低标签噪声的影响，提升代理模型的泛化能力。
• 提出PESNet++，一种受限的神经波函数模型，通过引入双占据轨道作为归纳偏置，灵感源自受限哈特ree-fock理论。
• 将PESNet++与PlaNet框架结合，实现在多种分子几何构型下快速、准确且可扩展的推理。
• 使用训练好的代理模型进行推理，完全绕过蒙特卡洛积分，将推理时间从数小时缩短至毫秒级。

实验结果

研究问题

• RQ1 能否有效重用VMC过程中的中间能量，训练出一种快速、准确的代理模型，用于神经波函数模型中的能量推理？
• RQ2 在使用噪声大、动态变化的VMC估计值训练代理模型时，其在实现数个数量级加速的同时，精度能保持多高？
• RQ3 在多构型设置下，引入物理归纳偏置（即双占据轨道）如何提升神经波函数的性能？
• RQ4 PlaNet框架能否实现对乙醇等较大分子的高分辨率、多维势能面探索？
• RQ5 在从头算量子化学中，用学习到的代理模型替代蒙特卡洛积分时，推理速度与精度之间的权衡如何？

主要发现

• PlaNet将乙醇每种几何构型的能量推理时间从约2.7小时缩短至118微秒，实现了7个数量级的加速，且精度无损失。
• 即使在乙醇等大体系上，代理模型训练所占总训练时间也少于1%，其中甚至不足0.2%。
• 与PESNet相比，PESNet++在氮分子（N₂）上将能量误差降低了高达74%，创下基于VMC的神经波函数的新SOTA纪录。
• PlaNet能够准确复现高分辨率、多维势能面（包括2D扭转扫描），这些在传统神经波函数方法中此前无法实现。
• 该框架使此前不可行的应用成为可能，如全势能面扫描和需要数十万次能量评估的分子动力学模拟。
• 在代理模型训练过程中引入时间序列数据聚合，有效缓解了中间VMC能量带来的噪声，实现了稳定且准确的模型收敛。