Original: A B -I NITIO P OTENTIAL E NERGY S URFACES BY P AIRING GNN S WITH N EURAL W AVE F UNCTIONS
Revised: A B -I NITIO P OTENTIAL E NERGY S URFACES BY P AIRING GNN S WITH N EURAL W AVE F UNCTIONS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Solving the Schr¨odinger equation is key to many quantum mechanical properties.
Revised: Solving the Schr¨odinger equation is key to many quantum mechanical properties.

Original: However, an analytical solution is only tractable for single-electron systems.
Revised: However, an analytical solution is only tractable for single-electron systems.

Original: Recently, neural networks succeeded at modeling wave functions of many-electron systems.
Revised: Recently, neural networks succeeded at modeling wave functions of many-electron systems.

Original: Together with the variational Monte-Carlo (VMC) framework, this led to solutions on par with the best known classical methods.
Revised: Together with the variational Monte-Carlo (VMC) framework, this led to solutions on par with the best known classical methods.

Original: Still, these neural methods require tremendous amounts of computational resources as one has to train a separate model for each molecular geometry.
Revised: Still, these neural methods require tremendous amounts of computational resources as one has to train a separate model for each molecular geometry.

Original: In this work, we combine a Graph Neural Network (GNN) with a neural wave function to simultaneously solve the Schr ¨ odinger equation for multiple geometries via VMC.
Revised: In this work, we combine a Graph Neural Network (GNN) with a neural wave function to simultaneously solve the Schr ¨ odinger equation for multiple geometries via VMC.

Original: This enables us to model continuous subsets of the potential energy surface with a single training pass.
Revised: This enables us to model continuous subsets of the potential energy surface with a single training pass.

Original: Compared to existing state-of-the-art networks, our Potential Energy Surface Network (PESNet) speeds up training for multiple geometries by up totimes while matching or surpassing their accuracy.
Revised: Compared to existing state-of-the-art networks, our Potential Energy Surface Network (PESNet) speeds up training for multiple geometries by up totimes while matching or surpassing their accuracy.

Original: This may open the path to accurate and orders of magnitude cheaper quantum mechanical calculations.
Revised: This may open the path to accurate and orders of magnitude cheaper quantum mechanical calculations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: [Figure]
Revised: [Figure]

Original: Given these, the WFModel evaluates the electronic wave function ψ ( ⃗r ) .
Revised: Given these evaluates the electronic wave function ψ ( ⃗

Original: 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
Revised: 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

Original: Lilienfeld, 2021).
Revised: Lilienfeld, 2021).

Original: In particular, a lot of work has gone into building surrogate models to reproduce
Revised: In particular, a lot of work has gone into building surrogate models to reproduce

Original: QM properties, e.g., energies.
Revised: QM properties, e.g., energies.

Original: 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)
Revised: 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)

Original: (Chmiela et al., 2018).
Revised: (Chmiela et al., 2018).

Original: While this approach has shown great success in recovering the baseline calculations, it suffers from several disadvantages.
Revised: While this approach has shown great success in recovering the baseline calculations, it suffers from several disadvantages.

Original: 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 property.
Revised: 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 property.

Original: 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).
Revised: 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).

Original: In orthogonal research, neural networks have been used as wave function Ans¨atze to solve the stationary Schr ¨ odinger equation (Kessler et al., 2021; Han et al., 2019).
Revised: In orthogonal research, neural networks have been used as wave function Ans¨atze to solve the stationary Schr¨odinger equation (Kessler et al., 2021; Han et al., 2019).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The data (electron positions) are sampled from the wave function itself, and the loss is derived from the Schr¨odinger equation (Ceperley et al., 1977).
Revised: The data (electron positions) are sampled from the wave function itself, and the loss is derived from the Schr¨odinger equation (Ceperley et al., 1977).

Original: 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).
Revised: This approach has shown great success as multiple authors report results outperforming the traditional ‘gold-standard’ CCSD on various systems (Pfau et al., 2020; Hermann et al., 2020).

Original: However, these techniques require expensive training for each geometry, resulting in high computational requirements and, thus, limiting their application to small sets of configurations.
Revised: However, these techniques require expensive training for each geometry, resulting in high computational requirements and, thus, limiting their application to small sets of configurations.

Original: In this work, we accelerate VMC with neural wave functions by proposing an architecture that solves the Schr¨odinger equation for multiple systems simultaneously.
Revised: In this work, we accelerate VMC with neural wave functions by proposing an architecture that solves the Schr¨odinger equation for multiple systems simultaneously.

Original: 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 Schr¨odinger equation for a specific geometry.
Revised: 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 Schr¨odinger equation for a specific geometry.

Original: Previously, these parameters were obtained by optimizing a separate wave function for each geometry.
Revised: Previously, these parameters were obtained by optimizing a separate wave function for each geometry.

Original: We improve this procedure by generating the parameters with a GNN, as illustrated in Figure 1.
Revised: We improve this procedure by generating the parameters with a GNN, as illustrated in Figure 1.

Original: This enables us to capture continuous subsets of the potential energy surface in one training pass, removing the need for costly retraining.
Revised: This enables us to capture continuous subsets of the potential energy surface in one training pass, removing the need for costly retraining.

Original: 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¨utt et al., 2018).
Revised: 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¨utt et al., 2018).

Original: While these symmetries hold for observable metrics such as energies, the wave function itself may not have these symmetries.
Revised: While these symmetries hold for observable metrics such as energies, the wave function itself may not have these symmetries.

Original: We solve this issue by defining a coordinate system that is equivariant to the symmetries of the energy.
Revised: We solve this issue by defining a coordinate system that is equivariant to the symmetries of the energy.

Original: 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.
Revised: 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.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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.
Revised: 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.

Original: Classically, features were constructed by hand and fed into a machine learning model to predict target properties (Christensen et al., 2020;
Revised: Classically, features were constructed by hand and fed into a machine learning model to predict target properties (Christensen et al., 2020;

Original: Behler, 2011; Bart ´ ok et al., 2013).
Revised: Behler, 2011; Bart ´ ok et al., 2013).

Original: Recently, GNNs have proven to be more accurate and took over the field (Yang et al., 2019; Klicpera et al., 2019;
Revised: Recently, GNNs have proven to be more accurate and took over the field (Yang et al., 2019; Klicpera et al., 2019;

Original: Sch ¨ utt et al., 2018).
Revised: Sch ¨ utt et al., 2018).

Original: As GNNs approach the accuracy limit, recent work focuses on improving generalization by integrating calculations from computational chemistry.
Revised: As GNNs approach the accuracy limit, recent work focuses on improving generalization by integrating calculations from computational chemistry.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Despite their success, surrogate models lack reliability.
Revised: Despite their success, surrogate models lack reliability.

Original: 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).
Revised: 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).

Original: While such supervised models are architecturally related, they pursue a fundamentally different objective than PESNet.
Revised: While such supervised models are architecturally related, they pursue a fundamentally different objective than PESNet.

Original: Where surrogate models approximate QM calculations from data, this work focuses on performing the exact QM calculations from scratch.
Revised: Where surrogate models approximate QM calculations from data, this work focuses on performing the exact QM calculations from first principles.

Original: Neural wave function Ans¨atze 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 Schr ¨ odinger equation (Szabo & Ostlund, 2012).
Revised: Neural wave function Ans¨atze 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 Schr¨odinger equation (Szabo & Ostlund, 2012).

Original: However, early works were limited to small systems and low accuracy (Kessler et al., 2021; Han et al., 2019; Choo et al., 2020).
Revised: However, early works were limited to small systems and low accuracy (Kessler et al., 2021; Han et al., 2019; Choo et al., 2020).

Original: 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.
Revised: 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.

Original: To further improve accuracy, Wilson et al.
Revised: To further improve accuracy, Wilson et al.

Original: coupled FermiNet with diffusion Monte-Carlo (DMC).
Revised: coupled FermiNet with diffusion Monte-Carlo (DMC).

Original: But, all these methods need to be trained for each configuration individually.
Revised: But, all these methods need to be trained for each configuration individually.

Original: 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).
Revised: 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).

Original: In a concurrent work, Scherbela et al.
Revised: In a concurrent work, Scherbela et al.

Original: (2021) extend this idea to electronic wave functions.
Revised: (2021) extend this idea to electronic wave functions.

Original: 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.
Revised: 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.

Original: Other efforts have aimed at accelerating Ans¨atze by replacing costly determinant operations, but this comes at a significant loss of accuracy (Acevedo et al., 2020).

Original:  M ETHOD
Revised:  M ETHOD

Original: To build a model that solves the Schr¨odinger equation for many geometries simultaneously and accounts for the symmetries of the energy, we use three key ingredients.
Revised: To build a model that solves the Schr¨odinger equation for many geometries simultaneously and accounts for the symmetries of the energy, we use three key ingredients.

Original: MLP
Revised: MLP

Original: RBF
Revised: RBF

Original: Embedding
Revised: Embedding

Original: Interaction Interaction
Revised: Interaction Interaction

Original: Interaction
Revised: Interaction

Original: MLP
Revised: MLP

Original: MLP Embedding
Revised: MLP Embedding

Original: U pd a t e
Revised: U pd a t e

Original: Envelope U pd a t e U pd a t e U pd a t e
Revised: Envelope U pd a t e U pd a t e U pd a t e

Original: MetaGNN:
Revised: MetaGNN:

Original: Embedding:
Revised: Embedding:

Original: Interaction:
Revised: Interaction:

Original: WFModel:
Revised: WFModel:

Original: Update:
Revised: Update:

Original: Envelope:
Revised: Envelope:

Original: MLP:
Revised: MLP:

Original: MLP MLP
Revised: MLP MLP

Original: Firstly, to solve the Schr ¨ odinger 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.
Revised: Firstly, to solve the Schr ¨ odinger 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.

Original: The WFModel ψ θ ( −→ r ) : R N × 3 (cid:55)→ R is a function parametrized by θ that maps electron configurations to amplitudes.
Revised: The WFModel ψ θ ( −→ r ) : R N × 3 (cid:55)→ R is a function parametrized by θ that maps electron configurations to amplitudes.

Original: 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.
Revised: 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.

Original: As we cover in Section 3.4, the WFModel is essential for sampling electron configurations and computing energies.
Revised: As we cover in Section 3.4, the WFModel is essential for sampling electron configurations and computing energies.

Original: Secondly, we extend this to multiple geometries by introducing a GNN that reparametrizes the WFModel.
Revised: Secondly, we extend this to multiple geometries by introducing a GNN that reparametrizes the WFModel.

Original: In reference to meta-learning, we call this the MetaGNN.
Revised: In reference to meta-learning, we call this the MetaGNN.

Original: It takes the nuclei coordinates −→ R m and charges Z m and outputs subsets ω , ω m ⊂ θ , m ∈ { 1 , . . .
Revised: It takes the nuclei coordinates −→ R m and charges Z m and outputs subsets ω , ω m ⊂ θ , m ∈ { 1 , . . .

Original: , M } of WFModel’s parameters.
Revised: , M } of WFModel’s parameters.

Original: Thanks to message passing, the MetaGNN can capture the full 3D geometry of the nuclei graph
Revised: Thanks to message passing, the MetaGNN can capture the full 3D geometry of the nuclei graph

Original: Lastly, as we prove in Appendix A, to predict energies invariant to rotations and reflections the wave function needs to be equivariant.
Revised: Lastly, as we prove in Appendix A, to predict energies invariant to rotations and reflections the wave function needs to be equivariant.

Original: We accomplish this by constructing an equivariant coordinate system E 
Revised: We accomplish this by constructing an equivariant coordinate system E =

Original: [Equation] based on the principle component analysis (PCA).
Revised: [ −→ e 1 , −→ e 2 , −→ e 3 ] based on the principle component analysis (PCA).

Original: Together, these components form PESNet, whose architecture is shown in Figure 2.
Revised: Together, these components form PESNet, whose architecture is shown in Figure 2.

Original: Since sampling and energy computations only need the WFModel, we only need a single forward pass of the
Revised: Since sampling and energy computations only need the WFModel, we only need a single forward pass of the

Original: MetaGNN for each geometry during evaluation.
Revised: MetaGNN for each geometry during evaluation.

Original: 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;
Revised: 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;

Original: Kondor & Trivedi, 2018).
Revised: Kondor & Trivedi, 2018).

Original: Notation.
Revised: Notation.

Original: 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, −→ R m , Z m for nuclei coordinates and charge, respectively.
Revised: 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, −→ R m , Z m for nuclei coordinates and charge, respectively.

Original: [ ◦ , ◦ ] and [ ◦ ] Ni =1 denote vector concatenations.
Revised: [ ◦ , ◦ ] and [ ◦ ] Ni =1 denote vector concatenations.

Original: 3. W AVE F UNCTION M ODEL
Revised: 3. W AVE F UNCTION M ODEL

Original: We use the FermiNet (Pfau et al., 2020) architecture and augment it with a new feature construction that is invariant to reindexing nuclei.
Revised: We use the FermiNet (Pfau et al., 2020) architecture and augment it with a new feature construction that is invariant to reindexing nuclei.

Original: In the original FermiNet, the inputs to the first layer are simply concatenations of the electron-nuclei distances.
Revised: In the original FermiNet, the inputs to the first layer are simply concatenations of the electron-nuclei distances.

Original: This causes the features to permute if nuclei indexing changes.
Revised: This causes the features to permute if nuclei indexing changes.

Original: To circumvent this issue, we propose a new feature construction as follows:
Revised: To circumvent this issue, we propose a new feature construction as follows:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where z m is an embedding of the m -th nuclei and E ∈ R 3 × 3 is our equivariant coordinate system, see Section 3.3.
Revised: where z m is an embedding of the m -th nuclei and E ∈ R 3 × 3 is our equivariant coordinate system, see Section 3.3.

Original: By summing over all nuclei instead of concatenating we obtain the desired invariance.
Revised: By summing over all nuclei instead of concatenating we obtain the desired invariance.

Original: The features are then iteratively updated using the update rule from Wilson et al.
Revised: The features are then iteratively updated using the update rule from Wilson et al.

Original: [Equation]
Revised: [Equation]

Original: where σ is an activation function, A ↑ and A ↓ are the index sets of the spin-up and spin-down electrons, respectively.
Revised: where σ is an activation function, A ↑ and A ↓ are the index sets of the spin-up and spin-down electrons, respectively.

Original: We also add skip connections where possible.
Revised: We also add skip connections where possible.

Original: We chose σ := tanh since it must be at least twice differentiable to compute the energy, see Section 3.4.
Revised: We chose σ := tanh since it must be at least twice differentiable to compute the energy, see Section 3.4.

Original: After L WF many updates, we take the electron embeddings h L WF i and construct K orbitals:
Revised: After L WF many updates, we take the electron embeddings h L WF i and construct K orbitals:

Original: [Equation]
Revised: [Equation]

Original: where k ∈ { 1 , · · · , K } , α ∈ {↑ , ↓} , i,
Revised: where k ∈ { 1 , · · · , K } , α ∈ {↑ , ↓} , i,

Original: j ∈
Revised: j ∈

Original: A α , and p i , s i are free parameters.
Revised: A α , and p i , s i are free parameters.

Original: Here, we use the sigmoid and softplus functions to ensure that the wave function decays to 0 infinitely far away from any nuclei.
Revised: Here, we use the sigmoid and softplus functions to ensure that the wave function decays to 0 infinitely far away from any nuclei.

Original: To satisfy the antisymmetry to the exchange of same-spin electrons, the output is a weighted sum of determinants (Hutter, 2020)
Revised: To satisfy the antisymmetry to the exchange of same-spin electrons, the output is a weighted sum of determinants (Hutter, 2020)

Original: [Equation]
Revised: [Equation]

Original: 3. M ETA GNN
Revised: 3. M ETA GNN

Original: The MetaGNN’s task is to adapt the WFModel to the geometry at hand.
Revised: The MetaGNN’s task is to adapt the WFModel to the geometry at hand.

Original: It does so by substituting subsets, ω and ω m , of WFModel’s parameters.
Revised: It does so by substituting subsets, ω and ω m , of WFModel’s parameters.

Original: While ω m contains parameters specific to nuclei m , ω is a set of nuclei-independent parameters such as biases.
Revised: While ω m contains parameters specific to nuclei m , ω is a set of nuclei-independent parameters such as biases.

Original: To capture the geometry of the nuclei, the GNN embeds the nuclei in a vector space and updates the embeddings via learning message passing.
Revised: To capture the geometry of the nuclei, the GNN embeds the nuclei in a vector space and updates the embeddings via learning message passing.

Original: 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.
Revised: 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.

Original: Hence, our node embeddings are initialized by
Revised: Hence, our node embeddings are initialized by

Original: [Equation]
Revised: [Equation]

Original:  where G is a matrix of charge embeddings, Z m ∈ N + is the charge of nucleus m , f pos : R 3 (cid:55)→
Revised: [Equation] where G is a matrix of charge embeddings, Z m ∈ N + is the charge of nucleus m , f pos : R 3 (cid:55)→

Original: R N SBF · N RBF is our positional encoding function, and −→ R ′ m
Revised: R N SBF · N RBF is our positional encoding function, and −→ R ′ m

Original: E is the relative position of the m th nucleus in our equivariant coordinate system E (see Section 3.3).
Revised: E is the relative position of the m th nucleus in our equivariant coordinate system E (see Section 3.3).

Original: As positional encoding function, we use the spherical Fourier-Bessel basis functions ˜ a SBF ,ln from Klicpera et al.
Revised: As positional encoding function, we use the spherical Fourier-Bessel basis functions ˜ a SBF ,ln from Klicpera et al.

Original: [Equation]
Revised: [Equation]

Original: (8) with −→ e i being the i th axis of our equivariant coordinate system E .
Revised: (8) with −→ e i being the i th axis of our equivariant coordinate system E .

Original: Unlike Klicpera et al.
Revised: Unlike Klicpera et al.

Original: (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.
Revised: (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.

Original: A message passing layer consists of a message function f msg and an update function f update .
Revised: A message passing layer consists of a message function f msg and an update function f update .

Original: Together, one can compute an update to the embeddings as
Revised: Together, one can compute an update to the embeddings as

Original: [Equation]
Revised: [Equation]

Original: where e mn is an embedding of the edge between nucleus m and nucleus n .
Revised: where e mn is an embedding of the edge between nucleus m and nucleus n .

Original: We use Bessel radial basis functions to encode the distances between nuclei (Klicpera et al., 2019).
Revised: We use Bessel radial basis functions to encode the distances between nuclei (Klicpera et al., 2019).

Original: Both f msg and f update are realized by simple feed-forward neural networks with residual connections.
Revised: Both f msg and f update are realized by simple feed-forward neural networks with residual connections.

Original: After L GNN many message passing steps, we compute WFModel’s parameters on two levels.
Revised: After L GNN many message passing steps, we compute WFModel’s parameters on two levels.

Original: On the global level, f outglobal outputs the biases of the network and, on the node level, f outnode outputs nuclei specific parameters:
Revised: On the global level, f outglobal outputs the biases of the network and, on the node level, f outnode outputs nuclei specific parameters:

Original: [Equation]
Revised: [Equation]

Original: We use distinct feed-forward neural networks with multiple heads for the specific types of parameters estimated to implement f outnode and f outglobal .
Revised: We use distinct feed-forward neural networks with multiple heads for the specific types of parameters estimated to implement f outnode and f outglobal .

Original: 3. E QUIVARIANT C OORDINATE S YSTEMS
Revised: 3. E QUIVARIANT C OORDINATE S YSTEMS

Original: Incorporating symmetries helps to reduce the training space significantly.
Revised: Incorporating symmetries helps to reduce the training space significantly.

Original: 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.
Revised: 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.

Original: While this works for predicting observable metrics such as energies, it does not work for wave functions.
Revised: While this works for predicting observable metrics such as energies, it does not work for wave functions.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: To solve this issue, we introduce directionality in the form of a coordinate system that is equivariant to rotations and reflections.
Revised: To solve this issue, we introduce directionality in the form of a coordinate system that is equivariant to rotations and reflections.

Original: The axes of our coordinate system E =
Revised: The axes of our coordinate system E =

Original: [ −→ e 1 , −→ e 2 , −→ e 3 ] are defined by the principal components of the nuclei coordinates, −→e PCA
Revised: [ −→ e 1 , −→ e 2 , −→ e 3 ] are defined by the principal components of the nuclei coordinates, −→

Original: [Equation]
Revised: [Equation]

Original: Using PCA is robust to reindexing nuclei and ensures that the axes rotate with the system and form an orthonormal basis.
Revised: Using PCA is robust to reindexing nuclei and ensures that the axes rotate with the system and form an orthonormal basis.

Original: However, as PCA only returns directions up to a sign, we have to resolve the sign ambiguity.
Revised: However, as PCA only returns directions up to a sign, we have to resolve the sign ambiguity.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: As equivariant vector we use the difference between a weighted and the regular center of mass
Revised: As equivariant vector we use the difference between a weighted and the regular center of mass

Original: [Equation]
Revised: [Equation]

Original: With this construction, we obtain an equivariant coordinate system that defines directionality in space.
Revised: With this construction, we obtain an equivariant coordinate system that defines directionality in space.

Original: However, we are aware that PCA may not be robust, e.g., if eigenvalues of the covariance matrix are identical.
Revised: However, we are aware that PCA may not be robust, e.g., if eigenvalues of the covariance matrix are identical.

Original: A detailed discussion on such edge cases can be found in Appendix C.
Revised: A detailed discussion on such edge cases can be found in Appendix C.

Original: 3. O PTIMIZATION
Revised: 3. O PTIMIZATION

Original: 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 θ :
Revised: 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 θ :

Original: [Equation]
Revised: [Equation]

Original:  where H is the Hamiltonian of the Schr¨odinger equation
Revised: 14) where H is the Hamiltonian of the Schr¨odinger equation

Original: [Equation]
Revised: [Equation]

Original: ) with ∇ 2 being the Laplacian operator and V ( −→ r ) describing the potential energy.
Revised: with ∇ 2 being the Laplacian operator and V ( −→ r ) describing the potential energy.

Original: Given samples from the probability distribution ∼ ψ 2 θ ( −→ r ) , one can obtain an unbiased estimate of the gradient
Revised: Given samples from the probability distribution ∼ ψ 2 θ ( −→ r ) , one can obtain an unbiased estimate of the gradient

Original: [Equation]
Revised: [Equation]

Original: with E L denoting the local energy.
Revised: where E θ ( −→ r ) denotes the local energy of the WFModel with parameters θ for the electron configuration −→ r .

Original: One can see that for the energy computation, we only need the derivative of the wave function w.r.t.
Revised: One can see that for the energy computation, we only need the derivative of the wave function w.r.t.

Original: the electron coordinates.
Revised: the electron coordinates.

Original: As these are no inputs to the MetaGNN, we do not have to differentiate through the MetaGNN to obtain the local energies.
Revised: As these are no inputs to the MetaGNN, we do not have to differentiate through the MetaGNN to obtain the local energies.

Original: We clip the local energy as in Pfau et al.
Revised: We clip the local energy as in Pfau et al.

Original: and obtain samples from ∼ ψ 2 θ ( −→ r ) via MetropolisHastings.
Revised: and obtain samples from ∼ ψ 2 θ ( −→ r ) via Metropolis-Hastings.

Original: The gradients for the MetaGNN are computed jointly with those of the WFModel by altering Equation 17:
Revised: The gradients for the MetaGNN are computed jointly with those of the WFModel by altering Equation 17:

Original: [Equation]
Revised: [Equation]

Original: where Θ is the joint set of WFModel and MetaGNN parameters.
Revised: where Θ is the joint set of WFModel and MetaGNN parameters.

Original: To obtain the gradient for multiple geometries, we compute the gradient as in Equation 18 multiple times and average.
Revised: To obtain the gradient for multiple geometries, we compute the gradient as in Equation 18 multiple times and average.

Original: This joint gradient of the WFModel and the MetaGNN enables us to use a single training pass to simultaneously solve multiple Schr ¨ odinger equations.
Revised: This joint gradient of the WFModel and the MetaGNN enables us to use a single training pass to simultaneously solve multiple Schr ¨ odinger equations.

Original: 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).
Revised: 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).

Original: Here, we use natural gradient descent to enable the use of larger learning rates.
Revised: Here, we use natural gradient descent to enable the use of larger learning rates.

Original: So, instead of doing a regular gradient descent step in the form of Θ t +1 =
Revised: So, instead of doing a regular gradient descent step in the form of Θ t +1 =

Original: Θ t − η ∇ Θ L , where η is the learning rate, we add the inverse of the Fisher information matrix as a preconditioner
Revised: Θ t − η ∇ Θ L , where η is the learning rate, we add the inverse of the Fisher information matrix as a preconditioner

Original: [Equation]
Revised: [Equation]

Original: 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).
Revised: 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).

Original: To determine the convergence of the CG method, we follow Martens (2010) and stop based on the quadratic error.
Revised: To determine the convergence of the CG method, we follow Martens (2010) and stop based on the quadratic error.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: During pretraining, only the WFModel and the final biases of the MetaGNN are optimized.
Revised: During pretraining, only the WFModel and the final biases of the MetaGNN are optimized.

Original: 3. L IMITATIONS
Revised: 3. L IMITATIONS

Original: While PESNet is capable of accurately modeling complex potential energy surfaces, we have not focused on architecture search yet.
Revised: While PESNet is capable of accurately modeling complex potential energy surfaces, we have not focused on architecture search yet.

Original: For instance, both the WFModel and the MetaGNN are only as powerful as the Weisfeiler-Lehman graph isomorphism test (Xu et al., 2018).

Original: Furthermore, as we discuss in Section 4, PauliNet  still offers a better initialization and converges in fewer iterations than our network.
Revised: Furthermore, as we discuss in Section 4, PauliNet (Hermann et al., 2020) still offers a better initialization and converges in fewer iterations than our network.

Original: 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.
Revised: Lastly, PESNet is limited to geometries of the same set of nuclei with identical electron spin configurations, i.e., to access properties like the electron affinity one still needs to train two models.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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) 
Revised: 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.

Revised: Note, the concurrently developed DeepErwin was only recently released as a pre-print and still requires separate models and training for each configuration.

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

Original: [Figure]
Revised: [Figure]

Original: E i n E h
Revised: E i n E h

Original: [Figure]
Revised: [Figure]

Original: PESNet
Revised: PESNet

Original: H HH
Revised: H HH

Original: H
Revised: H

Original: H HH
Revised: H HH

Original: H
Revised: H

Original: H HH
Revised: H HH

Original: H
Revised: H

Original: [Figure]
Revised: [Figure]

Original: H
Revised: H

Original: H
Revised: H

Original: [Figure]

Original: H 
Revised: H H

Original: Since PESNet respects the symmetriesof the energy, we only trained on half of the config 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.
Revised: Since PESNet respects the symmetries

Original: Evaluation of ab-initio methods is still a challenge as true energies are rarely known, and experimental data are subject to uncertainties.
Revised: Evaluation of ab-initio methods is still a challenge as true energies are rarely known, and experimental data are subject to uncertainties.

Original: In addition, many energy differences may seem small due to their large absolute values, but chemists set the threshold for chemical accuracy to 1 kcal mol − 1 ≈ 1 . 6 m E h .
Revised: In addition, many energy differences may seem small due to their large absolute values, but chemists set the threshold for chemical accuracy to 1 kcal mol − 1 ≈ 1 . 6 m E h .

Original: Thus, seemingly small differences in energy are significant.
Revised: Thus, seemingly small differences in energy are significant.

Original: Therefore, to put all results into perspective, we always include highly accurate classical reference calculations.
Revised: Therefore, to put all results into perspective, we always include highly accurate classical reference calculations.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: The exact procedure and the general experimental setup are described in Appendix D.
Revised: The exact procedure and the general experimental setup are described in Appendix D.

Original: Additional ablation studies are available in Appendix E.
Revised: Additional ablation studies are available in Appendix E.

Original: Transition path of H +4 and weight sharing.
Revised: Transition path of H +4 and weight sharing.

Original: Scherbela et al.
Revised: Scherbela et al.

Original: (2021) use the first transition path of
Revised: (2021) use the first transition path of

Original: H +4 (Alijah & Varandas, 2008) to demonstrate the acceleration gained by weight-sharing.
Revised: H +4 (Alijah & Varandas, 2008) to demonstrate the acceleration gained by weight-sharing.

Original: But, they found their weight-sharing scheme to be too restrictive and additionally optimized each wave function separately.
Revised: But, they found their weight-sharing scheme to be too restrictive and additionally optimized each wave function separately.

Original: Unlike DeepErwin, our novel PESNet is flexible enough such that we do not need any extra optimization.
Revised: Unlike DeepErwin, our novel PESNet is flexible enough such that we do not need any extra optimization.

Original: In Figure 3, we see the DeepErwin results after their multi-step optimization and the energies of a single PESNet.
Revised: In Figure 3, we see the DeepErwin results after their multi-step optimization and the energies of a single PESNet.

Original: We notice that while both methods estimate similar transition barriers PESNet results in ≈ 0 .
Revised: We notice that while both methods estimate similar transition barriers PESNet results in ≈ 0 .

Original: 27 m E h smaller energies which agree with the very accurate
Revised: 27 m E h smaller energies which agree with the very accurate

Original: MRCI-D-F12 results ( ≈ 0 .
Revised: MRCI-D-F12 results ( ≈ 0 .

Original: 015 m E h ).
Revised: 015 m E h ).

Original: Hydrogen rectangle and symmetries.
Revised: Hydrogen rectangle and symmetries.

Original: The Hydrogen rectangle is a known failure case for CCSD and CCSD(T).
Revised: The Hydrogen rectangle is a known failure case for CCSD and CCSD(T).

Original: While the exact solution, FCI, indicates a local maximum at θ = 90 ◦ , both, CCSD and CCSD(T), predict local minima.
Revised: While the exact solution, FCI, indicates a local maximum at θ = 90 ◦ , both, CCSD and CCSD(T), predict local minima.

Original: Figure 4 shows that VMC methods such as FermiNet and our PESNet do not suffer from the same issue.
Revised: Figure 4 shows that VMC methods such as FermiNet and our PESNet do not suffer from the same issue.

Original: PESNet’s energies are identical to FermiNet’s
Revised: PESNet’s energies are identical to FermiNet’s

Original: ( ≈ 0 .
Revised: ( ≈ 0 .

Original: 014 m E h ) 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.
Revised: 014 m E h ) despite training only a single network on half of the configuration space thanks to our equivariant coordinate system.

Original: 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.
Revised: 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.

Original: Figure 5 shows the potential energy surface of the hydrogen chain computed by a range of methods.
Revised: Figure 5 shows the potential energy surface of the hydrogen chain computed by a range of methods.

Original: While our PESNet generally performs identical to FermiNet, we predict on average 0 .
Revised: While our PESNet generally performs identical to FermiNet, we predict on average 0 .

Original: 31 m E h lower energies.
Revised: 31 m E h lower energies.

Original: Further, we notice that our results are consistently better than PauliNet and DeepErwin despite only training a single model.[Figure] H [Figure]
Revised: Further, we notice that our results are consistently better than PauliNet and DeepErwin despite only training a single model.

Original: [Figure] 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).
Revised: 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).

Original: While the accurate r12-MR-ACPF method more closely matches the experimental results, it scales factorially (Gdanitz, 1998).
Revised: While the accurate r12-MR-ACPF method more closely matches the experimental results, it scales factorially (Gdanitz, 1998).

Original: After Pfau et al.
Revised: After Pfau et al.

Original: (2020) have shown that FermiNet is capable of modeling such complex triple bonds, we are interested in PESNet’s performance.
Revised: (2020) have shown that FermiNet is capable of modeling such complex triple bonds, we are interested in PESNet’s performance.

Original: 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.
Revised: To better represent both methods, we [Figure] [Figure] decided to compare both FermiNet as well as PESNet with 32 determinants due to a substantial performance gain for both methods.

Original: The results in Figure 6 show that PESNet agrees very well with FermiNet and is on average just 0 .
Revised: The results in Figure 6 show that PESNet agrees very well with FermiNet and is on average just 0 .

Original: 37 m E h higher, despite training only a single model for less than 140 of FermiNet’s training time, see Section 4.1.
Revised: 37 m E h higher, despite training only a single model for less than 147 of FermiNet’s training time, see Section 4.1.

Revised: In addition, the ablation of PESNet without the

Revised: MetaGNN shows a significant loss of accuracy of 4 .

Revised: 3 m E h on average.

Original: Cyclobutadiene and the MetaGNN.
Revised: Cyclobutadiene and the MetaGNN.

Original: 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).
Revised: 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).

Original: In contrast, PauliNet and FermiNet had success at modelling this challenging system.
Revised: In contrast, PauliNet and FermiNet had success at modelling this challenging system.

Original: Naturally, we are interested in how well PESNet can estimate the transition barrier.
Revised: Naturally, we are interested in how well PESNet can estimate the transition barrier.

Original: To be comparable to Spencer et al.
Revised: To be comparable to Spencer et al.

Original: (2020), we increased the number of determinants for the experiment to 32 and the single-stream size to 512.
Revised: (2020), we increased the number of determinants toand the single-stream size to 512.

Revised: Similar to PauliNet Hermann et al.

Revised: , we found PESNet to occasionally converge to a higher energy for the transition state depending on the initialization and pretraining.

Revised: To avoid this, we pick a 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.

Original: 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.
Revised: 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.

Original: But, they require different numbers of training steps and result in different total energies.
Revised: But, they require different numbers of training steps and result in different total energies.

Original: PauliNet generally converges fastest, but results in the highest energies, whereas FermiNet’s transition barrier converges slower but its energies are 70 m E h smaller.
Revised: PauliNet generally converges fastest, but results in the highest energies, whereas FermiNet’s transition barrier converges slower but its energies are 70 m E h smaller.

Original: Lastly, PESNet’s transition barrier converges similar PauliNet’s, but its energies are 54 m E h lower than PauliNet’s, placing it closer to FermiNet than PauliNet in terms of accuracy.
Revised: Lastly, PESNet’s transition barrier converges similar PauliNet’s, but its energies are 54 m E h lower than PauliNet’s, placing it closer to FermiNet than PauliNet in terms of accuracy.

Original: 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.
Revised: Considering that PESNet has only been trained for ≈ 16 of FermiNet’s time (see Section 4.1), we are confident that additional optimization would further narrow the gap to FermiNet.

Original: 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.
Revised: In an additional ablation study, we compare to PESNet without the MetaGNN, i.e, we still train a single model for both states of cyclobutadiene but without weight adaptation.

Original: 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 18 m E
Revised: While the results in Figure 7 show that the truncated network’s energies continuously decrease, it fails to reproduce the same transition barrier and its energies are 18 m E

Original: h worse compared to the full PESNet.
Revised: h worse compared to the full PESNet.

Original: 4. T RAINING T IME
Revised: 4. T RAINING T IME

Original: While the previous experiments have shown that our model’s accuracy is on par with FermiNet,
Revised: While the previous experiments have shown that our model’s accuracy is on par with FermiNet,

Original: PESNet’s main appeal is its capability to fit multiple geometries simultaneously.
Revised: PESNet’s main appeal is its capability to fit multiple geometries simultaneously.

Original: Here, we study the training times for all systems from the previous section.
Revised: Here, we study the training times for all systems from the previous section.

Original: 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.
Revised: We compare the official JAX (Bradbury et al., 2018) 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 implementation of PESNet.

Revised: We use the same hyperparameters as in the experiments or the defaults from the respective works.

Revised: All measurements have been conducted on a machine with 16 AMD EPYC 7543 cores and a single Nvidia A100 GPU.

Original: We use the same hyperparameters as in the
Revised: Here, we only mea- [Figure] [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: r g y
Revised: r g y

Original: [Figure]
Revised: [Figure]

Original: h
Revised: h

Original: Paulinet FermiNet
Revised: Paulinet FermiNet

Original: PESNet PESNet, no GNN
Revised: PESNet PESNet, no GNN

Original: [Figure]
Revised: [Figure]

Original: ∆ E i n k c a l m o l
Revised: ∆ E i n k c a l m o l

Original: PauliNet FermiNet
Revised: PauliNet FermiNet

Original: PESNet PESNet, no GNN
Revised: PESNet PESNet, no GNN

Original: CCSD(T) MR-CC experiment
Revised: CCSD(T) MR-CC experiment

Original: C C C C
Revised: C C C C

Original: H
Revised: H

Original: H H
Revised: H H

Original: H C
Revised: H C

Original: C C C
Revised: C C C

Original: H
Revised: H

Original: H H
Revised: H H

Original: H C
Revised: H C

Original: C C C
Revised: C C C

Original: H
Revised: H

Original: H H
Revised: H H

Original: H
Revised: H

Original: H +4 (Fig. 3)
Revised: H +4 (Fig. 3)

Original: H4 (Fig. 4)
Revised: H4 (Fig. 4)

Original: H10 (Fig. 5)
Revised: H10 (Fig. 5)

Original: N2 (Fig. 6)
Revised: N2 (Fig. 6)

Original: Cyclobutadiene (Fig. 7)
Revised: Cyclobutadiene (Fig. 7)

Original: PauliNet 43h * 34h * 153h854h * 437h
Revised: PauliNet [Table] [Table] * 153h

Revised: [Table]

Revised: * 437h

Original: DeepErwin 34h 27h * 111h — —
Revised: DeepErwin [Table] * 111h — —

Original: FermiNet 154h * 163h 820h 4196h 1120h PESNet 20h 24h 65h 89h 336h 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.
Revised: [Table]

Original: 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.
Revised: 381h sure 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.

Original: Furthermore, note that the timings refer to training all models of the respective experiments.
Revised: Furthermore, note that the timings refer to training all models of the respective experiments.

Original: Table 1 shows the GPU hours to train the models of the last section.
Revised: Table 1 shows the GPU hours to train the models of the last section.

Original: It is apparent that PESNet used the fewest GPU hours across all systems.
Revised: It is apparent that PESNet used the fewest GPU hours across all systems.

Original: Compared to the similar accurate FermiNet, PESNet is up to 47 times faster to train.
Revised: Compared to the similar accurate FermiNet, PESNet is up to 47 times faster to train.

Original: This speedup is especially noticeable if many configurations shall be evaluated, e.g., 39 nitrogen geometries.
Revised: This speedup is especially noticeable if many configurations shall be evaluated, e.g., 39 nitrogen geometries.

Original: Compared to the less accurate PauliNet and DeepErwin,
Revised: Compared to the less accurate PauliNet and DeepErwin,

Original: PESNet’s speed gain shrinks, but our training times are still consistently lower while achieving significantly better results.
Revised: PESNet’s speed gain shrinks, but our training times are still consistently lower while achieving significantly better results.

Original: 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.
Revised: 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.

Original: Thus, if one is interested in additional configurations, PESNet’s speedup grows linearly in the number of configurations.
Revised: Thus, if one is interested in additional configurations, PESNet’s speedup grows linearly in the number of configurations.

Original: 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.
Revised: 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.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We presented a novel architecture that can simultaneously solve the Schr¨odinger equation for multiple geometries.
Revised: We presented a novel architecture that can simultaneously solve the Schr¨odinger equation for multiple geometries.

Original: 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.
Revised: 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.

Original: The integration of physical symmetries enables us to reduce our training space.
Revised: The integration of physical symmetries enables us to reduce our training space.

Original: Finally, our results show that a single model can capture a continuous subset of the potential energy surface.
Revised: Finally, our results show that a single model can capture a continuous subset of the potential energy surface.

Original: This acceleration of neural wave functions opens access to accurate quantum mechanical calculations to a broader audience.
Revised: This acceleration of neural wave functions opens access to accurate quantum mechanical calculations to a broader audience.

Original: 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.
Revised: 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.

Original: Ethics and reproducibility.
Revised: Ethics and reproducibility.

Original: Advanced computational chemistry tools may have a positive impact in chemistry research, for instance in material discovery.
Revised: Advanced computational chemistry tools may have a positive impact in chemistry research, for instance in material discovery.

Original: However, they also pose the risk of misuse, e.g., for the development of chemical weapons.
Revised: However, they also pose the risk of misuse, e.g., for the development of chemical weapons.

Original: To the best of our knowledge, our work does not promote misuse any more than general computational chemistry research.
Revised: To the best of our knowledge, our work does not promote misuse any more than general computational chemistry research.

Original: To reduce the likelihood of such misuse, we publish our source code under the Hippocratic license (Ehmke, 2019).
Revised: To reduce the likelihood of such misuse, we publish our source code under the Hippocratic license (Ehmke, 2019) 1 .

Original: To facilitate reproducibility, the source code includes simple scripts to reproduce all experiments from
Revised: To facilitate reproducibility, the source code includes simple scripts to reproduce all experiments from

Original: Section 4.
Revised: Section 4.

Original: 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.
Revised: 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.

Revised: Acknowledgements.

Revised: 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 Klicpera and Leon Hetzel for their valuable feedback.

Revised: 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¨ander.

Original: A I NVARIANT E NERGIES R EQUIRE E QUIVARIANT W AVE F UNCTIONS
Revised: A I NVARIANT E NERGIES R EQUIRE E QUIVARIANT W AVE F UNCTIONS

Original: Here, we prove that a wave function needs to be equivariant with respect to rotations and reflections if the energy is invariant.
Revised: Here, we prove that a wave function needs to be equivariant with respect to rotations and reflections if the energy is invariant.

Original: Recall, our goal is to solve the stationary Schr ¨ odinger equation
Revised: Recall, our goal is to solve the stationary Schr ¨ odinger equation

Original: [Equation]
Revised: [Equation]

Original: (21) where ψ : R 3 N (cid:55)→ R is the electronic wave function.
Revised: (21) where ψ : R 3 N (cid:55)→ R is the electronic wave function.

Original: Since the Hamiltonian H encodes the molecular structure via its potential energy term, see Equation 15, a rotation or reflection U ∈ O (3) of the system results is a unitary transformation applied to the Hamiltonian H → UHU † .
Revised: Since the Hamiltonian H encodes the molecular structure via its potential energy term, see Equation 15, a rotation or reflection U ∈ O (3) of the system results is a unitary transformation applied to the Hamiltonian H → UHU † .

Original: Since the energy E is invariant to rotation and reflection, we obtain the transformed equation
Revised: Since the energy E is invariant to rotation and reflection, we obtain the transformed equation

Original: [Equation]
Revised: [Equation]

Original: One can see that if ψ is a solution to Equation 21, the equivariantly transformed ψ → U ψ solves
Revised: One can see that if ψ is a solution to Equation 21, the equivariantly transformed ψ → U ψ solves

Original: Equation
Revised: Equation

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: (25) with the same energy.
Revised: (25) with the same energy.

Original: □
Revised: □

Original: B E QUIVARIANT N EURAL N ETWORKS A S W AVE F UNCTIONS
Revised: B E QUIVARIANT N EURAL N ETWORKS A S W AVE F UNCTIONS

Original: Here, we want to briefly discuss why equivariant neural networks as proposed by Thomas et al.
Revised: Here, we want to briefly discuss why equivariant neural networks as proposed by Thomas et al.

Original: (2018) or Batzner et al.
Revised: (2018) or Batzner et al.

Original: (2021) are no alternative to our equivariant coordinate system.
Revised: (2021) are no alternative to our equivariant coordinate system.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: However, if one would construct the orbitals ϕ ( −→ r ) =
Revised: However, if one would construct the orbitals ϕ ( −→ r ) =

Original: [ ϕ 1 ( −→ r ) , . . .
Revised: [ ϕ 1 ( −→ r ) , . . .

Original: , ϕ N ( −→ r )] by concatenating
Revised: , ϕ N ( −→ r )] by concatenating

Original: E equivariant SO(3) representations ϕ ( −→ r )
Revised: E equivariant SO(3) representations ϕ ( −→ r )

Original: =
Revised: =

Original: [ ϕ 1 ( −→ r ) , . . . , ϕ E ( −→ r )] with (cid:80) Ee =1 dim( ϕ e ( −→ r i ))
Revised: [ ϕ 1 ( −→ r ) , . . . , ϕ E ( −→ r )] with (cid:80) Ee =1 dim( ϕ e ( −→ r i ))

Original: =
Revised: =

Original: N , any resulting real-valued wave function ψ ( −→ r ) = det ϕ ( −→ r ) would be spherically symmetric, i.e., ψ ( −→ r R ) =
Revised: N , any resulting real-valued wave function ψ ( −→ r ) = det ϕ ( −→ r ) would be spherically symmetric, i.e., ψ ( −→ r R ) =

Original: ψ ( −→ r ) , ∀ R ∈ SO (3) .
Revised: ψ ( −→ r ) , ∀ R ∈ SO (3) .

Original: The proof is as follows: If one rotates the electrons −→ r ∈ R N × 3 by any rotation matrix R ∈ SO (3) , the orbital matrix changes as[Equation]
Revised: The proof is as follows: If one rotates the electrons −→ r ∈ R N × 3 by any rotation matrix R ∈ SO (3) , the orbital matrix changes as

Original: [Equation]
Revised: [Equation]

Original: 27) where D R ∈ R N × N is a block-diagonal matrix and D Re is the Wigner-D matrix induced by rotation
Revised: where D R ∈ R N × N is a block-diagonal matrix and D Re is the Wigner-D matrix induced by rotation

Original: R corresponding to the e -th SO(3) representation.
Revised: R corresponding to the e -th SO(3) representation.

Original: Since Wigner-D matrices are unitary and we restrict our wave function to real-valued
Revised: Since Wigner-D matrices are unitary and we restrict our wave function to real-valued

Original: [Equation]
Revised: [Equation]

Original: □

Original: C E DGE C ASES OF THE E QUIVARIANT C OORDINATE S YSTEM
Revised: C E DGE C ASES OF THE E QUIVARIANT C OORDINATE S YSTEM

Original: [Figure]
Revised: [Figure]

Original: While the definition of the coordinate system in Section 3.3 works in most cases.
Revised: While the definition of the coordinate system in Section 3.3 works in most cases.

Original: There still exist edge cases where the coordinate system may not be unique.
Revised: There still exist edge cases where the coordinate system may not be unique.

Original: To maximize transparency, we discuss some of these cases, when they occur, how we handle them, and what their implications are.
Revised: To maximize transparency, we discuss some of these cases, when they occur, how we handle them [Figure] and what their implications are.

Original: For certain geometries, two eigenvalues of the nuclei covariance matrix might be identical.
Revised: For certain geometries, two eigenvalues of the nuclei covariance matrix might be identical.

Original: If that is the case the PCA axes are not unique.
Revised: If that is the case the PCA axes are not unique.

Original: This occurs for any geometry where the nuclei coordinates are distributed regularly on a sphere around the center of mass.
Revised: This occurs for any geometry where the nuclei coordinates are distributed regularly on a sphere around the center of mass.

Original: Examples of such geometries are regular polygons such as the pentagram depicted in Figure 8a.
Revised: Examples of such geometries are regular polygons such as the pentagram depicted in Figure 8a.

Original: In such cases, we compute the PCA on pseudo coordinates which we obtain by stretching the graph in the direction of the shortest edge −→ R m −−→ R n .
Revised: In such cases, we compute the PCA on pseudo coordinates which we obtain by stretching the graph in the direction of the largest Coulomb potential [Equation]

Original: 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.
Revised: 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.

Original: The actual direction does not matter, as it is simply a rotation or reflection of the whole system.
Revised: The actual direction does not matter, as it is simply a rotation or reflection of the whole system.

Original: While regular spherical patterns are not the only case where this issue arises they are the most common cause.
Revised: While regular spherical patterns are not the only case where this issue arises they are the most common cause.

Original: 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.
Revised: 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.

Original: In such cases, the signs of the axes do not matter as reflections result in identical geometries.
Revised: In such cases, the signs of the axes do not matter as reflections result in identical geometries.

Original: However, there also exist other geometries for which Equation 12 is 0 .
Revised: However, there also exist other geometries for which Equation 12 is 0 .

Original: An example is depicted in Figure 8b.
Revised: An example is depicted in Figure 8b.

Original: But these cases are rare and can be resolved by applying a nonlinear function on the distances in Equation 12.
Revised: But these cases are rare and can be resolved by applying a nonlinear function on the distances in Equation 12.

Original: The occurrence of these edge cases leads us to question: ‘Why can we not design a unique coordinate system for each geometry?’.
Revised: The occurrence of these edge cases leads us to question: ‘Why can we not design a unique coordinate system for each geometry?’.

Original: While we ideally would want an equivariant coordinate system that changes smoothly with changes in the geometry, it is impossible.
Revised: While we ideally would want an equivariant coordinate system that changes smoothly with changes in the geometry, it is impossible.

Original: We show a counter example of this in Figure 8c where we see a system of 3 nuclei.
Revised: We show a counter example of this in Figure 8c where we see a system of 3 nuclei.

Original: In their starting configuration, one can uniquely define two axes as indicated by the colored arrows.
Revised: In their starting configuration, one can uniquely define two axes as indicated by the colored arrows.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: So, we neither can define a smoothly changing coordinate system nor a unique one for every system.
Revised: So, we neither can define a smoothly changing coordinate system nor a unique one for every system.

Original: However, in practice, we do not need a smoothly changing one but only a unique one.
Revised: However, in practice, we do not need a smoothly changing one but only a unique one.

Original: 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.
Revised: 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.

Original: Considering these aspects, we believe that our coordinate system definition is sufficient in most scenarios.
Revised: Considering these aspects, we believe that our coordinate system definition is sufficient in most scenarios.

Original: D E XPERIMENTAL S ETUP
Revised: D E XPERIMENTAL S ETUP

Original: Hyperparameters.
Revised: Hyperparameters.

Original: If not otherwise specified we used the hyperparameters from Table 2.
Revised: If not otherwise specified we used the hyperparameters from Table 2.

Original: These result in a similarly sized WFModel to FermiNet from Pfau et al.
Revised: These result in a similarly sized WFModel to FermiNet from Pfau et al.

Original: For cyclobutadiene we did not train for the full 60000 iterations but stopped the training after 2 weeks.
Revised: For cyclobutadiene, we did not train for the full 60000 iterations but stopped the training after 2 weeks.

Original: 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.
Revised: 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.

Original: Furthermore, we compute the final output of the WFModel in log-domain and use the log-sum-exp trick.
Revised: Furthermore, we compute the final output of the WFModel in log-domain and use the log-sum-exp trick.

Original: Learning continuous subsets.
Revised: Learning continuous subsets.

Original: 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.
Revised: 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.

Original: Specifically, we subdivide the potential energy surface into even-sized bins and place a random walker within each bin.
Revised: Specifically, we subdivide the potential energy surface into even-sized bins and place a random walker within each bin.

Original: These walkers slightly alter the molecular structure after each step by moving within their bin.
Revised: These walkers slightly alter the molecular structure after each step by moving within their bin.

Original: This procedure ensures that our model is not restricted to a discrete set of configurations.
Revised: This procedure ensures that our model is not restricted to a discrete set of configurations.

Original: Therefore, after training, our model can be evaluated at any arbitrary configuration within the training domain without retraining.
Revised: Therefore, after training, our model can be evaluated at any arbitrary configuration within the training domain without retraining.

Original: But, we only evaluate PESNet on configurations where reference calculations are available.
Revised: But, we only evaluate PESNet on configurations where reference calculations are available.

Original: This procedure is done for the hydrogen rectangle, the hydrogen chain, and the nitrogen molecule.
Revised: This procedure is done for the hydrogen rectangle, the hydrogen chain, and the nitrogen molecule.

Original: For H +4 and cyclobutadiene, we train on discrete sets of geometries from the literature (Scherbela et al., 2021; Kinal & Piecuch, 2007).
Revised: For H +4 and cyclobutadiene, we train on discrete sets of geometries from the literature (Scherbela et al., 2021; Kinal & Piecuch, 2007).

Original: 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.
Revised: 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.

Original: However, in cases where we optimize for a continuous set of configurations, the optimal energy value depends on the current batch of geometries.
Revised: However, in cases where we optimize for a continuous set of configurations, the optimal energy value depends on the current batch of geometries.

Original: To still access convergence, we use the fact that the local energy E L of any eigenfunction (including the ground-state) has 0 variance.
Revised: To still access convergence, we use the fact that the local energy E L of any eigenfunction (including the ground-state) has 0 variance.

Original: So, our convergence criteria is the expected variance of the local energy E 
Revised: So, our convergence criteria is the expected variance of the local energy E −→

Original: [Equation]
Revised: [Equation]

Original: where the optimal value is 0.
Revised: where the optimal value is 0.

Original: [Table]
Revised: [Table]

Revised: Value

Original: Optimization
Revised: Optimization

Original: Local energy clipping
Revised: Local energy clipping

Original: [Table]
Revised: [Table]

Original: Optimization
Revised: Optimization

Original: Batch size
Revised: Batch size

Original: Optimization
Revised: Optimization

Original: Iterations
Revised: Iterations

Original: Optimization #Geometry random walker
Revised: [Table] #Geometry [Table] walker

Original: Optimization
Revised: Optimization

Original: Learningrate η
Revised: Learning

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: Natural 
Revised: Natural Gradient

Revised: Damping

Original: [Table]
Revised: [Table]

Original: Damping 10 − 4 · Std [ E L ]
Revised: E L ]

Original: [Table]
Revised: [Table]

Original: Gradient

Original: CG max steps
Revised: CG max steps

Original: WFModel
Revised: WFModel

Original: Nuclei embedding dim
Revised: Nuclei embedding dim

Original: WFModel
Revised: WFModel

Original: Single-stream width
Revised: Single-stream width

Original: WFModel
Revised: WFModel

Original: Double-stream width
Revised: Double-stream width

Original: WFModel #Update layers
Revised: WFModel #Update layers

Original: WFModel #Determinants
Revised: WFModel #Determinants

Original: MetaGNN #Message passings
Revised: MetaGNN #Message passings

Original: MetaGNN
Revised: MetaGNN

Original: Embedding dim
Revised: Embedding dim

Original: MetaGNN
Revised: [Table]

Original: Message dim
Revised: Message dim

Original: MetaGNN
Revised: MetaGNN

Original: N SBF
Revised: N SBF

Original: MetaGNN
Revised: MetaGNN

Original: N RBF
Revised: N RBF

Original: MetaGNN
Revised: MetaGNN

Original: MLP depth
Revised: MLP depth

Original: MCMC
Revised: MCMC

Original: Proposal step size 0.
Revised: Proposal step size 0.

Original: MCMC
Revised: MCMC

Original: Steps between updates
Revised: Steps between updates

Original: Pretraining
Revised: [Table]

Original: Iterations
Revised: Iterations

Original: [Table]
Revised: [Table]

Original: Learning
Revised: Learning

Original: [Table]
Revised: [Table]

Original: Method
Revised: Method

Original: RHF
Revised: UHF

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Numbers in brackets indicate the standard error at the last digit(s).
Revised: - bers in brackets indicate the standard error at the last digit(s).

Original: In cases without the MetaGNN, we still train a single model for all configurations of a system.
Revised: In cases without the MetaGNN, we still train a single model for all configurations of a system.

Original: E A BLATION S TUDIES
Revised: E A BLATION S TUDIES

Original: As hyperparameters often play a significant in the machine learning community, we present some ablation studies in this appendix.
Revised: As hyperparameters often play a significant in the machine learning community, we present some ablation studies in this appendix.

Original: All the following experiments only alter one variable at a time, while the rest is as in Table 2.
Revised: All the following experiments only alter one variable at a time, while the rest is as in Table 2.

Original: The results in the tables are averaged over the same configurations as in the main body.
Revised: The results in the tables are averaged over the same configurations as in the main body.

Original: Table 3 presents results with and without the MetaGNN.
Revised: Table 3 presents results with and without the MetaGNN.

Original: It is noticeable that the gain of the
Revised: It is noticeable that the gain of the

Original: MetaGNN is little to none for small molecules consisting of simple hydrogen atoms.
Revised: MetaGNN is little to none for small molecules consisting of simple hydrogen atoms.

Original: But, for more #Dets
Revised: But, for more #Dets

Original: H + H H N
Revised: [Table]

Original: Cyclobutadiene -1.849363(6)
Revised: Cyclobutadiene - [Table]

Original: ∆ E -0.000021(7) -0.000019(8) 0.000083(20) 0.00136(10) -0.0015(4)

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: ∆ E -

Original: [Figure] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table][Figure] [Table] [Figure]
Revised: [Table]

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table][Figure] [Table] [Figure]
Revised: [Table]

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: complex molecules such as nitrogen and cyclobutadiene
Revised: complex molecules such as nitrogen and cyclobutadiene

Original: , we notice significant improvements of 2 .
Revised: , we notice significant improvements of 2 .

Original: 5 m E h and 10 .
Revised: 5 m E h and 10 .

Original: 5 m E h , respectively.
Revised: 5 m E h , respectively.

Original: Moreover, the MetaGNN enables us to account for symmetries of the energy while the network without is only invariant w.r.t.
Revised: Moreover, the MetaGNN enables us to account for symmetries of the energy while the network without is only invariant w.r.t.

Original: to translation but not to rotation, reflection, and reindexing of nuclei.
Revised: to translation but not to rotation, reflection, and reindexing of nuclei.

Original: Table 4 shows the impact of the number of determinants on the average energy for the systems from Section 4.
Revised: 4 shows the impact of the number of determinants on the average energy for the systems from Section 4.

Original: For small hydrogen-based systems, the number of determinants is mostly irrelevant, while larger numbers of determinants improve performance for nitrogen.
Revised: For small hydrogen-based systems, the number of determinants is mostly irrelevant, while larger numbers of determinants improve performance for nitrogen.

Original: But, this does not seem to carry over to cyclobutadiene.
Revised: But, this does not seem to carry over to cyclobutadiene.

Original: While the total estimated energy is higher for the larger model, we noticed that it is significantly faster at converging the transition barrier.
Revised: While the total estimated energy is higher for the larger model, we noticed that it is significantly faster at converging the transition barrier.

Original: Figure 9 illustrates this by comparing the convergence of the 16 and 32 determinant models.
Revised: Figure 9 illustrates this by comparing the convergence of the 16 and 32 determinant models.

Original: Increasing the single-stream size does not seem to result in any benefit for most models as Tableshows.
Revised: Increasing the single-stream size does not seem to result in any benefit for most models as Tableshows.

Original: Again, the hydrogen systems are mostly unaffected by this hyperparameter while nitrogen benefits from larger hidden dimensions but cyclobutadiene converges worse.
Revised: Again, the hydrogen systems are mostly unaffected by this hyperparameter while nitrogen benefits from larger hidden dimensions but cyclobutadiene converges worse.

Original: We suspect that this is due to the optimization problem becoming significantly harder.
Revised: We suspect that this is due to the optimization problem becoming significantly harder.

Original: Firstly, increasing the WFModel increases the number of parameters the MetaGNN has to predict.
Revised: Firstly, increasing the WFModel increases the number of parameters the MetaGNN has to predict.

Original: 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.
Revised: 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.

Original: F T IME PER I TERATION
Revised: F T IME PER I TERATION

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: While we find these numbers to be misleading we still want to disclose them to support open research.
Revised: While we find these numbers to be misleading we still want to disclose them to support open research.

Original: 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
Revised: 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

Original: K-FAC (Martens & Grosse, 2015), and PESNet with CG-computed natural gradient (Neuscamman et al., 2012).
Revised: K-FAC (Martens & Grosse, 2015), and PESNet with CG-computed natural gradient (Neuscamman et al., 2012).

Original: This has implications on the number of iterations one has to train.
Revised: This has implications on the number of iterations one has to train.

Original: 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),
Revised: 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 (

Original: H +
Revised: H +

Original: H H
Revised: H H

Original: N
Revised: N

Original: Cyclobutadiene
Revised: Cyclobutadiene

Original: PauliNet 0.83s 1.13s 5.51s 8.09s 175s
Revised: PauliNet [Table] .83s

Original: DeepErwin [Table]
Revised: [Table] [Table]

Original: .92s

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: 175s

Original: FermiNet
Revised: DeepErwin

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: and PESNet for 60000 iterations.
Revised: and PESNet for 60000 iterations.

Original: 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.
Revised: 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.

Original: 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
Revised: 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

Original: PESNet strictly outperform PauliNet and DeepErwin on all systems.
Revised: PESNet strictly outperform PauliNet and DeepErwin on all systems.

Original: Another potential issue arises due to the choice of deep learning framework the models are implemented in.
Revised: Another potential issue arises due to the choice of deep learning framework the models are implemented in.

Original: 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.
Revised: 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.

Original: G C ONVERGENCE
Revised: G C ONVERGENCE

Original: When choosing a method to solve the Schr ¨ odinger equation one has to pick between accuracy and speed.
Revised: When choosing a method to solve the Schr ¨ odinger equation one has to pick between accuracy and speed.

Original: For classical methods, this might be the difference between choosing DFT, CCSD(T), or FCI.
Revised: For classical methods, this might be the difference between choosing DFT, CCSD(T), or FCI.

Original: In the neural VMC setting, one way to reflect this is by choosing the number of training steps.
Revised: In the neural VMC setting, one way to reflect this is by choosing the number of training steps.

Original: To better investigate this, we present convergence graphs for H4, H +4 , H10, and N2 in Figure 10.
Revised: To better investigate this, we present convergence graphs for H4, H +4 , H10, and N2 in Figure 10.

Original: For cyclobutadiene, we can see the convergence of different configurations of PESNet in Figure 9.
Revised: For cyclobutadiene, we can see the convergence of different configurations of PESNet in Figure 9.

Original: One can see that our network converges quite quickly on hydrogen-based systems such as the hydrogen rectangle, H +4 , or the hydrogen chain.
Revised: One can see that our network converges quite quickly on hydrogen-based systems such as the hydrogen rectangle, H +4 , or the hydrogen chain.

Original: For these small systems, PESNet surpasses PauliNet and
Revised: For these small systems, PESNet surpasses PauliNet and

Original: DeepErwin accuracy in less than 8000 training steps, reducing the training times to 2.9h, 3.2h, and 8.7h, respectively.
Revised: DeepErwin accuracy in less than 8000 training steps, reducing the training times to 2.9h, 3.2h, and 8.7h, respectively.

Original: In less than 16000 steps, about 17.4h of training, PESNet surpasses FermiNet on the hydrogen chain.
Revised: In less than 16000 steps, about 17.4h of training, PESNet surpasses FermiNet on the hydrogen chain.

Original: For more complicated molecules such as N2 and cyclobutadiene, the training requires more iterations to converge.
Revised: For more complicated molecules such as N2 and cyclobutadiene, the training requires more iterations to converge.

Original: So, we may also expect the extrapolated numbers from Tablefor PauliNet to be an optimistic estimate.
Revised: So, we may also expect the extrapolated numbers from Tablefor PauliNet to be an optimistic estimate.