[论文解读] Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|Oct 11, 2021

Machine Learning in Materials Science参考文献 47被引用 5

一句话总结

该论文提出PESNet，一种新颖的架构，将图神经网络（GNNs）与神经型波函数相结合，通过变分蒙特卡洛（VMC）方法在单次训练过程中同时求解多个分子几何构型下的薛定谔方程。通过使用GNN在单次训练中预测不同构型的波函数参数，PESNet相较于最先进的方法（如FermiNet）将训练速度提升了最高40倍，同时在势能面扫描中保持或超越其精度。

ABSTRACT

Solving the Schrödinger equation is key to many quantum mechanical properties. However, an analytical solution is only tractable for single-electron systems. Recently, neural networks succeeded at modeling wave functions of many-electron systems. Together with the variational Monte-Carlo (VMC) framework, this led to solutions on par with the best known classical methods. Still, these neural methods require tremendous amounts of computational resources as one has to train a separate model for each molecular geometry. In this work, we combine a Graph Neural Network (GNN) with a neural wave function to simultaneously solve the Schrödinger equation for multiple geometries via VMC. This enables us to model continuous subsets of the potential energy surface with a single training pass. Compared to existing state-of-the-art networks, our Potential Energy Surface Network PESNet speeds up training for multiple geometries by up to 40 times while matching or surpassing their accuracy. This may open the path to accurate and orders of magnitude cheaper quantum mechanical calculations.

研究动机与目标

• 解决在变分蒙特卡洛（VMC）模拟中为每个分子几何构型单独训练神经型波函数所带来的高计算成本问题。
• 通过单一神经网络实现实时势能面（PES）的连续建模，避免为每个构型重新训练。
• 将物理对称性（平移、旋转、反射）整合到波函数参数化中，以提升模型的泛化能力和不变性。
• 开发一种可扩展且高效的框架，在精度上匹配或超越现有神经型波函数模型的同时，大幅减少训练时间。

提出的方法

• 使用元图神经网络（MetaGNN）基于分子的核几何构型预测神经型波函数（WFModel）的参数。
• 通过VMC框架训练MetaGNN，使其生成的波函数参数能最小化每种输入构型的变分能量。
• 通过为GNN输入设计协变坐标系，强制实现对物理对称性（平移、旋转、反射）的不变性。
• 利用变分蒙特卡洛框架，通过电子构型的随机采样来优化波函数。
• 在单次前向传播中端到端训练整个PESNet架构，实现在多个分子几何构型上的同步学习。
• 采用共享的波函数架构（如受FermiNet启发的结构），其参数由GNN输出条件控制，从而实现在连续PES区域上的泛化能力。

实验结果

研究问题

• RQ1 单一神经网络架构能否同时准确预测多个分子几何构型的电子波函数？
• RQ2 与为每个构型单独训练最先进的波函数相比，基于GNN的波函数模型在多大程度上能减少训练时间？
• RQ3 通过协变坐标系将物理对称性整合到GNN输入中，能在多大程度上提升预测势能面的泛化能力和不变性？
• RQ4 该模型能否在无需重新训练的情况下泛化到连续PES区域内未见过的新构型？
• RQ5 在建模多构型势能面时，训练效率与预测精度之间存在怎样的权衡？

主要发现

• PESNet在包括H₄⁺、H₄、H₁₀、N₂和环丁二烯在内的多个体系中，相较于FermiNet将训练时间减少了最多47倍。
• 对于H₄⁺，PESNet在A100 GPU上仅需20 GPU小时，而FermiNet需127小时，且精度保持一致或更优。
• PESNet在精度和训练效率上均持续优于PauliNet和DeepErwin，H₄的训练时间分别为24小时与118小时，N₂的训练时间分别为89小时与4196小时。
• 该模型可在单次训练中捕获势能面的连续子区域，实现构型数量的线性可扩展性。
• PESNet在所有测试体系中均达到最先进精度，包括环丁二烯等复杂分子，且计算成本显著降低。
• 通过协变坐标系整合物理对称性，增强了模型的稳定性和泛化能力，同时未损害预测性能。