Generalizing Neural Wave Functions

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.