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

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 modelling 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.

Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?

“…They use the output of CASSCF (Complete Active Space Self Consistent Field, a sophisticated conventional quantum-chemistry method) as Ω and use a relatively small (∼ 100k weights) neural network for Λ. FermiNet on the other hand uses a simple exponential function as envelope Ω and uses a large (∼ 700k weights) neural network for Λ. Both approaches have been applied with great success to many different systems and properties, such as energies of individual molecules [4,9,8], ionization energies [9], potential energy surfaces [10,11], forces [10], excited states [12], model-systems for solids [13,14] and actual solids [15]. Several approaches have been proposed to increase accuracy or decrease computational cost, most notably architecture simplifications [16], alternative antisymmetric layers [17], effective core potentials [18] and Diffusion Monte Carlo (DMC) [19,20].…”

“…FermiNet commonly reaches lower (i.e. more accurate) energies than PauliNet [9], but PauliNet has been observed to converge faster [11]. It has been proposed [8] that combining the embedding of FermiNet and the physical prior knowledge of PauliNet could lead to a superior architecture.…”

“…Similar local geometries should generate similar local features, mostly independent of changes to the geometry far from the particle in question. Published architectures have so far not been able to address all three points: PauliNet [9] uses only distances as input features, making them invariant and local, but not sufficiently expressive, as demonstrated by [11]. FermiNet [9] uses raw distances and differences, making the inputs expressive and local, but not invariant under rotations.…”

“…FermiNet [9] uses raw distances and differences, making the inputs expressive and local, but not invariant under rotations. PESNet [11] proposes a global coordinate system along the principle axes of a molecule, making the inputs invariant and sufficiently expressive, but not local.…”

“…Effectively this amounts to applying a rotation matrix U J to the raw electron-nucleus differences: ρiJ = U J ρ iJ . As opposed to [11] where U is constant for all nuclei, in our approach U J can be different for every nucleus. These electronnucleus differences ρiJ are invariant under rotations, contain all the information contained in raw Cartesian coordinates and depend primarily on the local environment of an atom.…”