Recent neural network-based wave functions have achieved state-of-the-art accuracies in modeling ab-initio ground-state potential energy surface. However, these networks can only solve different spatial arrangements of the same set of atoms. To overcome this limitation, we present Graph-learned Orbital Embeddings (Globe), a neural network-based reparametrization method that can adapt neural wave functions to different molecules. We achieve this by combining a localization method for molecular orbitals with spatial message-passing networks. Further, we propose a locality-driven wave function, the Molecular Oribtal Network (Moon), tailored to solving Schr\"odinger equations of different molecules jointly. In our experiments, we find Moon requiring 8 times fewer steps to converge to similar accuracies as previous methods when trained on different molecules jointly while Globe enabling the transfer from smaller to larger molecules. Further, our analysis shows that Moon converges similarly to recent transformer-based wave functions on larger molecules. In both the computational chemistry and machine learning literature, we are the first to demonstrate that a single wave function can solve the Schr\"odinger equation of molecules with different atoms jointly.

翻译：近年来，基于神经网络的波函数在从头算基态势能面建模中取得了最先进的精度。然而，这些网络仅能求解相同原子集合的不同空间排布。为突破此局限，我们提出图学习轨道嵌入（Globe）方法——一种基于神经网络的重新参数化技术，可调整神经波函数以适应不同分子。该方法通过将分子轨道局域化方法与空间消息传递网络相结合实现。进一步地，我们提出一种局域性驱动的波函数——分子轨道网络（Moon），专为联合求解不同分子的薛定谔方程而设计。实验表明，在联合训练不同分子时，Moon达到与先前方法相似精度所需步数减少8倍，而Globe则实现了从小分子到大分子的迁移学习。此外，我们的分析显示，Moon在大分子上的收敛特性与近期基于Transformer的波函数相似。在计算化学与机器学习领域，我们首次证明单一波函数可联合求解含不同原子的分子薛定谔方程。