[Paper Review] Generalizing Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|Feb 8, 2023

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper introduces Globe, a graph-learned orbital embedding method that enables neural wave functions to generalize across diverse molecules by learning localized electronic structures via spatial message passing in GNNs. It further proposes Moon, a size-consistent neural wave function that achieves faster convergence—up to 4.5× faster—than prior methods and maintains accurate, additive energy scaling across system sizes, enabling joint solution of Schrödinger equations for different molecules with a single model.

ABSTRACT

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. Globe learns representations of local electronic structures that generalize across molecules via spatial message passing by connecting molecular orbitals to covalent bonds. Further, we propose a size-consistent wave function Ansatz, the Molecular orbital network (Moon), tailored to jointly solve Schrödinger equations of different molecules. In our experiments, we find Moon converging in 4.5 times fewer steps to similar accuracy as previous methods or to lower energies given the same time. Further, our analysis shows that Moon's energy estimate scales additively with increased system sizes, unlike previous work where we observe divergence. In both computational chemistry and machine learning, we are the first to demonstrate that a single wave function can solve the Schrödinger equation of molecules with different atoms jointly.

Motivation & Objective

• To overcome the limitation of existing neural wave functions, which can only generalize across different geometries of the same molecule, but not across different molecular compositions.
• To enable a single neural wave function to solve the Schrödinger equation for multiple distinct molecules simultaneously, including those with varying numbers of atoms and orbitals.
• To achieve size consistency in neural wave functions, ensuring energy scales additively with system size, a key requirement for physical accuracy in quantum chemistry.
• To accelerate convergence in variational quantum Monte Carlo training by designing a wave function architecture that leverages spatial message passing and nuclear anchor points.

Proposed method

• Globe uses spatial message passing in graph neural networks to embed molecular orbitals into 3D space, enabling localized learning of electronic structure across different molecules.
• Orbital embeddings are learned by connecting molecular orbitals to covalent bonds, allowing the model to generalize across molecules with different numbers of orbitals.
• Moon, the wave function architecture, enforces size consistency by focusing on local electron-nucleus interactions and using nuclei as anchor points for message passing.
• The wave function is constructed as a sum of determinants with orbital weights reparametrized by Globe, enabling joint training across diverse molecules.
• The method uses a two-level neural network: one for reparametrizing orbital weights (Globe), and another for computing the final wave function (Moon).
• Training is performed via variational Monte Carlo with energy minimization, where Globe’s reparametrization is computed once per step, decoupling it from sampling and energy evaluation.

Figure 1: Globe steps. 1) Atom embeddings are obtained by message passing between atoms. 2) Orbital locations and embeddings are determined by core orbitals and molecular bonds. 3) Orbital embeddings are updated via a unidirectional message passing. For clarity, we omitted messages from each atom to

Experimental results

Research questions

• RQ1 Can a single neural wave function be trained to solve the Schrödinger equation for multiple distinct molecules with different atomic compositions?
• RQ2 How can size consistency be enforced in a differentiable, scalable neural wave function architecture?
• RQ3 Can spatial message passing in GNNs enable generalization of orbital embeddings across molecules with varying numbers of orbitals?
• RQ4 Does joint training of multiple molecules accelerate convergence compared to individual training?
• RQ5 Does the energy estimate of a neural wave function scale additively with system size when generalized across molecules?

Key findings

• Moon converges in 4.5 times fewer training steps than previous methods while achieving similar or lower energy estimates, demonstrating faster optimization convergence.
• Moon’s energy estimate scales additively with system size, unlike prior methods that exhibit divergent scaling, confirming its size consistency.
• Globe enables transfer to larger molecular systems without requiring additional self-consistent field (SCF) calculations, enabling direct evaluation on larger chains.
• On hydrogen chains, Globe with Moon maintains a constant energy per atom as system size increases, indicating robust extensivity and physical consistency.
• Globe with Moon achieves convergence on larger molecules like ethene and cyclobutadiene without pretraining, whereas other methods fail to converge without it.
• Moon outperforms FermiNet and matches PsiFormer in accuracy with 600k fewer parameters, demonstrating high parameter efficiency.

Figure 2: Illustration of Moon. 1) We initialize electron embeddings by aggregating their local neighborhood of electrons. 2) Nuclei aggregate electron embeddings via message passing. 3) Nuclei embeddings are iteratively updated. 4) Nuclei embeddings are structurally diffused towards the electrons v