Year: 2021

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

text: 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.

heading: 2 RELATED WORK
text: 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).

heading: 3 METHOD
text: 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.

heading: 3.1 WAVE FUNCTION MODEL
text: 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)

heading: 3.2 METAGNN
text: The MetaGNN’s task is to adapt the WFModel to the geometry at hand. It does so by substituting subsets, ω and ωm, of WFModel’s parameters. While ωm contains parameters specific to nuclei m, ω is a 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 coordinate system when initializing the node embeddings to avoid identical embeddings in symmetric structures. Hence, our node embeddings are initialized by
l1m = [ GZm , fpos( −→ R ′ mE) ] (7)
where G is a matrix of charge embeddings, Zm ∈ N+ is the charge of nucleus m, fpos : R3 7→ RNSBF·NRBF is our positional encoding function, and −→ R ′ mE 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-Bessel basis functions ãSBF,ln from Klicpera et al. (2019)
fpos( −→x ) = 3∑ i=1 [ ãSBF,ln(∥−→x ∥,∠(−→x ,−→e i)) ] l∈{0,...,NSBF−1},n∈{1,...,NRBF} (8)
with −→e i being the ith axis of our equivariant coordinate system E. Unlike Klicpera et al. (2019), we are working on the fully connected graph and, thus, neither include a cutoff nor the envelope function that decays to 0 at the cutoff.
A message passing layer consists of a message function fmsg and an update function fupdate. Together, one can compute an update to the embeddings as
lt+1m = f l update ( ltm, ∑ n f tmsg(l t m, l t n, emn) ) (9)
where emn is an embedding of the edge between nucleus m and nucleus n. We use Bessel radial basis functions to encode the distances between nuclei (Klicpera et al., 2019). Both fmsg and fupdate are realized by simple feed-forward neural networks with residual connections.
After LGNN many message passing steps, we compute WFModel’s parameters on two levels. On the global level, f outglobal outputs the biases of the network and, on the node level, f out node outputs nuclei specific parameters:
ω = [ b1single/double, . . . , b ↑/↓ 1 , . . . ,w ] := f outglobal [∑ m ltm ]LGNN t=1  , ωm = [ zm, s 1,↑/↓ m , . . . ,p 1,↑/↓ m , . . . ] := f outnode ([ ltm ]LGNN t=1 ) .
(10)
We use distinct feed-forward neural networks with multiple heads for the specific types of parameters estimated to implement f outnode and f out global.

heading: 3.3 EQUIVARIANT COORDINATE SYSTEMS
text: Incorporating symmetries helps to reduce the training space significantly. In GNNs this is done by only operating on inter-nuclei distances without a clear directionality in space, i.e., without x, y, z coordinates. While this works for predicting observable metrics such as 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, as we show in Appendix B, recently proposed equivariant GNNs (Thomas et al., 2018; Batzner et al., 2021) also suffer from the same limitation.
To solve this issue, we introduce directionality in the form of a coordinate system that is equivariant to rotations and reflections. The axes of our coordinate system E = [−→e 1,−→e 2,−→e 3] are defined by the principal components of the nuclei coordinates, −→e PCA1 ,−→e PCA 2 ,
−→e PCA3 . 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 to a sign, we have to resolve the sign ambiguity. We do this by computing an equivariant vector −→v , i.e, a vector that rotates and reflects with the system, and defining the direction of the axes as
−→e i = {−→e PCAi , if −→v T−→e PCAi ≥ 0, −−→e PCAi , else.
(11)
As equivariant vector we use the difference between a weighted and the regular center of mass
−→v := 1 M M∑ m=1 ( M∑ n=1 ∥−→Rm − −→ Rn∥2 ) Zm −→ R ′ m, (12)
−→ R ′ m = −→ Rm − 1
M M∑ m=1 −→ Rm. (13)
With this construction, we obtain an equivariant coordinate system that defines directionality in space. However, we are aware that PCA may not be robust, e.g., if eigenvalues of the covariance matrix are identical. A detailed discussion on such edge cases can be found in Appendix C.

heading: 3.4 OPTIMIZATION
text: We use the standard VMC optimization procedure (Ceperley et al., 1977) where we seek to minimize the expected energy of a wave function ψθ parametrized by θ:
L = ⟨ψθ|H|ψθ⟩⟨ψθ|ψθ⟩ (14)
where H is the Hamiltonian of the Schrödinger equation
H = −1 2 N∑ i=1 ∇2i + N∑ i=1 N∑ j=i
1
∥−→r i −−→r j∥ − N∑ i=1 M∑ m=1
1
∥−→r i − −→ Rm∥ + M∑ m=1 M∑ n=m
ZmZn
∥−→Rm − −→ Rn∥︸ ︷︷ ︸
V (−→r ) (15)
with ∇2 being the Laplacian operator and V (−→r ) describing the potential energy. Given samples from the probability distribution ∼ ψ2θ(−→r ), one can obtain an unbiased estimate of the gradient
EL( −→r ) = ψ−1θ (−→r )Hψθ(−→r )
= −1 2 N∑ i=1 3∑ k=1 [ ∂2 log |ψθ| ∂−→r 2ik + ∂ log |ψθ| ∂−→r ik 2 ] + V (−→r ),
(16)
∇θL = E[(EL − E[EL])∇θ log |ψθ|] (17) with EL denoting the local energy. One can see that for the energy computation, we only need the derivative of the wave function w.r.t. the electron coordinates. As these are no inputs to the MetaGNN, we do not have to differentiate through the MetaGNN to obtain the local energies. We clip the local energy as in Pfau et al. (2020) and obtain samples from ∼ ψ2θ(−→r ) via MetropolisHastings. The gradients for the MetaGNN are computed jointly with those of the WFModel by altering Equation 17:
∇ΘL = E[(EL − E[EL])∇Θ log |ψΘ|] (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 gradient of the WFModel and the MetaGNN enables us to use a single training pass to simultaneously solve multiple Schrödinger equations.
While Equation 18 provides us with a 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, we use natural gradient descent to enable the use of larger learning rates. So, instead of doing a regular gradient descent step in the form of Θt+1 = Θt − η∇ΘL, where η is the learning rate, we add the inverse of the Fisher information matrix as a preconditioner
Θt+1 = Θt − ηF−1∇ΘL, (19) F = E[∇Θ log |ψΘ|∇Θ log |ψΘ|T ]. (20)
Since the Fisher F scales quadratically with the numbers of parameters, we approximate F−1∇ΘL via the conjugate gradient (CG) method (Neuscamman 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 η, we clip the norm of the preconditioned gradient F−1∇ΘL (Pascanu et al., 2013) and use a fixed learning rate for all systems.
We pretrain all networks with the Lamb optimizer (You et al., 2020) on Hartree-Fock orbitals, i.e., we match each of the K orbitals to a Hartree-Fock orbital of a different configuration. During pretraining, only the WFModel and the final biases of the MetaGNN are optimized.

heading: 3.5 LIMITATIONS
text: While PESNet is capable of accurately modeling complex potential energy surfaces, we have not focused on architecture search yet. For instance, both the WFModel and the MetaGNN are only as powerful as the Weisfeiler-Lehman graph isomorphism test (Xu et al., 2018). Furthermore, as we discuss in Section 4, PauliNet still offers a better initialization and converges in fewer iterations than our network. Lastly, PESNet is limited to geometries with identical electron spin configurations and, thus, the same number of electrons. We plan on addressing these issues in future work.

heading: 4 EXPERIMENTS
text: To investigate PESNet’s accuracy and training time benefit, we compare it to FermiNet (Pfau et al., 2020; Spencer et al., 2020), PauliNet (Hermann et al., 2020), and DeepErwin (Scherbela et 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 on energies one should be aware that, except for PESNet, all methods must be trained separately for each configuration resulting in significantly higher training times, as discussed in Section 4.1.
Evaluation of ab-initio methods is still a challenge as true energies are rarely known, and experimental 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 1 kcalmol−1 ≈ 1.6mEh. 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, interpretation is simpler: lower is always better as VMC energies are upper bounds (Szabo & Ostlund, 2012).
To analyze PESNet’s ability to capture continuous subsets of the potential energy surface, we train on the continuous energy surface rather than on a discrete set of configurations for potential energy surface scans. The exact procedure and the general experimental setup are described in Appendix D. Additional ablation studies are available in Appendix E.
Transition path of H+4 and weight sharing. Scherbela et al. (2021) use the first transition path of H+4 (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 DeepErwin, our novel PESNet is flexible enough such that we do not need any extra optimization. In Figure 3, we see the DeepErwin results after their multi-step optimization and the energies of a single PESNet. We notice that while both methods estimate similar transition barriers PESNet results in ≈ 0.27mEh smaller energies which agree with the very accurate MRCI-D-F12 results (≈ 0.015mEh). Hydrogen rectangle and symmetries. The Hydrogen rectangle is a known failure case for CCSD and CCSD(T). While the exact solution, FCI, indicates a local maximum at θ = 90◦, both, CCSD 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 (≈ 0.014mEh) despite training only a single network on half of the configuration space. Thanks to our equivariant coordinate system, we only have to train on the first half of the potential energy surface.
Hydrogen chain. The hydrogen chain is a very common benchmark geometry that allows us to compare our method to a range of classical methods (Motta et al., 2017) as well as to FermiNet, PauliNet, and DeepErwin. Figure 5 shows the potential energy surface of the hydrogen chain computed by a range of methods. While our PESNet generally performs identical to FermiNet, we predict on average 0.31mEh lower energies. Further, we notice that our results are consistently better than PauliNet and DeepErwin despite only training a single model.
The nitrogen molecule. The nitrogen molecule poses a challenge as 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, it scales factorially (Gdanitz, 1998). After Pfau et al. (2020) have shown that FermiNet is capable of modeling such complex triple bonds, we are interested in PESNet’s performance. To better represent both methods, we decided to compare both FermiNet as well as PESNet with 32 determinants due to a substantial performance gain for both methods. The results in Figure 6 show that PESNet agrees very well with FermiNet and is on average just 0.37mEh higher, despite training only a single model for less than 140 of FermiNet’s training time, see Section 4.1.
Cyclobutadiene and the MetaGNN. The automerization of cyclobutadiene is challenging due to its multi-reference nature, i.e., single reference methods such as CCSD(T) overestimate the transition 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), we increased the number of determinants for the experiment to 32 and the single-stream size to 512.
As shown in Figure 7, all neural methods converge to the same transition barrier which aligns with the highest MR-CC results at 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 70mEh smaller. Lastly, PESNet’s transition barrier converges similar PauliNet’s, but its energies are 54mEh lower than PauliNet’s, placing it closer to FermiNet than PauliNet in terms of accuracy. Considering that PESNet has only been trained for 13 of FermiNet’s time (see Section 4.1), we are confident that additional optimization would further narrow the gap to FermiNet.
In an additional ablation study, we compare to PESNet without the MetaGNN. Note, that we still train a single model for both states of cyclobutadiene. While the results in Figure 7 show that the truncated network continuously lowers its energies, it fails to reproduce the same transition barrier and its energies are 18mEh worse compared to the full PESNet.

heading: 4.1 TRAINING TIME
text: While the previous experiments have shown that our model’s accuracy is 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 implementation of FermiNet (Spencer et al., 2020), the official PyTorch implementation of PauliNet (Hermann et al., 2020), the official TensorFlow implementation of DeepErwin (Scherbela et al., 2021), and our JAX (Bradbury et al., 2018) implementation of PESNet. We use the same hyperparameters as in the
experiments or the defaults from the respective works. All measurements have been conducted on a machine with 16 AMD EPYC 7543 cores and a single Nvidia A100 GPU. Here, we only measure the VMC training time and explicitly exclude the time to perform the SCFs calculations or 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.
Table 1 shows the GPU hours to train the models of the last section. It 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 a continuous subset of the energy surface. Thus, if 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 F and G for additional discussion on training and convergence.

heading: 5 DISCUSSION
text: We presented a novel architecture that can simultaneously solve the Schrödinger equation for multiple geometries. Compared to the existing state-of-the-art networks, our PESNet accelerates the training for many configurations by up to 40 times while often achieving better accuracy. The integration of physical symmetries enables us to reduce our training space. Finally, our results show that a single model can capture a continuous subset of the potential energy surface. This acceleration of neural wave functions opens access to accurate quantum mechanical calculations to a broader audience. For instance, it may enable significantly higher-resolution analyses of complex potential energy surfaces with potential applications in generating new datasets with unprecedented accuracy or even integration into molecular dynamics simulations.
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 of misuse, e.g., for the development of chemical weapons. To the best of our knowledge, our work does not promote misuse any more than general computational chemistry research. To reduce the likelihood of such misuse, we publish our source code under the Hippocratic license (Ehmke, 2019). To facilitate reproducibility, the source code includes simple scripts to reproduce all experiments from Section 4. Furthermore, we provide a detailed schematic of the computational graph in Figure 2 and additional details on the experimental setup including all hyperparameters in Appendix D.

heading: B EQUIVARIANT NEURAL NETWORKS AS WAVE FUNCTIONS
text: Here, we want to briefly discuss why equivariant neural networks as proposed by Thomas et al. (2018) or Batzner et al. (2021) are no alternative to our equivariant coordinate system. The issue is the same as for regular GNNs (Klicpera et al., 2019), namely that such networks can only represent spherically symmetric functions for atomic systems which, as discussed in Section 3.3, is insufficient for wave functions. While this is obvious for regular GNNs, as they operate only on inter-particle distances rather than vectors, equivariant neural networks take advantage of higher SO(3) representations. However, if one would construct the orbitals ϕ(−→r ) = [ϕ1(−→r ), . . . , ϕN (−→r )] by concatenating E equivariant SO(3) representations ϕ(−→r ) = [ϕ1(−→r ), . . . , ϕE(−→r )] with ∑E e=1 dim(ϕe(
−→r i)) = N , any resulting real-valued wave function ψ(−→r ) = detϕ(−→r ) would be spherically symmetric, i.e., ψ(−→r R) = ψ(−→r ),∀R ∈ SO(3). The proof is as follows: If one rotates the electrons −→r ∈ RN×3 by any rotation matrix R ∈ SO(3), the orbital matrix changes as
ϕ(−→r R) = ϕ(−→r )DR, (26) DR = diag(DR1 , . . . ,D R E) (27)
where DR ∈ RN×N is a block-diagonal matrix and DRe is the Wigner-D matrix induced by rotation R corresponding to the e-th SO(3) representation. Since Wigner-D matrices are unitary and we restrict our wave function to real-valued
ψ(−→r R) = detϕ(−→r R) = det(ϕ(−→r )DR) = detϕ(−→r ) detDR
= detϕ(−→r ) E∏ e=1 detDRe = detϕ(−→r ) = ψ(−→r ).
(28)
□
C EDGE CASES OF THE EQUIVARIANT COORDINATE SYSTEM
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 we handle them, and what their implications are.
For certain geometries, two eigenvalues 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 on a sphere around the center of mass. Examples of such geometries are regular polygons such as the pentagram depicted in Figure 8a. In such cases, we compute the PCA on pseudo coordinates which we obtain by stretching the graph in the direction of the shortest edge−→ Rm− −→ Rn. In the example from Figure 8a, this is equivalent to stretching the graph along one of the outer edges as they are all of the same lengths. The actual direction does not matter, as it is simply a rotation or reflection of 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 = 0, this occurs for any geometry which is point symmetric in the center of mass such as the pentagram in Figure 8a. In such cases, the signs of the axes do not matter as reflections result in identical geometries. However, there also exist other geometries for which Equation 12 is 0. An example is depicted in Figure 8b. But these cases are rare and can be resolved by applying a nonlinear function on the distances in Equation 12.
The occurrence of these edge cases leads us to 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 a counter example of this in Figure 8c 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 to 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 a unique one for every system. However, in practice, we do not need a smoothly changing one but only a unique one. While this is already unobtainable 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.

heading: D EXPERIMENTAL SETUP
text: 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 of the MetaGNN f outnode and f outglobal such that the biases play the dominant role. Furthermore, we 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 on a 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 a 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 PESNet on configurations where reference calculations are available. This procedure is done for the hydrogen rectangle, the hydrogen chain, and the nitrogen molecule. For H+4 and cyclobutadiene, we train on discrete sets of geometries from the literature (Scherbela et al., 2021; Kinal & Piecuch, 2007).
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 EL of any eigenfunction (including the ground-state) has 0 variance. So, our convergence
criteria is the expected variance of the local energy E−→ R∼PES [ E−→r ∼ψ2−→
R
[( EL − E−→r ∼ψ2−→
R
[EL] )2]]
where the optimal value is 0.

heading: E ABLATION STUDIES
text: As hyperparameters often play a significant in the machine learning community, we present some ablation studies in this appendix. All the following experiments only alter one variable at a time, while the rest is as in Table 2. The results in the tables are averaged over the same configurations 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
complex molecules such as nitrogen and cyclobutadiene, we notice significant improvements of 2.5mEh and 10.5mEh, respectively. Moreover, the MetaGNN enables us to account for symmetries of the energy while the network without is only invariant w.r.t. to translation but not to rotation, reflection, and reindexing of nuclei.
Table 4 shows the impact of the number of determinants on the average energy for the systems from Section 4. 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 at converging the transition barrier. Figure 9 illustrates this by comparing the convergence of the 16 and 32 determinant models.
Increasing the single-stream size does not seem to result in any benefit for most models as Table 5 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 a finite fixed-sized batch but the Fisher grows quadratically with the number of parameters which also grow quadratically with the single size stream.

heading: F TIME PER ITERATION
text: While the main document already covers the time it took to reproduce the results from the figures, we want to use this appendix to provide more details on this. Table 6 lists the time per iteration for a single model instead of the whole training time to reproduce the potential energy surfaces as in Table 1. 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 on the number of iterations one has to train. For Table 1 in the main body, we assumed that PauliNet is trained for 10000 iterations (Hermann et al., 2020), DeepErwin for 7000 iterations (Scherbela et al., 2021), FermiNet for 200000 iterations (Pfau et al., 2020),
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 is 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 an update.

heading: G CONVERGENCE
text: When choosing a method to solve the Schrödinger equation one 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 to reflect this is by choosing the number of training steps. To better investigate this, we present convergence graphs for H4, H+4 , H10, and N2 in Figure 10. For cyclobutadiene, we can see the convergence of different configurations of PESNet in Figure 9. One can see that our network converges quite quickly on hydrogen-based systems such as the hydrogen rectangle, H+4 , or the hydrogen chain. For these small systems, PESNet surpasses PauliNet and DeepErwin 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 of training, PESNet surpasses FermiNet on the hydrogen chain. For more complicated molecules such as N2 and cyclobutadiene, the training requires more iterations to converge. So, we may also expect the extrapolated numbers from Table 1 for PauliNet to be an optimistic estimate.

Review summaries:
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.

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.