Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions

1 INTRODUCTION In recent years , machine learning gained importance in computational quantum physics and chemistry to accelerate material discovery by approximating quantum mechanical ( QM ) calculations ( Huang & von Lilienfeld , 2021 ) . In particular , a lot of work has gone into building surrogate models to reproduce QM properties , e.g. , energies . These models learn from datasets created using classical techniques such as density functional theory ( DFT ) ( Ramakrishnan et al. , 2014 ; Klicpera et al. , 2019 ) or coupled clusters ( CCSD ) ( Chmiela et al. , 2018 ) . While this approach has shown great success in recovering the baseline calculations , it suffers from several disadvantages . Firstly , due to the tremendous success of graph neural networks ( GNNs ) in this area , the regression target quality became the limiting factor for accuracy ( Klicpera et al. , 2019 ; Qiao et al. , 2021 ; Batzner et al. , 2021 ) , i.e. , the network ’ s prediction is closer to the data label than the data label is to the actual prop- erty . Secondly , these surrogate models are subject to the usual difficulties of neural networks such as overconfidence outside the training domain ( Pappu & Paige , 2020 ; Guo et al. , 2017 ) . In orthogonal research , neural networks have been used as wave function Ansätze to solve the stationary Schrödinger equation ( Kessler et al. , 2021 ; Han et al. , 2019 ) . These methods use the variational Monte Carlo ( VMC ) ( McMillan , 1965 ) framework to iteratively optimize a neural wave function to obtain the ground-state electronic wave function of a given system . Chemists refer to such methods as ab-initio , whereas the machine learning community may refer to this as a form of self-generative learning as no dataset is required . The data ( electron positions ) are sampled from the wave function itself , and the loss is derived from the Schrödinger equation ( Ceperley et al. , 1977 ) . This approach has shown great success as multiple authors report results outperforming the tradi- tional ‘ gold-standard ’ CCSD on various systems ( Pfau et al. , 2020 ; Hermann et al. , 2020 ) . However , these techniques require expensive training for each geometry , resulting in high computational requirements and , thus , limiting their application to small sets of configurations . In this work , we accelerate VMC with neural wave functions by proposing an architecture that solves the Schrödinger equation for multiple systems simultaneously . The core idea is to predict a set of parameters such that a given wave function , e.g. , FermiNet ( Pfau et al. , 2020 ) , solves the Schrödinger equation for a specific geometry . Previously , these parameters were obtained by optimizing a separate wave function for each geometry . We improve this procedure by generating the parameters with a GNN , as illustrated in Figure 1 . This enables us to capture continuous subsets of the potential energy surface in one training pass , removing the need for costly retraining . Additionally , we take inspiration from supervised surrogate networks and enforce the invariances of the energy to physical symmetries such as translation , rotation , and reflection ( Schütt et al. , 2018 ) . While these symmetries hold for observable metrics such as energies , the wave function itself may not have these symmetries . We solve this issue by defining a coordinate system that is equivariant to the symmetries of the energy . In our experiments , our Potential Energy Surface Network ( PESNet ) consistently matches or surpasses the results of the previous best neural wave functions while training less than 140 of the time for high-resolution potential energy surface scans . 2 RELATED WORK . Molecular property prediction has seen a surge in publications in recent years with the goal of predicting QM properties such as the energy of a system . Classically , features were constructed by hand and fed into a machine learning model to predict target properties ( Christensen et al. , 2020 ; Behler , 2011 ; Bartók et al. , 2013 ) . Recently , GNNs have proven to be more accurate and took over the field ( Yang et al. , 2019 ; Klicpera et al. , 2019 ; Schütt et al. , 2018 ) . As GNNs approach the accuracy limit , recent work focuses on improving generalization by integrating calculations from computational chemistry . For instance , QDF ( Tsubaki & Mizoguchi , 2020 ) and EANN ( Zhang et al. , 2019 ) approximate the electron density while OrbNet ( Qiao et al. , 2020 ) and UNiTE ( Qiao et al. , 2021 ) include features taken from QM calculations . Another promising direction is ∆-ML models , which only predict the delta between a high-accuracy QM calculation and a faster lowaccuracy one ( Wengert et al. , 2021 ) . Despite their success , surrogate models lack reliability . Even if uncertainty estimates are available ( Lamb & Paige , 2020 ; Hirschfeld et al. , 2020 ) , generalization outside of the training regime is unpredictable ( Guo et al. , 2017 ) . While such supervised models are architecturally related , they pursue a fundamentally different objective than PESNet . Where surrogate models approximate QM calculations from data , this work focuses on performing the exact QM calculations from scratch . Neural wave function Ansätze in combination with the VMC framework have recently been proposed as an alternative ( Carleo & Troyer , 2017 ) to classical self-consistent field ( SCF ) methods such as Hartree-Fock , DFT , or CCSD to solve the Schrödinger equation ( Szabo & Ostlund , 2012 ) . However , early works were limited to small systems and low accuracy ( Kessler et al. , 2021 ; Han et al. , 2019 ; Choo et al. , 2020 ) . Recently , FermiNet ( Pfau et al. , 2020 ) and PauliNet ( Hermann et al. , 2020 ) presented more scalable approaches and accuracy on par with the best traditional QM computations . To further improve accuracy , Wilson et al . ( 2021 ) coupled FermiNet with diffusion Monte-Carlo ( DMC ) . But , all these methods need to be trained for each configuration individually . To address this issue , weight-sharing has been proposed to reduce the time per training , but this was initially limited to non-fermionic systems ( Yang et al. , 2020 ) . In a concurrent work , Scherbela et al . ( 2021 ) extend this idea to electronic wave functions . However , their DeepErwin model still requires separate models for each geometry , does not account for symmetries and achieves lower accuracy , as we show in Section 4 . Other efforts have aimed at accelerating Ansätze by replacing costly determinant operations , but this comes at a significant loss of accuracy ( Acevedo et al. , 2020 ) . 3 METHOD . To build a model that solves the Schrödinger equation for many geometries simultaneously and accounts for the symmetries of the energy , we use three key ingredients . Firstly , to solve the Schrödinger equation , we leverage the VMC framework , i.e. , we iteratively update our wave function model ( WFModel ) until it converges to the ground-state electronic wave function . The WFModel ψθ ( −→r ) : RN×3 7→ R is a function parametrized by θ that maps electron configurations to amplitudes . It must obey the Fermi-Dirac statistics , i.e. , the sign of the output must flip under the exchange of two electrons of the same spin . As we cover in Section 3.4 , the WFModel is essential for sampling electron configurations and computing energies . Secondly , we extend this to multiple geometries by introducing a GNN that reparametrizes the WFModel . In reference to meta-learning , we call this the MetaGNN . It takes the nuclei coordinates−→ Rm and charges Zm and outputs subsets ω , ωm ⊂ θ , m ∈ { 1 , . . . , M } of WFModel ’ s parameters . Thanks to message passing , the MetaGNN can capture the full 3D geometry of the nuclei graph . Lastly , as we prove in Appendix A , to predict energies invariant to rotations and reflections the wave function needs to be equivariant . We accomplish this by constructing an equivariant coordinate system E = [ −→e 1 , −→e 2 , −→e 3 ] based on the principle component analysis ( PCA ) . Together , these components form PESNet , whose architecture is shown in Figure 2 . Since sampling and energy computations only need the WFModel , we only need a single forward pass of the MetaGNN for each geometry during evaluation . Furthermore , its end-to-end differentiability facilitates optimization , see Section 3.4 , and we may benefit from better generalization thanks to our equivariant wave function ( Elesedy & Zaidi , 2021 ; Kondor & Trivedi , 2018 ) . Notation . We use bold lower-case letters h for vectors , bold upper-case W letters for matrices , −−−−→arrows to indicate vectors in 3D , −→r i to denote electron coordinates , −→ Rm , Zm for nuclei coordinates and charge , respectively . [ ◦ , ◦ ] and [ ◦ ] Ni=1 denote vector concatenations . 3.1 WAVE FUNCTION MODEL . We use the FermiNet ( Pfau et al. , 2020 ) architecture and augment it with a new feature construction that is invariant to reindexing nuclei . In the original FermiNet , the inputs to the first layer are simply concatenations of the electron-nuclei distances . This causes the features to permute if nuclei indexing changes . To circumvent this issue , we propose a new feature construction as follows : h1i = M∑ m=1 MLP ( W [ ( −→r i − −→ Rm ) E , ∥−→r i − −→ Rm∥ ] + zm ) , ( 1 ) g1ij = ( ( −→r i −−→r j ) E , ∥−→r i −−→r j∥ ) ( 2 ) where zm is an embedding of the m-th nuclei and E ∈ R3×3 is our equivariant coordinate system , see Section 3.3 . By summing over all nuclei instead of concatenating we obtain the desired invariance . The features are then iteratively updated using the update rule from Wilson et al . ( 2021 ) ht+1i = σ W tsingle hti , ∑ j∈A↑ gtij , ∑ j∈A↓ gtij + btsingle +W tglobal ∑ j∈A↑ htj , ∑ j∈A↓ htj , ( 3 ) gt+1ij = σ ( W tdoubleg t ij + b t double ) ( 4 ) where σ is an activation function , A↑ and A↓ are the index sets of the spin-up and spin-down electrons , respectively . We also add skip connections where possible . We chose σ : = tanh since it must be at least twice differentiable to compute the energy , see Section 3.4 . After LWF many updates , we take the electron embeddings hLWFi and construct K orbitals : ϕkαij = ( w kα i h LWF j + b kα orbital , i ) M∑ m πkαim exp ( −σkαim∥−→r j − −→ Rm∥ ) , πkαim = Sigmoid ( p kα im ) , σkαim = Softplus ( s kα im ) ( 5 ) where k ∈ { 1 , · · · , K } , α ∈ { ↑ , ↓ } , i , j ∈ Aα , and pi , si are free parameters . Here , we use the sigmoid and softplus functions to ensure that the wave function decays to 0 infinitely far away from any nuclei . To satisfy the antisymmetry to the exchange of same-spin electrons , the output is a weighted sum of determinants ( Hutter , 2020 ) ψ ( −→r ) = K∑ k=1 wk detϕ k↑ detϕk↓ . ( 6 )

Summary: The paper develops a neural network based variational ansatz for modeling wave functions. Authors build their model on top of FermiNet architecture with a few modifications: they use a different feature embedding approach that is invariant with respect to basic spatial symmetries. In addition, authors use a GNN “hypernetwork” that predicts the parameters of the variational wave function for a given configuration of the system. With these modifications authors show that the model can be optimized on a range of system configurations simultaneously. The optimization is performed in a standard VMC setting.  The results are compared on commonly studied problems such as variations of energy of the H4 molecule, hydrogen chain system, nitrogen molecule and cyclobutadiene. Authors find that their approach generally achieves similar results to those of a FermiNet model, while training on multiple system configurations simultaneously, reducing the time needed to obtain variational energies for multiple configurations.