Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions
Authors: Nicholas Gao
Stephan Günnemann
Venue: ICLR 2022 Spotlight (2022)
Keywords: Graph Neural Networks
Computational Physics
Self-Generative Learning
Machine Learning for Science

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.

Published as a conference paper at ICLR 2022
AB-INITIO POTENTIAL ENERGY SURFACES BY PAIRING GNNS WITH NEURAL WAVE FUNCTIONS
Nicholas Gao & Stephan Günneemann
Department of Informatics & Munich Data Science Institute Technical University of Munich, Germany
{gaoni, guennemann ) @1 tum.de de
ABSTRACT
Solving the Schrödinger equation is key many quantum mechanical properties. However, analytical solution is only tractable for single-electron systems. Re- cently, neural networks succeeded at modeling 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 has
methods require tremendous amounts of computational resources train a separate model for each molecular geometry. In this work, we combine Graph Neural Network (GNN) with neural wave function to simultaneously solve the Schrödinger equation for multiple geometries via VMC. This enables us model continuous subsets of the potential energy surface with single training pass. Compared to existing state-of-the-ar networks, our Potential Energy Sur- face 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.
INTRODUCTION
In recent years, machine learning gained importance in computational quantum physics and chemistry to accelerate material discovery by approximating quan- tum mechanical (QM) calculations (Huang & von Lilienfeld, 2021). In particular, 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 let al.] 2014 Klicpera et al. 2019) coupled clus- ters (CCSD) (Chmicla et al. 2018). While this ap- proach has shown great success recovering the baseline calculations, suffers from several disad- vantages. Firstly, due the tremendous success of graph neural networks (GNNs) in this area, the regres- sion target quality became the limiting factor for accu- racy (Klicpera et al., 2019 Qiao et al., 2021 Batzner al.) 2021),i.e. the network's 's prediction is closer to the data label than the data label is to the actual QM property. Secondly, these surrogate models are subject to the usual difficulties of neural networks
MctaGNN
Wm
W
WFModel
Figure 1: Schematic of PESNet. For each molecular structure (top row), the MetaGNN takes the nuclei graph and parametrizes the WFModel via w and w.- Given these, the WFModel evaluates the electronic wave function (F)).
such overconfidence outside the training domain (Pappu & Paige 2020 Guo et al. 2017).
In orthogonal research, neural networks have been used as wave function Ansätze to solve the sta- tionary Schrödinger equation (Kessler et al.] 2021 Han et al., 2019). These methods use the vari- ational Monte Carlo (VMC) (McMillan 1965) framework to iteratively optimize neural wave function obtain the ground-state electronic wave function of given system. Chemists refer to such methods as ab-initio, whereas the machine learning community may refer to this as a form of self- self-generative learning dataset is required. The data (electron positions) sampled from the

---

Published as a conference paper at ICLR 2022
wave function itself, and the loss is derived from the Schrö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 re- quirements 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 Schrödinger equation for multiple systems simultaneously. The core idea is to predict of pa- rameters such that given wave function, e.g., FermiNet (Pfau et al.) 2020), solves the Schrödinger equation for a specific geometry. Previously, these parameters were obtained by optimizing sep- arate wave function for each geometry. We improve this procedure by generating the parameters with GNN, as illustrated in Figure T This enables us to capture continuous subsets 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 (Schütt et al. [00 8). While these symmetries hold for observable metrics such as energies, the wave function itself may not have these symme- tries. We solve this issue by defining coordinate system that is equivariant to the symmetries of the energy. 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 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 the energy of system. Classically, features constructed by hand and fed into a machine learning model to predict target properties (Christensen et al. 2020 Behler 2011: Bartók et al.) 2013). Lately, GNNs have proven be more accurate and took over the field Yang et al. 2019 Klicpera et al. 2019; Schü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 (020) and EANN (Zhang al. 2019) approximate the electron density while OrbNet (Qiao et al. (020)) and UNiTE (Qiao itt al. 2021) include features taken from QM calculations. Another promising direction is A-M models, which only predict the delta between a high- aechracc QM calculation and a faster low- accuracy one (Wengert et al.) (02T). Despite their success, surrogate models lack reliability. Even 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 first principles.
Neural wave function Ansätze in combination with the VMC framework have recently been pro- posed as an alternative (Carleo & Troyer, (017)) to classical self-consistent field (SCF) methods such as Hartree-Fock, DFT, CCSD to solve the Schrödinger equation (Szabo & Ostlund 2012). How- ever, early works were limited to small systems and low accuracy (Kessler et al.) 2021; Han et al. Choo 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. (202T) coupled FermiNet with diffusion Monte-Carle (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 concurrent work, Scherbela et al. (2021) extend this idea to electronic wave functions. However, their DeepErwin model still requires sepa- rate models for each geometry, does not account for symmetries and achieves lower accuracy, show in Section
3 METHOD
To build a model that solves the Schrödinger equation for many geometries simultaneously and accounts for the symmetries of the energy, we three key ingredients.
2

---

Published as a conference paper at ICLR 2022
WFModel::
MetaGNN:
Zer
Eubediing
MLEP
Ray
Esbobliax
Uplac
Σ
as det
xXF)
Interaction
Interacionn
Ban Rall
XBF
Isteracionn
Embedding: Re:
Update:
Interaction:
<WWic 6)
Envelope:
own
o(Wio
Pro
Wea Max
Wa++
MLP
Dealed
MLP
Menge
e1:3
MEP
Figure 2: PESNet's architecture is split into two main components, the MetaGNN and the WFModel. Circles indicate parameter-free free and rectangles parametrized functions, o/lo denotes the vector catenation, AT and A- denote the index sets of the spin-up and spin-down electrons, respectively. To avoid clutter, we left out residual connections.
Firstly, to solve the Schrödinger equation, we leverage the VMC framework, i.c., iteratively update our wave function model (WFModel) until it converges to the ground-state electronic wave function. The WFModel v T") Nx3
is function parametrized by that maps electron configurations to amplitudes. It must obey the Fermi-Diran statistics, i.e., the sign of the output must flip under the exchange two electrons of the same spin. As we cover Section 3.4 WFModel is essential for sampling electron configurations and computing energies.
Secondly, we extend this to multiple geometries by introducing a GNN that reparametrizes the WF- Model. In reference to meta-leamiin, we call this the MetaGNN, It takes the nuclei coordinates R., and charges Zn and outputs subsets w, wm C 0, 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 predict energies invariant rotations and reflections the wave function needs to be equivariant. We accomplish this by constructing equivariant coordinate system é 1, 2, é a] based on the principle component analysis (PCA).
Together, these components form PESNet, whose architecture is shown in Figure 2 Since sam- pling and energy computations only need the WFModel, single forward pass of the MetaGNN sufficient for each geometry during evaluation. Furthermore, its end-to-end differentiability facil- itates 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 for vectors, bold upper-case W letters for matrices, arrow's to indicate vectors in 3D, 7, to denote electron coordinates, Rm, Zm for nuclei coordinates and charge, respectively. [o, o] and [o] denote vector concatenations.
3.1 WAVE FUNCTION MODEL.
We use the FermiNet (Paau et al. 2020) architecture and augment it with a new feature construction that invariant reindexing nuclei. In the original FermiNet, the inputs 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:
M
h
MLP W ( Rm
E, 7
Rmmll
+ 2m
(1)
m=1
9
((
7. E. 7. - -7,11)
(2)
3

---

Published as a conference paper at ICLR 2022
where 2m is embedding of the m-th nuclei and E R3x3 is our equivariant coordinate sys- tem, 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)
h 1+1 o
W single
h
Σ
gijs
Σ gi
Usingl + W
global
Σ JEA
h2
(3)
JEAT
jEAL
jEAt
9 = a (W1 9ij
+
b &asubble)
(4)
where a is an activation function, A and A are the index sets of the spin-up and spin-down elec- trons, respectively. We also add skip connections where possible. We chose a lanh since must be at least twice differentiable to compute the energy, see Section 3.4 After LWF many updates, we take the electron emheddings Low and construct K orbitals:
(w ko h
Σ
R sm exp(-oomi Ko P, Rmll),
ka
bebital,
Sigmoid(pmm), ka), where k {1,..... K}. a {f,4), i Aa, and p:, 8; free parameters. Here, the sigmoid and softplus functions to ensure that the wave function decays to O infinitely far away from any nuclei. To satisfy the antisymmetry the exchange of same-spin electrons, the output weighted sum of determinants (Hutter 2020)
kor
Softplus(s.m) kor)
π Tmm
Wk det r det s4
3.2 METAGNN
The MetaGNN's task is adapt the WFModel to the geometry hand. It does so by substituting subsets, w and Wm. of WFModel's parameters. While wm contains parameters specific to nuclei m, is set of nuclei-independent parameters such as biases. To capture the geometry of the nuclei, the GNN embeds the nuclei in a vector space and updates the embeddings via learning message passing. Contrary to surrogate GNNs, we also account for the position in our equivariant coor- dinate system when initializing the node embeddings to avoid identical embeddings in symmetric structures. Hence, our node embeddings are initialized by
1 GZ... Spos: RT,EE
where G is a matrix of charge embeddings, Zm E NN is the charge of nucleus m, fpos RR3 +
R Naar- N is our positional encoding function, and R E is the relative position of the mth nucleus in our equivariant coordinate system E (see Section 3.3 As positional encoding function, we use the spherical Fourier-Besse basis functions äSAF,In fn from Klicpera et al. 2019
foxx
( 21)) l8 Nor (I......
(8)
1=1
with Z, being the ith axis of our equivariant coordinate system E. Unlike Klicpera et al. (0019) are working the fully connected graph and, thus, neither include cutoff nor the envelope function that decays to at the cutoff.
A message passing layer consists of a message function fmmg and an update function Supdate. Together, one can compute an update the embeddings as
1++1 m
Supalate
I'm E..ESing
Fing (I'm 1' emn)
(9)
where Cmn is embedding of the edge between nucleus and nucleus We Bessel radial basis functions to encode the distances between nuclei (Klicpera et al. 2019). Both freg and fupdate realized by simple feed-forward neural networks with residual connections.
4

---

Published as a conference paper at ICLR 2022
After LGNN many message passing steps, compute WFModelI's parameters two levels. On the global level. four specific parameters:
gloobal outputs the biases of the network and, on the node level, J nodo outputs nuclei
LeNeN
w
D Osingleddobble
67/4
w
= fout globall
Σ pm
(10)
m
1=l
Wm
2m
1.1/L m
Pm
1.t/L
,....
node
(
t=1
We use distinct feed-forward neural networks with multiple heads for the specific types of parame- ters estimated implement mode and four
3.3 EQUIVARIANT COORDINATE SYSTEMS
Incorporating symmetries helps reduce the training space significantly. In GNNs this is done by only operating on inter-nuclei distances without clear directionality in space, i.e., without V, coordinates. While this works for predicting observable metries such energies, it does not work for wave functions. For instance, any such GNN could only describe spherically symmetric wave functions for the hydrogen atom despite all excited states (the real spherical harmonics) not having such symmetries. Unfortunately,
show in Appendix [ recently proposed equivariant GNNs (Thomas et al. 2018 Batzner et al. (021)) also suffer from the same limitation.
To solve this issue, introduce directionality in the form of a coordinate system that is equivariant to rotations and reflections. The axes of our coordinate system E = c é,, 2'll are defined by
the principal components of the nuclei coordinates, PCA PCA E2 e PCA Using PCA is robust to reindexing nuclei and ensures that the axes rotate with the system and form an orthonormal basis. However, as PCA only returns directions up sign, have resolve the sign ambiguity. We do this by computing an equivariant vector o. i.e, vector that rotates and reflects with the system, and defining the direction of the axes as
PCA PCA
if ग PCA 2 0,
(11)
clse.
As equivariant vector we use the difference between a weighted and the regular center mass
M
1
Σ
Rm
12
Zm R
(12)
M
m=1
M
R
Rm-
(13)
With this construction, we obtain an equivariant coordinate system that defines directionality in space. However, are aware that PCA may not be robust, e.g., if eigenvalues of the covariance matrix are identical. detailed discussion on such edge cases can found Appendix C
3.4 OPTIMIZATION
We the standard VMC optimization procedure (Ceperley et al. 1977) where seek to minimize the expected energy of a wave function tho parametrized by A:
wvo H (/o (ve (hbe)
L
(14)
where H is the Hamiltonian of the Schrödinger equation
1 N
M
M
M
Zm Zn
2
Rm
Rmm
R.!l
v(7)
(15)
5

---

Published as a conference paper at ICLR 2022
with V being the Laplacian operator and V(T) describing the potential energy. Given samples
from the probability distribution ~ ( one can obtain an unbiased estimate of the gradient
Eo = v00
(7)
H Hool (
N
0² log Woo (T))
a aog Ioo 2)
(16)
+ +V(T),
2
07
0T sk
A=1
VeL= B7~03
(E((T) B [Eo (p)) Ve log logl (800 7)]
(17)
where Ee( r') denotes the local energy of the WFModel with parameters 0 for the electron config- uration F. One can see that for the energy computation, we only need the derivative the wave function w.r.t the electron coordinates. As these
inputs the MetaGNN, to differentiate through the MetaGNN to obtain the local energies. We clip the local energy as in Pfau al. (2020) and obtain samples from ~ v 7²) via Metropolis-Hastingg The gradients for the MetaGNN are computed jointly with those of the WFModel by altering Equation 17
do not have
VeC Ex (Eo((T) - E Lo [Eo(() (7)) Ve log Itee [
(18)
where θ is the joint set of WFModel and MetaGNN parameters. To obtain the gradient for multiple geometries, we compute the gradient as in Equation 18 multiple times and average. This joint gra- dient of the WFModel and the MetaGNN enables to single training pass to simultaneously solve multiple Schrödinger equations.
provides with raw estimate of the gradient, different techniques have been used to construct proper updates to the parameters (Hermann et al. 2020) Pfau et al.) (2020). Here, use natural gradient descent enable the use of larger learning rates. So, instead of doing regular gradient descent step in the form of @ = o VeL, where 1) is the learning rate, we add the inverse of the Fisher information matrix as a preconditioner
While Equation 18
r++ = O' nf FVF'
(19)
F= = Ve log wee (ee (7) log / (7)1).
(20)
Since the Fisher F scales quadratically with the numbers of parameters, we approximate F VV via the conjugate gradient (CG) method [Ncuscamman et al. 2012). To determine the convergence of the CG method, we follow Martens (2010) and stop based on the quadratic error. To avoid tuning the learning rate 7). clip the norm of the preconditioned gradient F-1Va4 (Pascanu al. 2013) and use a fixed learning rate for all systems.
We pretrain all networks with the Lamb optimizer You et al. 2020) Hartree-Foci orbitals, i.e., we match each of the K orbitals to a Hartree-Food orbital of a different configuration. During pretraining, only the WFModel and the final biases of the MetaGNN optimized.
3.5 LIMITATIONS
While PESNet is capable of accurately modeling complex potential energy surfaces, we have not fo- cused on architecture search yet. Furthermore, as we discuss in Section PauliNet (Hermann et al.) 2020) still offers a better initialization and converges in fewer iterations than our network. Lastly, PESNet is limited to geometries the same set of nuclei with identical electron spin configurations, ic., to access properties like the electron affinity one still needs to train two models.
4 EXPERIMENTS
To investigate PESNet's accuracy and training time benefit, we compare it to FermiNet (Paau et al., 2020 Spencer et al., 2020). PauliNet (Hermann et al. 2020), and DeepErwin (Scherbela al. 2021 on diverse systems ranging from 3 to 28 electrons. Note, the concurrently developed DeepErwin was only recently released as a pre-print and still requires separate models and training for each configuration. While viewing the results energies one should be aware that, except for PESNet, all methods must be trained separately for each configuration resulting in significantly higher training times, discussed in Section

---

Published as a conference paper at ICLR 2022
-2016
-.844
-2.18
-1,888
22020
-1,8-9
4
CCSD CBS
F
2.002
(CSD(T) CBS
1.850
MRCI-D-F12 DrepErwin
FCT CRS
-1.611
-2.26
PeruiiNee
-1802
PESNet
PFSNet
Training domaiic
-1.00
-0.75
0.00
0.26
0.70
1.00
90
04
Reaction coordinate
Figure 3: The energy of HJ along the first re- action path (Alijah & Varandas, 2008). While PESNet and DeepErwin match the barrier height estimate of the MRCI-D-F12 calculation, PES- Net estimates = 0.27 m E, lower energies. Ref- erence data taken from Scherbela al. (202T).
Figure 4: Potential energy surface sean of the hydrogen rectangle. Similar to FermiNet, PES- Net does not produce the fake minimum at 90% Since PESNet respects the symmetries of the en- ergy, we only trained on half of the config space. Reference data is taken from Pfau et al. (2020).
Evaluation of ab-in ab-initio methods is still a challenge as true energies are rarely known, and ex- perimental data are subject to uncertainties. In addition, many energy differences may seem small due to their large absolute values, but chemists set the threshold for chemical accuracy to keal mol x 1.6 m Thus, seemingly small differences in energy are significant. Therefore, to put all results into perspective, we always include highly accurate classical reference calculations. When comparing VMC methods such as PESNet, FermiNet, PauliNet, and DeepErwin, interpreta- tion is simpler: lower is always better as VMC energies are upper bounds (Szabo & Ostlund 2012).
To analyze PESNet's s ability to capture continuous subsets of the potential energy surface, we train on the continuous energy surface rather than on discrete set of configurations for potential energy surface scans. The exact procedure and the general experimental setup are described in Appendix D
Additional ablation studics are available in Appendix
Transition path of H and weight sharing. Scherbela et al. 2021) use the first transition path of HJ (Alijah & Varandas, 2008) to demonstrate the acceleration gained by weight-sharing. But, they found their weight-sharing scheme to be too restrictive and additionally optimized each wave function separately. Unlike DecpErwin, our novel PESNet is flexible enough such that we do not need any extra optimization. In Figure 3] we see the DecpErwin results after their multi-step opti- mization and the energies of a single PESNet. We notice that while both methods estimate similar transition barriers PESNet results in ss 0.27 EA smaller energies which match the very accurate MRCI-D-F12 results (< 0.015 mBr).
Hydrogen rectangle and symmetries. The Hydrogen rectangle is known failure case for CCSD and CCSD(T). While the exact solution, FCI, indicates local maximum and CCSD(T), predict local minima. Figure 4 shows that VMC methods such as FermiNet and our PESNet do not suffer from the same issue. PESNet's energies are identical to FermiNet's (se 0.014 mE)) despite training only a single network on half of the configuration space thanks to our equivariant coordinate system.
909, both, CCSD)
The hydrogen chain is very common benchmark geometry that allows us to compare our method range of classical methods (Motta al.) (017) as well as to FermiNet, PauliNet, and DeepErwin. Figure 5 shows the potential energy surface of the hydrogen chain computed by range of methods. While our PESNet generally performs identical to FermiNet, we predict on average 0.31 F, lower energies. Further, notice that our results are consistently better than PauliNet and DeepErwin despite only training single model.
The nitrogen molecule poses challenge classical methods such as CCSD or CCSD(T) fail to reproduce the experimental results (Lyakh et al., 2012 Le Roy et al.] 2006). While the accurate r12-MR- ACPF method more closely matches the experimental results, scales factorially (Gdanitz) 1998). After [fau et al. (2020) have shown that FermiNet is capable of modeling such complex triple bonds, are interested in PESNet's performance. To better represent both methods,
7

---

Published as a conference paper at ICLR 2022
0.025
MRCI+Q-FLL F12 CBS
DeepErwin PouliNee
Expert
FeraiNa PESNet
0.000
ARQMC CIB
0.020
UOOSD CBS
UOCSD CBB
FermiNNtt PESNet
DOCSDIT; CRS
w/o MetalENN
0.015
UDCSD(T) CRRS
0.015
MR. ACFACFF
Fict in
-Eat
0.010
0.005
0.005
0.000
N
0.007
1.0
1.5
2.0
2.5
3.0
3.5
seperationn It oo
Boad length
Figure 5: Potential energy surface scan of the hydrogen chain with 10 atoms. We find our PESNet to outperform PauliNet and DeepErwin strictly while matching the results of FermiNet across all configurations. Reference data is taken from Hermann et al. (2020); Pfau al. (2020);
Figure 6: Potential energy surface sean of the nitrogen molecule. PESNet yields very simi- lar but slightly higher ( 0.37 m energies than FermiNet. Without the MetaGNN the ac- curacy drops significantly by R 4.3 m/h on av- erage. Reference data is taken from Le Roy al. 2006); Gdanitz (1998); Pfau et al. (2020).
Scherbela et al. (2021) Motta et al. (2017)
decided to compare both FermiNet as well as PESNet with 32 determinants due a substantial performance gain for both methods. The results in Fio that PESNet agrees very well with FermiNet and is on average just 0.37 m Fl higher, despite training only single model for less than 37 of FermiNet's training time, see Section 4.1 In addition, the ablation of PESNet without the MetaGNN shows a significant loss of accuracy of 4.3 mln on average.
Cyclobutadiene and the MetaGNN. The automerization of cyclobutadiene is challenging due to its multi-referrnce nature, i.e., single reference methods such as CCSD(T) overestimate the transi- tion barrier (Lyakh et al., 2012). In contrast, PauliNet and FermiNet had success at modelling this challenging system. Naturally, we are interested in how well PESNet can estimate the transition barrier. To be comparable to (Spencer et al. (2020), increased the number of determinants to 32 and the single-stream size to 512. Similar to PauliNet Hermann et af. (2020), we found PESNet to occasionally converge to a higher energy for the transition state depending on the initialization and pretraining. To avoid this, we pick well-initialized model by training 5 models for 1000 iterations and then continue the rest of the optimization with the model yielding the lowest energy.
As shown in Figure 7 all neural methods converge to the same transition barrier which aligns with the highest MR-CC results the upper end of the experimental range. But. they require different numbers of training steps and result in different total energies. PauliNet generally converges fastest, but results in the highest energies, whereas FermiNet's transition barrier converges slower but its energies are 70 F smaller. Lastly, PESNet's transition barrier converges similar PauliNet's, but its energies are 54 E. lower than PauliNet's, s, placing closer to FermiNet than PauliNet in terms of accuracy. Considering that PESNet has only been trained for 8 of FermiNet's time (see Sec- tion 4.1, we are confident that additional optimization would further narrow the gap to FermiNet.
In additional ablation study, compare to PESNet without the MetaGNN, i.c still train single model for both states of cyclobutadiene but without weight adaptation. While the results in Figure 7 show that the truncated network's energies continuously decrease, it fails reproduce the same transition barrier and its energies are 18 F worse compared to the full PESNet.
4.1 TRAINING TIME
While the previous experiments have shown that our model's accuracy on par with FermiNet, PESNet's main appeal is its capability to fit multiple geometries simultaneously. Here, we study the training times for all systems from the previous section. We compare the official JAX (Bradbury) et al. 2018) implementation of FermiNet (Spencer et al. 2220)) the official Py implementation of PauliNet (Hermann et al.) 2020), the official TensorFlow implementation of DeepErwin (Scher- bela et al. 2021), and our JAX implementation of PESNet. We use the same hyperparameters in the experiments or the defaults from the respective works. All measurements have been conducted on machine with 16 AMD EPYC 7543 cores and single Nvidia A A100 GPU. Here, we only mea-
8

---

Published as a conference paper at ICLR 2022
-14..00
-100.00
Ponlinet
-16..0
FerniNet
PESNet
-164.00
PESNet, mu GNN
-154.67
30
PamliNot
PermiNes
PRSNet, GNN
CCSD(T)
MR-0C
-10
expetimarat
10
10
Itenations
Figure 7: Comparison between the ground and transition states of cyclobutadiene. The top figure shows the total energy plotted in log scale zeroed at 154.68 En with light colors for the ground state and darker colors for the transition state. The bottom figure shows the estimate of the transition barrier. Both figures use a logarithmic x-axis. All neural methods estimate the same transition barriers in line with the highest MR-CC results the upper end of the experimental data. Reference energies are taken from Hermann et al. (2020): Spencer et al. (2020) Shen & Piecuch (0022).
H (Fig.)
H4 (Fig-
H10 (Fig.
5
N2 (Fig.
Cyclobutadiene (Fig.
PauliNet
43h
34h
153h
854h
437h
DeepErwin
34h
27h"
111h
FermiNet
127h"
118h
594h
4196h
2309h
PESNet
20h
24h
65h
89h
381h
Table 1: Total GPU (A100) hours to train all models of the respective figures. 'Experiments not included in the original works and timings are measured with the default parameters for the respective models. Larger molecules did not work with DeepErwin.
sure the VMC training time and explicitly exclude the time to perform the SCFs calculations any pretraining as these take up less than 1% of the total training time for all methods. Furthermore, note that the timings refer to training all models of the respective experiments.
the GPU hours train the models the section. is apparent that PESNet used the fewest GPU hours across all systems. Compared to the similar accurate FermiNet, PESNet is up to 47 times faster to train. This speedup is especially noticeable if many configurations shall be evaluated, e.g., 39 nitrogen geometries. Compared to the less accurate PauliNet and DeepErwin, PESNet's speed gain shrinks, but our training times are still consistently lower while achieving significantly better results. Additionally, for H4, H10, and N2, PESNet is not fitted to the plotted discrete set of configurations but instead on continuous subset of the energy surface. Thus, one is interested in additional configurations, PESNet's speedup grows linearly in the number of configurations. Still, the numbers in this section do not tell the whole story, thus, we would like to refer the reader to Appendix and G for additional discussion on training and convergence.
DISCUSSION
We presented novel architecture that can simultaneously solve the Schrödinger equation for mul- tiple geometries. Compared to the existing state-of-the-art networks, our PESNet accelerates the training for many configurations by up 40 times while often achieving better accuracy. The inte- gration of physical symmetries enables us to reduce our training space. Finally, our results show that a single model can capture continuous subset of the potential energy surface. This acceleration of neural wave functions opens access to accurate quantum mechanical calculations to broader audience. For instance, it may enable significantly higher-resolution analyses complex poten- tial energy surfaces with foresecable applications in generating datasets with unprecedented accuracy as well as possible integration into molecular dynamics simulations.

---

Published as a conference paper at ICLR 2022
Ethics and reproducibility. Advanced computational chemistry tools may have a positive impact in chemistry research, for instance in material discovery. However, they also pose the risk misuse, e.g,, for the development of chemical weapons. To the best of our knowledge, promote misuse any more than general computational chemistry research. To reduce the likelihood of such misuse, publish facilitate reproducibility, the source code includes simple scripts to reproduce all experiments from Section Furthermore, we provide detailed schematic of the computational graph in Figure 2 and additional details on the experimental setup including all hyperparameters in Appendix D
work does not
source code under the Hippocratic license (Ehmke) [0199] To
Acknowledgements. We thank David Pfau, James Spencer. Jan Hermann, and Rafael Reisenhofer for providing their results and data, Johannes Margraf, Max Wilson and Christoph Scheurer for helpful discussions, and Johannes Kliepera and Leon Hetzel for their valuable feedback.
Funded by the Federal Ministry of Education and Research (BMBF) and the Free State of Bavaria under the Excellence Strategy of the Federal Government and the Länder.
REFERENCES
Alexander Alijah and António J. C. Varandas. H4+: What do we know about it? The Journal of Chemical Physics, 12933),33303, July 2008. ISSN 0021-9606, 1089-7690. doi: 10.1063/1. 2953571.
Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B, 87(18))184115, May 2013. ISSN 1098-0121, 15500,,,155554 1550-235X. doi: 10.1103/PhysRevB.81 184115.
Simon Batzner, Tess Smidt, Lixin Sun, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Moli- nari, and Boris Kozinsky. SE(3)-Equivariant Graph Neural Networks for Data-Efficient and Ac- curate Interatomic Potentials. arXiv:2101.031 [cond-mat, physics:physics/. January 2021.
Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. The Journal of Chemical Physics, 134(7):074106, February 2011. ISSN 0021-9606, 1089-7690. doi: 00.1063/1.355371
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
Giuseppe Carleo and Matthias Troyer. Solving the Quantum Many-Body Problem with Artificial Neural Networks. Science, 355(6325):602-60) February 2017. ISSN 0036-8075, 1095-9203. doi: 10. 121//siiencc..ag2202..
D. Ceperley, G. Chester, and M. H. Kalos. Monte Carlo simulation of many-fermion study. Physical Review B, 16(7))3081-3099 October 1977. doi: 10. 10.1103/hyssR<:B..6.008 16.33.8
Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards ex- molecular dynamics simulations with machine-learned force fields. Nature Communications,
9(1):3887, September 2018. ISSN 2041-1723, doi: 10.
Kenny Choo, Antonio Mezzacapo, and Giuseppe Carleo. Fermionic neural-network states for ab- initio electronic structure. Nature Communications, 11(1):2368, December 2020. ISSN 2041- 1723. doi: 000003//41467-00000 1572--99-
Anders S. Christensen, Lars A. Bratholm, Felix A. Faber, and O. Anatole von Lilienfeld. FCHL revisited: Faster and more accurate quantum machine learning. The Journal of chemical physics, 154(4))044107, 2020.
Coraline Ada Ehmke. The Hippocratic License 2.1: An Ethical License for Open Source. hutps:/fiistdonohrmm.dy 2019.
http::/wwww.dmm.in..mm wwwp://wwwwwwwmini.....
10

---

Published as a conference paper at ICLR 2022
Bryn Elesedy and Sheheryar Zaidi. Provably Strict Generalisation Benefit for Equivariant Models. In Proceedings of the 38th International Conference on Machine Learning, pp. 2959 2959-2969. PML,, July 2021.
Robert J. Gdanitz. Accurately solving the electronic Schrödinger equation of atoms and molecules using explicitly correlated (r12-)MR-CI: The ground state potential energy curve of N2. Chem- ical Physics Letters, (83(5)25532661 February 1998. ISSN 0009-2614. doi: 10.1016/
50009-2614497)0132222
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On Calibration of Modern Neural Networks. In Proceedings of the 34th International Conference Machine Learning, pp. 1321- 1330. PMLR, July 2017.
Jiequn Han, Linfeng Zhang, and Weinan E. Solving many-clectron Schrödinger equation using deep neural networks. Journal of Computational Physics, 399:108929) December 2019. ISSN 0021-9991. doi: 10.10166f.jep.2019. 108929.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-networkk solution of the electronic Schrödinger equation. Nature Chemistry, 12(10):891-897, October 2020. ISSN 1755-4330, 1755-4349. doi: 10.10388/41155-0000054444
Lior Hirschfeld, Kyle Swanson, Kevin Yang, Regina Barzilay, and Connor W. Coley. Uncertainty Quantification Using Neural Networks for Molecular Property Prediction. Journal of Chemical Information and Modeling, 60(3):3770-3780, August 2020. ISSN 1549-9596. doi: 10.1021/acs.
jcim.0c00500
Bing Huang and O. Anatole Lilienfeld. Ab Initio Machine Learning in Chemical Compound Space. Chemical Reviews, 121((66:10001110066, August 2021. ISSN 0009-2665. doi: 10.1021/ www.chemmev.0000100.
Marcus Hutter. On Representing (Anti)Symmetric Functions. arXiv:2007.15294 [quant-ph]. July 2020.
Jan Kessler, Francesco Calcavecchia, and Thomas D. Kühne. Artificial Neural Networks as Trial Wave Functions for Quantum Monte Carlo. Advanced Theory and Simulations, 4(4):2000269, 2021. ISSN 2513-0390. doi: 10. (000//adts202200022)
Armagan Kinal and Piotr Piecuch. Computational Investigation of the Conrotatory and Disrotatory Isomerization Channels of Biccycoo[1...0jbun 1,0bbute to Outa-1,3-diene A Completely Renormalized Coupled-Cluster Study. The Journal of Physical Chemistry A, 111(4):734-72, February 2007.
ISSN 1089-5639, 1520-5215. doi: 10. 10.1021/jp065721K.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In 3rd Interna- lional Conference for Learning Representations, December 2014.
Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molec- ular Graphs. In International Conference on Learning Representations, September 2019.
Risi Kondor and Shubhendu Trivedi. On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups. In International Conference on Machine Learning, pp. 2747-2755. PMLR, July 2018.
George Lamb and Brooks Paige Bayesian Graph Neural Networks for Molecular Property Predic- tion. arXiv:2012.02088 [es, q-biol, November 2020.
Robert J. Le Roy, Yiye Huang, and Calvin Jary. An accurate analytic potential function for ground- state N2 from direct-potential-fit analysis of spectroscopic data. The Journal of Chemical Physics, 125(16):1643111 October 2006. ISSN 0021-9606, 1089-7690.. doi: 10. 10.1063/1.. 22354522.
Dmitry Lyakh, Monika Musiat, Victor F. Lotrich, and Rodney J. Bartlett. Multireference Nature of Chemistry: The Coupled-Cluster View. Chemical Reviews, 112(1):182-22) 2233, January 2012. ISSN 0009-2665, 1520-6890. doi: 10. 10110//20001111.

---

Published as a conference paper at ICLR 2022
James Martens. Deep learning via Hessian-free optimization. In Proceedings of the 27th Interna- tional Conference on International Conference on Machine Learning, ICML'10, pp. 735-742, Madison, WI, USA. June 2010. Omnipress. ISBN 001-66558-977771
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approxi- mate curvature. In Proceedings of the 32nd International Conference on International Conference on Machine Learning- Volume 37, pp. 2408-2417, 2015.
W. L. McMillan. Ground State of Liquid He4, Physical Review, 138(2A))444-A4551 April 1965.
doi: 10. 10.1103/PyysRev.. www.sRevv........ A442.
Mario Motta, David M. Ceperley, Garnet Kin-Lic Chan, John A. Gomez, Emanuel Gull, Sheng Guo, Carlos A. Jiménez-Hoyos, Tran Nguyen Lan, Jia Li, Fengjie Ma, Andrew J. Millis, Nikolay V. Prokof" ev, Ushnish Ray, Gustavo E. Scuseria, Sandro Sorella, Edwin M. Stoudenmire, Qiming Sun, Igor S. Tupitsyn, Steven R. White, Dominika Zgid, Shiwei Zhang, and Simons Collaboration on the Many-Electron Problem. Towards the Solution of the Many-Electron Problem in Real Materials: Equation of State of the Hydrogen Chain with State oaate-of-the-Aee Artt-hh-Arr Many-Body Methods. Physical Review X. 7(3):03 1059, September 2017. ISSN 2160-3308. doi: 10.1 Re.. 031059.
Eric Neuscamman, C. J. Umrigar, and Garnet Kin-Lic Chan. Optimizing large parameter sets in variational quantum Monte Carlo. Physical Review B, 85(4):045103, January 2012. ISSN 1098- 0121, 1550-235X. doi: 10. Do.113/PtysRevB.88..... 88..004503.
Aneesh Pappu and Brooks Paige. Making Graph Neural Networks Worth It for Low-Data Molecular Machine Learning. arXiv:2011.1220; [cs], November 2020.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning, pp. 1310-1318. PMLR, 2013.
David Pfau, James S. Spencer, Alexander G. D. G. Matthews, and W. M. C. Foulkes. Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, September 2020. doi: 10.1 httssPPhyyReeeeeeerrr2.........
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller III. OrbNet: Deep Learning for Quantum Chemistry Using Symmetry-Adapter Atomic- Orbital Features. The Journal of Chemical Physics, 153(12):124111, September 2020. ISSN
0021-9606, 1089-7690.
doi: 111066/5500229555
Zhuoran Qiao, Anders S. Christensen, Frederick R. Manby, Matthew Welborn, Anima Anandkumar, and Thomas F. Miller III. UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry. arXiv:2105. 14655 [physics], May 2021.
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, December 2014. ISSN 2052-4463. doi: 10. 10.1038/slata.2014.22
Michael Scherbela, Rafael Reisenhofer, Leon Gerard, Philipp Marquetand, and Philipp Grohs. Solv- ing the electronic Schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks. arXiv:2105.0835) [physics], May 2021.
K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko. and K.-.. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24): 241722, June 2018. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.501977
Jun Shen and Piotr Piecuch. Combining active-space coupled-cluster methods with moment energy corrections via the CC( P; Q methodology, with benchmark calculations for biradical transition states. The Journal of Chemical Physics, (6611):144104 April 2012. ISSN 0021-9606, 1089- 7690. doi: 10. 0360///.70002.
James S. Spencer, David Pfau, Aleksandar Botev, and W. M. C. Foulkes. Better, Faster Fermionic Neural Networks. 3rd NeurIPS Workshop on Machine Learning and Physical Science, November 2020.
12

---

Published as a conference paper at ICLR 2022
Attila Szabo and Neil S. Ostlund. Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory. Courier Corporation, 2012.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and transtation-equivrriant neural networks for 3D point [cs], May 2018.
clouds.
Masashi Tsubaki and Teruyasu Mizoguchi. Quantum Deep Field: Data-Driven Wave Function, Electron Density Generation, and Atomization Energy Prediction and Extrapolation with Machine Learning. Physical Review Letters, pp. 6, 2020.
Simon Wengert, Gábor Csányi, Karsten Reuter, and Johannes Margraf. Data-efficient machine learning for molecular crystal structure prediction. Chemical Science, pp- 10. N0.1039.D0SC0576633 2021. ISSN 2041-6520, 2041-6539.
doi: 10. 10.00//DSCCO576GG.
Max Wilson, Nicholas Gao, Filip Wudarski, Eleanor Rieffel, and Norm M. Tubman. Simula- tions of state-of-the-ar fermionic neural network wave functions with diffusion Monte Carlo. arXiv:2103. 12570 [physics, physics: March 2021.
Kevin Yang. Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, Andrew Palmer, Volker Settels, Tommi Jaakkola, Klavs Jensen, and Regina Barzilay. Analyzing Learned Molecular Representations for Property Prediction. Journal Chemical information and Modeling, 59(8):3370-333 August 2019. ISSN 1549-9596, 1549-960X. doi: 10. hacp///as.ccimm00200....
Li Yang, Wenjun Hu, and Li Li. Scalable variational Monte Carlo with graph neural ansatz. In NeurlPS Workshop on Machine Learming and the Physical Sciences, November 2020.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. Eighth International Conference on Learning Repre- sentations, April 2020.
Yaolong Zhang, Ce Hu, and Bin Jiang. Embedded Atom Neural Network Potentials: Efficient and Accurate Machine Learning with Physically Inspired Representation. The Journal of Physical Chemistry Letters, 10(17)4996-44667 September 2019. ISSN 1948-7185, 1948-7185. doi: 10. 1022//acs.ipcceet.9000777.
A INVARIANT ENERGIES REQUIRE EQUIVARIANT WAVE FUNCTIONS
Here, we prove that wave function needs to be equivariant with respect rotations and reflections the energy is invariant, Recall, our goal is to solve the stationary Schrödinger equation
Hag
Ev
(21)
where : 12 IR is the electronic wave function. Since the Hamiltonian H encodes the molecular structure via its potential energy term, see Equation 15 a rotation or reflection U 0(3) of the system results is a unitary transformation applied to the Hamiltonian H UHU Since the energy E is invariant to rotation and reflection, we obtain the transformed equation
UHU Exy.
(22)
One can see that if & is a solution Equation 21 the equivariantly transformed Equation 22
U solves
UHU FUV.
(23)
UH = EU
(24)
He = E
(25)
with the same energy.
13

---

Published as a conference paper at ICLR 2022
B EQUIVARIANT NEURAL NETWORKS As WAVE FUNCTIONS
Here, we want to briefly discuss why equivariant neural networks as proposed by Thomas et al. alternative
equivariant coordinate system. The issue is the same as for regular GNNs KKlicpera et al. 2019), namely that such networks can only represent spherically symmetric functions for atomic systems which, discussed Section (.3 is insufficient for wave functions. While this is obvious for regular GNNs, as they operate only on inter- snter-particle distances rather than vectors, equivariant neural networks take advantage of higher SO(3) representa- tions. However, if one would construct the orbitals o( T' [0 ()).......)) ).... ΦN T by concatenating
(2018) or Batzner al. 202T)
E equivariant SO(3) representations o( P) Jon (7) (),-... ..... .............. og r ] with D dim(
N, any resulting real-valued wave function 2 r det oi would be spherically symmetric, ic., و () = 2 ( VR E SO(3).
The proof is follows: If one rotates the electrons the orbital matrix changes as
E I N by any rotation matrix R 8O(3),
Φ (F = (
(26)
D = diag
D
(27)
where D E NNN is block-diagonal matrix and D the Wigner-D matrix induced by rotation R corresponding to the e-th SO(3) representation. restrict our wave function to real-valued
Since Wigner-D matrices unitary and
/ 7 R) = det o( ((PR) det(o(()) RR)
det o( (') det D
E
(28)
det o( (')
II
det D
det d(T))
(7)).
C EDGE CASES OF THE EQUIVARIANT COORDINATE SYSTEM
2
(a) Example a regular poly- gon. For any regular polygon on plane, two cigenvalues of the covariance matrix are go- ing to be identical.
(b) Example of き = 0.
(c) Example why one can not construct unique equivariant coordinate system that changes smoothly with changes in the geometry.
Figure 8: Edge cases in the construction of our equivariant coordinate system. Circles indicate nuclei and the numbers their charges.
While the definition of the coordinate system in Section 3.3 works in most cases. There still exist edge cases where the coordinate system may not be unique. To maximize transparency, we discuss some of these cases, when they occur, how handle them, and what their implications
14

---

Published as a conference paper at ICLR 2022
For certain geometries, two cigenvalues of the nuclei covariance matrix might be identical. If that is the case the PCA axes are not unique. This occurs for any geometry where the nuclei coordinates are distributed regularly regular polygons such as the pentagram depicted in Figure 8a In such cases, we compute the PCA on pseudo coordinates which obtain by stretching the graph in the direction of the largest Coulomb In the example from Figure Sa this is equivalent to stretching the graph along potential R.... 12
sphere around the center of mass. Examples of such geometries are
7m Zmbn
one of the outer edges they are all of the same lengths. The actual direction does not matter, as it is simply a rotation or reflection the whole system. While regular spherical patterns are not the only case where this issue arises they are the most common cause.
Another potential issue arises if V = O, this occurs for any geometry which is point symmetric in the center of mass such the pentagram in Figure Sa In such cases, the signs of the axes do not matter as reflections result in identical geometries. However, there also exist other geometrics for which Equation T2 is O. An example is depicted in Figure 8b But these cases rare and can be resolved by applying a nonlinear function on the distances in Equation T2
The occurrence of these edge cases leads us question: "Why can we not design a unique coordinate system for each geometry?". While we ideally would want an equivariant coordinate system that changes smoothly with changes in the geometry, it is impossible. We show counter example of this in Figure [c where we see a system of 3 nuclei. In their starting configuration, one can uniquely define two axes as indicated by the colored arrows. But, when moving the leftmost nuclei such that it is in line with the other two nuclei one is left with only one uniquely defined direction as there is no way to differentiate between any orthogonal vector and the blue one. By moving the center nuclei to the right, we again can define two axes for this system, though, one axis is flipped compared to the initial state. So, we neither can define a smoothly changing coordinate system nor unique one for every system. However, in practice, we do not need smoothly changing one but only a unique one. While this is already unattainable as shown by the central figure, we also want to stress that any arbitrary orthogonal vector is equivalent due to the symmetries of the system. Considering these aspects, we believe that our coordinate system definition is sufficient in most scenarios.
D EXPERIMENTAL SETUP
Hyperparameters. If not otherwise specified we used the hyperparameters from Table 2 These result in a similarly sized WFModel to FermiNet from Pfau et al. (2020). For cyclobutadiene, we did not train for the full 60000 iterations but stopped the training after 2 weeks.
Numerical Stability To stabilize the optimization, we initialize the last layers the MetaGNN foute and global such that the biases play the dominant role. Furthermore, compute the final output of the WFModel in log-domain and use the log-sum-exp trick.
Learning continuous subsets. To demonstrate PESNet's ability to capture continuous subsets of the potential energy surface, we train PESNet dynamic set of configurations along the energy surface. Specifically, we subdivide the potential energy surface into even-sized bins and place a random walker within each bin. These walkers slightly alter the molecular structure after each step by moving within their bin. This procedure ensures that our model is not restricted to discrete set of configurations. Therefore, after training, our model can be evaluated at any arbitrary configuration within the training domain without retraining. But, we only evaluate PÊSNet on confi configurations where reference calculations hydrogen chain, and the nitrogen molecule. For H and cyclobutadiene, we train on discrete sets of geometries from the literature (Scherbela al. 2021 Kinal & Piecuch 2007).
available. This procedure is done for the hydrogen rectangle, the
Convergence is easy to detect in cases where we optimize for a fixed set of configurations as the energy will slowly converge to an optimal value, specifically, the average of the optimal energies. However, in cases where we optimize for a continuous set of configurations, the optimal energy value depends on the current batch of geometries. To still access convergence, we use the fact that the local energy F, of any eigenfunction (including the ground-state) has 0 variance. So, our convergence
27
criteria is the expected variance of the local energy EE ER
E E7
EE - E (Ezl)
where the optimal value is O.
15

---

Published as a conference paper at ICLR 2022
Parameter
Value
Optimization
Local energy clipping
5.0
Optimization
Batch size
4096
Optimization
Iterations
60000
Optimization
#Geometry random walker
16
Optimization
Learning rate 7)
0.1
(1+0/1000)
Natural Gradient
Damping
10-4. Std Fr.]
Natural Gradient
CG max steps
100
WFModel
Nuclei embedding dim
64
WFModel
Single-stream width
256
WFModel
Double-stream width
32
WFModel
#Update layers
WFModel
#Determinants
16
MetaGNN
#Message passings
MetaGNN
Embedding dim
64
MetaGNN
Message dim
32
MetaGNN
NsuF
7
MetaGNN
NRBF
MetaGNN
MLP depth
MCMC
Proposal step size
0.02
MCMC
Steps between updates
40
Pretraining
Iterations
2000
Pretraining
Learning rate
0.003
Pretraining
Method
UHF
Pretraining
Basis
STO-6G
Evaluation
#Samples
106
Evaluation
MCMC Steps
200
Table 2: Default hyperparameters.
H
H4
H10
N2 Cyclobutadiene
no MetaGNN
1.849286(9)
-2.016199(5)
-53328944(14)
-109.28322(9)
(4444111))11
MetaGNN
-1.849333()
-2.016208(5)
-5.22891((55)
-109.28570(7)
154.65469(27)
AE
0.000077(11) 0.000009(7)
-..00002((1)) 0.00248(11)
0.0105(4)
Table Energies in En averaged over the PES for PESNets with and without the MetaGNN. Num- bers in brackets indicate the standard eтror at the last digit(s). In cases without the MetaGNN, we still train single model for configurations of system.
E ABLATION STUDIES
As hyperparameters often play significant in the machine learning community, present some ablation studies in this appendix. All the following experiments only alter one variable a time, while the rest is fixed in Table The results in the tables are averaged over the same configura- tions as in the main body.
Table 3 presents results with and without the MetaGNN. It is noticeable that the gain of the MetaGNN is little to none for small molecules consisting of simple hydrogen atoms. But, for more
#Dets
H
H4
H10
N2 Cyclobutadiene
(.849363(6)
-2.016208(5)
-5.328916(15)
-109228570(7)
154.65469(27)
32
11849342(4)
-2.016188(6)
-5.328999(13)
-109.28706(7)
-154.65322(27)
AE
-0.00002(7)) -0.000019(8)
0.000083(20)
0.00136(10)
-0.0015(4)
Table 4: Energies in Eg averaged over the PES for different number of determinants in our PESNet model. Numbers in brackets indicate the standard error the last digit(s).
16

---

Published as a conference paper at ICLR 2022
dim (h.)
H
H4
H10
N2
Cyclobutadiene
256
-1 (6)
-2.016208(5)
-5.328916(15)
109.28570(7)
-154.654699 (54.6556)())
512
-1.849354((28)
-2.016190(7)
-5.328794(17)
-19.28662(6)
(.54.65042(28))
AE
-0.000009(7) -0.000017(8) -..00012((3))
0.00092(9)
-0.0043(4)
Table 5: Energies in Es averaged over the PES for different single-stream sizes in our PESNet model. Numbers in brackets indicate the standard error the last digit(s).
-140.00
150.00
Btoegge Bq.
PESNet (26/256)
154.00
PRS (32/256)
Net (16/512)
-154.00
-144.67
PESNet (16/256)
AE leat
NeS (32/256)
PESNet (36/512)
CCSD(T)
0-
MR-OC
experizent
10
Iterations
Figure 9: Comparison of different PESNet configurations on cyclobutadiene. The configurations are named (#dcterminann/inggl--ttte width) with light colors for the ground state and darker colors for the transition state.
complex molecules such as nitrogen and cyclobutadiene. we notice significant improvements of 2.5 F2 and 10.5 m Fr, respectively. Moreover, the MetaGNN enables us to account for symme- tries of the energy while the WFModel itself is only invariant w.r.t. to translation but not to rotation, reflection, and reindexing nuclei.
Table 4 shows the impact of the number of determinants on the average energy for the systems from Section For small hydrogen-based systems, the number of determinants is mostly irrelevant, while larger numbers of determinants improve performance for nitrogen. But, this does not seem to carry over to cyclobutadiene. While the total estimated energy is higher for the larger model, we noticed that it is significantly faster converging the transition barrier. Figure illustrates this by comparing the convergence of the 16 and 32 determinant models.
Increasing the single-stream size does not seem result in any benefit for most models Table shows. Again, the hydrogen systems are mostly unaffected by this hyperparameter while nitrogen benefits from larger hidden dimensions but cyclobutadiene converges worse. We suspect that this is due to the optimization problem becoming significantly harder. Firstly, increasing the WFModel increases the number of parameters the MetaGNN has to predict. Secondly, we estimate the inverse of the Fisher with finite fixed-sized batch but the Fisher grows quadratically with the number of parameters which also grow quadratically with the single size stream.
F TIME PER ITERATION
While the main document already covers the time it took reproduce the results from the figures, we want to use this appendix to provide more details on this. Table the time per iteration for single model instead of the whole training time to reproduce the potential energy surfaces as in Ta- ble [] While we find these numbers to be misleading we still want to disclose them to support open research. The main issue with these numbers is that they do not take the quality of the update into account, e.g., PauliNet and DeepErwin are trained with Adam (Kingma & Ba 2014), FermiNet with K-FAC (Martens & Grosse 2015). and PESNet with CG-computed natural gradient (Neuscamman et al. 2012). This has implications
the number of iterations has to train. For Table in the main body. we assumed that PauliNet is trained for 10000 iterations (Hermann et al. 2020 Deep-
Erwin for 7000 iterations (Scherbela et al. 202T). FermiNet for 200000 iterations (Pfau al. 2020
17

---

Published as a conference paper at ICLR 2022
H
H4
H10
N2
Cyclobutadiene
PauliNet
0,83s
1.13s
5.51s
8.09s
175s
DrepErwin
0.92s
1.28s
4.88s
FermiNet
0.12s
0.19s
1.07s
1.99s
20.8s
PESNet
1.19s
1.42s
3.91s
5.32s
33.2s
Table 6: Time per training step.
Chemical accuracy
Elaxco in Eh
10-2
H4
H4plus
10-2
H10
N2
10"
E
10-6
1000
2000
4000
10000
32000
Stops
Figure 10: Convergence behavior of PESNet. Error bars indicate the standard error the mean.
and PESNet for 60000 iterations. Though, one can assume that neither PauliNet nor DeepErwin are going to produce similar results to FermiNet and PESNet on the more complex nitrogen molecule or cyclobutadiene in 10000 or 7000 iterations. This further discussed in Appendix G When viewing these results, one also has to keep in mind the different quality of the results, e.g.. FermiNet and PESNet strictly outperform PauliNet and DeepErwin on all systems. Another potential issue arises due to the choice of deep learning framework the models are implemented in. It has been shown that JAX works very well for computing the kinetic energy (Spencer et al. 2020) which usually is the largest workload when computing update.
G CONVERGENCE
When choosing method to solve the Schrödinger equation has to pick between accuracy and speed. For classical methods, this might be the difference between choosing DFT, CCSD(T), or FCI. In the neural VMC setting, one way reflect this by choosing the number of training steps. To better investigate this, we present convergence graphs for H4, HJ H10, and N2 in Figure 10 For cyclobutadiene,
can the convergence of different configurations of PESNet in Figure One can see that our network converges quite quickly on hydrogen-based systems such as the hydrogen rectangle, H or the hydrogen chain. For these small systems, PESNet surpasses PauliNet's and DeepErwin's 's accuracy in less than 8000 training steps, reducing the training times to 2.9h, 3.2h, and 8.7h, respectively. In less than 16000 steps, about 17.4h training, PESNet surpasses FermiNet on the hydrogen chain. For more complicated molecules such N2 and cyclobutadiene, the training requires more iterations to converge. So, we may also expect the extrapolated numbers from Table for PauliNet to be optimistic estimate.
18


Reviews:
[{"review_id": "ckqLR4Ai82b", "reviewer": "Reviewer_3XzF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, the manuscript is well written and theory and experiments are presented in a concise fashion. The idea of coupling neural wavefunctions for different nuclear configurations with a meta-GNN is promising, since this allows for simultaneous solution with VMC. The experiments demonstrate the utility of such an approach in terms of accuracy and computational speed up.  There are, however, some issues that should be addressed. \n\nFirst, the comparison of computational timings between the different methods is hard to follow at points. Is it correct to assume that the reported timings are those needed to solve for every configuration of a particular molecule? How is this handled in case of the systems where PESNet is solved for a continuous subset of structures (e.g. are there global convergence criteria)? Regarding the time per iteration reported in the SI, it is unclear how exactly these are computed. For example, in Table 1, it is assumed that 10000 iterations with PauliNet for H4+ take 43 hours, which should give a time per iteration of ~15 s instead of the 0.83 s reported. Similarly, solving PESNet for the same system takes only half the time in Table 1, while the SI reports 60000 iterations @ 1.29 s. Finally, there is  the question whether the pretraining on HF orbitals is included in the timings.\n\nSecond, the comparison between different approaches could profit from additional information. While not a completely reliable measure, it would be helpful to report the parameter count of the different models as used in the experiments in order to gain an impression of their overall flexibility. Moreover, an ablation study, where the MetaGNN and WFModel are optimized for the different configurations of a system individually would add to the analysis of the model. This would help to identify how the MetaGNN influences the wavefunction solution of a single conformer and if e.g. the structural dependence of the weights introduced in the MetaGNN is beneficial.\n\nFinally, regarding the the equivariant coordinate frame. The PCA based frame can potentially become unstable if structures are strongly distorted (e.g. principal axis change order). Perhaps the authors can comment on how this is handled and whether it presents a problem for the wavefunction model or limits the range of structural configurations that can be studied with a single metaGNN. Is the same coordinate frame used for all configurations of a system or are individual frames employed? Finally, it would be interesting to motivate why this particular approach of introducing SE(3) equivariance was chosen in favor of more conventional SE(3) equivariant graph networks.", "summary_of_the_paper": "The authors introduce an approach based on using graph neural networks to parametrize wavefunctions represented by an equivariant neural network. This allows them simultaneously solve for the ground state energy of multiple molecular configurations with variational Monte Carlo.", "summary_of_the_review": "The manuscript is well written and the proposed approach is sufficiently novel. Its advantages are further clearly demonstrated in a series of experiments. As such, I recommend to accept this contribution.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IUAR7nMBCfn", "reviewer": "Reviewer_bjEg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**:\n* Comprehensive set of tests and comparison to other relevant work\n* Study of relevant hyperparameters and convergence analysis\n* Variational energies are respectable across all studied problems\n\n**Weaknesses**:\n* A major improvement of the method (compared to existing approaches like FermiNet) is that the model can be optimized on multiple similar configurations simultaneously to achieve faster net runtime. While it is an interesting observation that GNN is capable of conditioning the model parameters on the geometry to adjust the modeled ground state wave function, other methods could use optimization results from one configuration as a starting point for others. This generally speeds up the number of iterations needed to converge to slightly different solutions. Such an estimate would be a more fair approach to runtime comparison. I believe that even if the runtime advantage is noticeably reduced it is still an interesting result.\n* The other stress test to justify that the GNN model has learned something non-trivial would be to include a baseline method where the parameters of the wave function are not changing with the system. (e.g. show the results of PESNet for Fig. 3 and Fig. 4 with single parameters w; w_m predicted by GNN intermediate regime). This would enable a clear assessment of how relevant the modifications to the ground state are.\n", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "The paper presents sufficiently strong results and an interesting demonstration of how a single network can be used to simultaneously approximate a family of ground state wave functions. Results are compared to relevant baselines and with simple additional checks suggested in the main review should be a good contribution to variational methods in quantum chemistry.\n\n**Update**\nAfter the discussion period many of my concerns have been cleared. Changing my score from 6 to 8. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3dD1z7pmu2A", "reviewer": "Reviewer_Zkwk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n1. The proposed PESNet obtains SOTA or near-SOTA results compared to previous methods while being much faster to train. \n2. Using meta learning, the PESNet is able to handle multiple geometries simultaneously.\n3. The use of meta learning in this manner is novel and interesting.\n\nWeakness\n1. Recently, a number of datasets based on DFT have been created (e.g. OC20) and various GNN models have been proposed to train on such datasets (e.g. Dimenet++, Gemnet, etc). This paper, in contrast, presents an ab initio method that does not require a training dataset. While the introduction section discusses the difference in these approaches, it is unclear how the PESNet compares to the the trained methods when they are trained on very large and diverse datasets like OC20. It would be  helpful to add some discussion on this.\n2. The paper uses a lot of terminology that is difficult to follow for somebody without a background in chemistry. Since ICLR is a general DL conference, I would urge the authors to change the text to make it easier to follow.\n", "summary_of_the_paper": "The paper presents a new meta-learning method for solving the Schrodingers equation using neural networks. This method combines an existing neural wave function model called FermiNet together with a GNN (MetaGNN) to solve the Schrodingers equation for multiple geometries simultaneously. The MetaGNN takes the atomic graph as input and outputs a set of parameters that capture the 3D geometry of the system, which are then input to the FermiNet. The resulting method can is applicable to multiple geometries while being significantly faster to train.\n", "summary_of_the_review": "This paper presents a novel ab initio method for solving the Schrodingers equation using meta-learning by combining existing neural wave function models with GNNs. The resulting method can solve the Schrodingers equation for multiple geometries while training significantly faster.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper builds on the success of the FermiNet neural wave function framework by pairing it with a graph neural network which predicts the parameters of neural wave function from the geometry. The resulting PESNet trains significantly faster, with no loss of accuracy. This method constitutes an important advance in ML-powered quantum mechanical calculations.

The reviewers unanimously recommend acceptance.
Decision: Accept (Spotlight)