Sampling-free Inference for Ab-Initio Potential Energy Surface Networks
Authors: Nicholas Gao
Stephan Günnemann
Venue: ICLR 2023 poster (2023)
Keywords: Graph Neural Networks
Computational Physics
Self-Generative Learning
Machine Learning for Science
Online Learning
Self-Supervised Learning
Molecules

Abstract: Recently, it has been shown that neural networks not only approximate the ground-state wave functions of a single molecular system well but can also generalize to multiple geometries. While such generalization significantly speeds up training, each energy evaluation still requires Monte Carlo integration which limits the evaluation to a few geometries. In this work, we address the inference shortcomings by proposing the Potential learning from ab-initio Networks (PlaNet) framework, in which we simultaneously train a surrogate model in addition to the neural wave function. At inference time, the surrogate avoids expensive Monte-Carlo integration by directly estimating the energy, accelerating the process from hours to milliseconds. In this way, we can accurately model high-resolution multi-dimensional energy surfaces for larger systems that previously were unobtainable via neural wave functions. Finally, we explore an additional inductive bias by introducing physically-motivated restricted neural wave function models. We implement such a function with several additional improvements in the new PESNet++ model. In our experimental evaluation, PlaNet accelerates inference by 7 orders of magnitude for larger molecules like ethanol while preserving accuracy. Compared to previous energy surface networks, PESNet++ reduces energy errors by up to 74%.

Published as a conference paper at ICLR 2023
SAMPLING-FREE INFERENCE FOR AB-INITIO POTENTIAL ENERGY SURFACE NETWORKS
Nicholas Gao, Stephan Günnemann
Department of Computer Science & Munich Data Science Institute
Technical University of Munich, Germany
[n.gao, s guennemann ) um. ue
ABSTRACT
Recently, has been shown that neural networks not only approximate the ground- state wave functions of single molecular system well but can also generalize to multiple geometries. While such generalization significantly speeds up training, each energy evaluation still requires Monte Carlo integration which limits the eval- uation to a few geometries. this work, we address the inference shortcomings by proposing the Potential learning from ab-initio Networks (PlaNet) framework, in which we simultaneously train a surrogate model in addition to the neural wave function. At inference time, the surrogate avoids expensive Monte-Carke integration by directly estimating the energy, accelerating the process from hours to millisec- onds. In this way, energy surfaces for larger systems that previously were unobtainable via neural wave functions. Finally, we explore an additional inductive bias by introducing physically mottivlyy restricted neural wave function models. We implement such function with several additional improvements the new PESNet+ model. our experimental evaluation, PlaNet accelerates inference by 7 orders of magnitude for larger molecules like ethanol while preserving accuracy. Compared to previous energy surface networks, PESNet+ + reduces energy errors by up 74 %.
accurately model high-resolution multi-dimensional
INTRODUCTION
Solving the Schrödinger equation is key to assess- ing a molecular system's quantum mechanical (QM) properties. As analytical solutions are unavailable, one must rely on expensive approximate solutions. Such methods are called ab-initio if they do not rely on empirical data. Recently, neural networks suc- ceeded in such ab-initio caleulations within the vari- ational Monte-Carlo (VMC) framework (Carleo Troyer 2017). Although they lead accurate ener- gies, training such neural wave functions proved to be computationally intensive (Pfau al. 2020).
Training
MetaONN
WF
8.
R_
PESSet
VMC
V Liuer
Inference V(R.)
Surogiie
To reduce the computational burden, Gao & Gün- nemann (2022) proposed the potential energy Tace network (PESNet) to simultaneously solve many Schrödinger equations, i.c., for different spatial ar- rangements of the nuclei in RR They use GNN to reparametrize the wave function model based on the molecular structure. While training significantly faster, afterward one only obtains neural wave function that generalizes over a domain of geometries, but not the associated energy surface. Obtaining the energy of a geometry remains costly requiring Monte-Carlo integration with complexity scaling as ONN*) in the number of electrons N. This high inference time prohibits many applications for neural wave functions. For instance, geometry opti- mization, free energy calculation, potential energy surface scans, or molecular dynamics simulations typically involve hundreds of thousands of energy evaluations (Jensen 2010 Hojj al. 2021).
Figure I: PlaNet framework. During training. the noisy energies obtained during the VMC optimization to fit our surrogate. At inference time only query the surrogate avoid costly numerical integration.

---

Published as a conference paper at ICLR 2023
To address the inference shortcomings, we propose the Potential learning from ab-initio Networks (PlaNet) framework, where we utilize intermediate results from the PESNet optimization to train a surrogate graph neural network (GNN) illustrated in Figure I In optimizing PESNet, one must compute approximate energy values to evaluate the loss. We propose to use these noisy intermediate energies, which otherwise would he discarded, as training labels for the surrogate. After training, the surrogate accelerates inference by directly estimating energies bypassing Monte Carlo integration.
As a second contribution, we adapt neural wave functions for closed-shell systems, i.e., systems without unpaired electrons, by introducing restricted neural wave functions. Specifically, we use doubly-occupied orbitals as inductive bias like restricted Hartree-Fock theory (Szabo & Ostlund 2012). Together with several other architectural improvements, we present the improved PESNet+ +
In our experiments, +ESNet+ significantly improves energy estimates on challenging molecules such as the nitrogen dimer where PESNet++ reduces errors by 74 % compared to PESNet. When analyzing PlaNet we find it to accurately reproduce complex energy surfaces well within chemical accuracy across range of systems while accelerating inference by 7 orders of magnitude for larger molecules such as ethanol. To summarize our contributions:
PlaNet: an orders of magnitude faster inference method for PESNet(++), enabling explo-
ration of higher dimensional energy surfaces at no loss accuracy.
PESNet+ an improved neural wave function for multiple geometries setting new state-of- the-art results on several energy surfaces while training only a single model.
2 RELATED WORK
Neural wave functions. Variational Monte Carlo calculations rely on expressive function forms to approximate the ground-state wave function. While early single determinants with linear combinations of basis functions without any explicit electron-electron interaction were used (Slater 1929), later backflow (Feynman & Cohen 1956) and Jastrow factors (Jastrow 1955) directly introduced electron correlation effects. While the combination of both has shown success (Brown et al., 2007), Carleo & Troyer (2017) were the first to systematically exploit the expressiveness of neural networks as wave functions. While initial works targeted small systems (Kessler et al. 2021] Han et al.) 2019
Choo et al. 2020), FermiNet (Pfau al.) 2020) (Spencer al.) 2020), and PauliNet (Hermann
ct al. 2020) introduced scalable approaches. Building on their success, recent works proposed integration into diffusion Monte-Carlo (DMC) (Wilson et al. 202T Ren et al. (022). introduction of
pseudopotentials (Li et al.] (022aa). architectural improvements (Gerard et al. 2022), and extension to
periodic systems (Wilson al. [202]] Li et al.) [022b) Cassella al.J [022)] Finally, two approaches to scaling neural wave functions to multiple geometries have been explored. Scherbela et al. (2022)
proposed weight-sharing scheme across geometries to reduce the number of iterations per geometry, while Gao & Günnemann] (2022) directly reparametrize the wave function with an additional neural network eliminating the need for retraining.
Machine learning potentials. The use of machine learning models as surrogates for quantum mechanical calculations has rich history, e.g., one of the first force fields has been the Merck Molecular Force Field (MMFF94) (Halgren 1996). Later, kernel methods were used (Behler 1 Bartók et al. 2013 Christensen et al. 2020). while graph neural networks currently obtain ofate-of- the-art results (Schütt al. 2018] Gasteiger et al. 2019 2021). Despite significant progress in the field by proving universality for certain model classes (Thomas et al. 2018 Gasteiger et al.) (021). the need for data and label quality remain limiting factors. Thus, the line between ab-initio methods and ML potentials has blurred in recent years, either by the integration of QM calculations into ML
models (Qiao al.) 2020 202T), AML methods (Wengert et al. (20TT. by integrating ML models Kirkpatrick et al. 2021).
into QM calculations (Snyder et al. 2012
3 BACKGROUND
Notation For consistency with previous work, we largely follow the notation by Gao & Günnemann (2022). We use 'geometries of molecule" to refer to different spatial arrangements of the same set of atoms in R We use C' to denote the number of geometries, B for the number of electron configurations per geometry, N for the number of electrons, and M for the number of nuclei. For
2

---

Published as a conference paper at ICLR 2023
electrons, use r RC x3x3 for the full tensor, r IR Bxxx3 for batch of electrons of single geometry, r E F N x for a single set of electrons, and r E R for a single electron. Similarly, we use R Mx3, and R R2 denote the nuclei tensor, nuclei geometry, and single nucleus, respectively. We assume all coordinates in 3D) space have already been transformed into the equivariant coordinate frame from Gao & Günnemann 2022) and drop the explicit transformation. We denote the local energy tensor as E E RCXD and refer to single energy value by Ee,i. Finally, we use bracketed superscripts to index sequences, e.g., layers in a network (1) or training steps(*).
R TERC RCXM x3 R
Variational Monte Carlo. To obtain the ground-state energy for fixed molecular system, one has to solve the associated time-independent Schrödinger equation
H(²) = =E
where th RNx3 + R is the wave function, E' E R the energy and H is the Hamiltonian operator
Zm rrys
Zm Zn
H
ryl
Rm
Rm Rn
m>n=
V(r)
with V2 the Laplacian operator. and V(r) describing the kinetic energy and potential energy. respectively. In linear algebra. Equation (I) is an eigenvalue problem where we are interested in the lowest eigenvalue Fg which also called the ground-state energy. VMC, we accomplish this by iteratively updating the parameters 0 of trial wave function the to approximate the ground state do (McMillan) 1965). The final accuracy of VMC is mostly determined by the function class of vhg- Luckily, trial function must only fulfill two requirements, first, it must be anti-symmetric to electron permutations, i.e., (r)) = - sigu( (T) yo (r) for all permutation m. and integral must be finite. This enables to neural networks as wave functions (Carleo & Troyer 2017). At each step, seek to minimize the energy
(wa H (b) (1bolyh)) (oo)
J (r)))())) 1/0 (r)drr
J (r) Hrs) (g(r)dr
LVMC.
(3)
Following the standard procedure (Ceperley et al. 1977p, we numerically approximate the integral via
*g((r) importance sampling by interpreting f (3(rr as a probability distribution of electron configurations
For each electron configuration, then compute the so-called local energy
0² log [r)(
i log (r)
(r) (r)
ΣΣ
+V(r). (4)
Or
Ori
We then use these local energies to obtain the gradients of our wave function parameters via
VoLuMe F (FL Bagg (F2 Vo log [ro(
where approximate all expectations with Monte-Carlo samples obtained via Metropolis Hastings, ie., via a Monte Carlo Markov Chain (MCMC) starting from the positions the last iteration.
As the trial wave function the approaches the ground state to. local energies approach the ground-state energy Eg. and their variance zero. Furthermore, our energy estimate Eng (x)] is lower bounded by Eg because the eigenfunctions of are a basis of all functions with the aforementioned properties.
While noisy estimates of the energy Eyg EE (r)] (r)] typically has to draw ((10)) samples (Gao & Günnemann 2022 Ren et al. 2022 to accurately model the energy for molecular system.
are sufficient for training, at inference time, one
Since one mostly interested in comparing energies of different geometries in chemistry, Gao & Günnemann (2022) proposed PESNet to avoid solving many Schrödinger equations indepen- dently but rather exploit the correlation between them. They accomplished this generalization by decomposing the problem into two neural networks, a wave function model (WFModel) representing the wave function dhg, and problem-adapting GNN (MetaGNN) that updates the WFModels's parameters 0 based the geometry. Thus, each VMC step becomes two-step process where one
PESNet.
3

---

Published as a conference paper at ICLR 2023
first obtains the wave function parameters 0. for each geometry CE {1..... C} via the MetaGNN and then computes the local energies and gradients with the WFModel Vs. via VMC. So, instead of obtaining energy estimates for the same fixed structure, during optimization, we get local energy estimates E E x where Fe,; = (r(rr,,) for a batch of C geometries R E CxXxx
Despite the unified training, one still needs perform independent numerical integration per geometry to obtain energies for multiple geometries. As the cost per evaluation O(N4) grows with the number of electrons, evaluations are typically restricted tens of geometries for small molecules due to computational constraints. Our PlaNet framework addresses these inference limitations enabling the evaluation of thousands of geometries, even for larger molecules, see Section 6.2
PLANET LEARNING POTENTIALS FROM NOISY INTERMEDIATE ENERGIES
To avoid the numerical integration for each geometry at inference time, we propose reusing the intermediate noisy energies E E RCB from the optimization of PESNet(++) (++) to train surrogate model. This surrogate model only acts on the nuclei and outputs the energy directly, i.c., in contrast to the wave function, there is no dependence on the electrons. At inference time, we only need to query the surrogate, reducing the inference time from hours to milliseconds.
As we optimize the energy throughout the training, we must deal with the change in target labels over time. We avoid this by treating the training process as an online learning problem where we use samples only once in the order in which they were generated. This online learning biases the network towards newer lower (better) energies.
To systematically incorporate physical invariances of the energy towards the Euclidean group E(3), i.e. translation, rotation, and reflections, we use a GNN as the surrogate. Compared to traditional internal coordinate-based polynomials (Jiang & Guo 2013 Li et al., 2013), GNNs offer straight- forward generalization to larger molecules without manual feature engineering. Though, as purely distance-based GNNs incomplete (?), we use the angle-encoding DimeNer++ (Gasteiger et al. 2020).
In each training step, we train the surrogate with the current geometries R and energies (!)) as input and target, respectively, As loss we optimize the weighted root mean squared error (RMSE)
C
(0)
VR(R.O) (2)
2 sur
(1)
OC
where E((0 = 1
6 )e and V3= RM3 + R denotes the forward pass of DimeNer++ with parameters X. Since we train in an online fashion, we perform Naur many gradient steps at each iteration / to better utilize the available data.
E(
B V
EB (
o,
Accounting for label noise. Since we obtain the energies from the VMC optimization described in Section 3 the labels are subject to noise. In fact, the observed noise is often orders of magnitude larger than the target energy differences we want to Iearn. To avoid the model oscillating between noisy labels,
encourage temporal data aggregation by applying an exponential moving average (EMA) to the weights as traditionally done in surrogate models (Gasteiger et al.) 2019 2021 SShüit al. 2021). However, picking the decay factor for the EMA is non-trivial as large decay values may cause the model to oscillate while small decay values may cause slow convergence. To resolve this issue, we recognize that the mean absolute error (MAE) between the surrogate and the VMC estimates is lower bound by the mean absolute deviation (MAD) of the energy distribution. We exploit this to identify two regimes, an initial regime where the PlaNet training error is higher than the energy label error and one where it is lower. In the first regime, we choose relatively high decay factor for the moving average to adapt quickly, while we shrink the decay in the second regime to encourage temporal data aggregation. To this end, control the decay factor via
70 Tbasc + Tigh 1 (cli) < (D())
with 0 < Thase + Thigh < 1 and 1 being hyperparameters and 1 being the indicator function. Since the true MAD ()() at iteration & is unknown, we must estimate it. By the central limit
4

---

Published as a conference paper at ICLR 2023
theorem, assume that the energies approximately normally distributed with standard deviation g+)) E de thus, approximate the MAD as D(4)
ozz
+ESNET+ IMPROVING MULTI-GBOMETRY NEURAL WAVE FUNCTIONS
PESNet+ builds on the recently proposed PESNet, discussed in Section In this work, we adopt several improvements to enable higher-accuracy energy surfaces at a similar cost. We use restricted neural wave functions for closed-shell systems, switch from block-diagonal orbital matrices to dense ones, add permutation invariant component to the wave function, and propose coordinate transformation to enable larger geometric perturbations during training. Appendix includes a complete definition of the new PESNet++ model.
Dense orbitals. Pfau et al. 2020); Hermann et al. (2020) and Gao & Günneeann 2022 define the wave function as linear combination of products of determinants of spin-up and spin-down orbitals
K
K
(ott ebb th E CRN"
(r)
Σ
wi det @ht det se
Σ
wk det
where a E {1, 4} indicates the spin and E R are learnable weights. Equivalently, may see the product of determinants as a determinant of an x N block-diagonal matrix. Instead of restricting the orbital matrix be block dock-diagonal, we adopt the improvements discovered by (Pfau al. 2020 in the official FermiNet repository and define a dense orbital matrix as
K
(T) =
det Φh xf
E ERNN Lka NRN
This is equivalent to picking N orbital functions for each spin instead of No ones for each spin Lin et al. )202I); Ren et al. (2022) and Gerard et al. (2022) have shown this to improve variational energies.
Restricted Neural Wave Funetions. The Fermi-Dirac statistic only applies to electrons of the same spin requiring to separate between spin-up t and spin-down +. electrons. But, most systems of interest, e.g., stable molecules, are closed-shell, i.e., systems without unpaired electrons, all orbitals will be doubly occupied (Szabo & Ostlund| 2012). We adopt this closed-shell formulation as inductive bias by enforcing identical orbital functions for both spins. One may see this as a neural analog restricted Hartree-Foor Fock (RHF), where one uses the same orbital functions for both spins, whereas previous work is related to unrestricted Hartree-Fock (UHF) with independent orbitals for both spins (Szabo & Ostlund| 2012). While UHF traditionally results in lower variational energies, we found this inductive bias to be better suited for neural wave functions, see Section G potentially due to its easier optimization by parameter sharing.
Speaking in symmetries, this restriction introduces invariance to the exchange of all spin-up and spin- down electrons. As both spin states are physically identical, all observables and, thus, the amplitude of the wave function stays the same under the exchange of spin states, i.e., [ r+]l = |+(rt, r. r-)]-
We implement this restriction by reformulating the interaction layers to only depend on the spin equality between two particles rather than their absolute spins. To incorporate this, we reformulate the update rules of our interaction layer to
h+ =a
Wi single
Σ
ghg
Σ
gij
+ b'inglc + Wi global
Σ jEAa
h
Σ h³ jEAn:
(10)
jEA
JEAS
1+1 = a (W a== I by. bl,=ay) Di
(11)
where h; and Pij are single and pairwise electron features, Oi is the spin of the i-the electron, ā denotes the opposing spin of a, and A is the index set of all a spin electrons. Subscripts a a= indicate different weights depending
whether the spins and B are identical. While Wilson jet al. (202T) and Gerard et al. (2022) studied similar update functions, we are the first to explore the restricted chosed-shell setting where one restricts the orbital functions to be identical between spins.
5

---

Published as a conference paper at ICLR 2023
Finally, it remains to reformulate the orbital construction
-otht
Toktt LoRLT
@att
(12)
kag Pi
(mw) (w (lwee)
) + b b4.0==
Σ
ka=8 am
exp(
0- im
ka=
||,
R
(13)
We initialize all w o with zeros to ensure identical initial behavior to block -iigonaa orbitals.
JJastrow factor. To case optimization, we follow classical quantum mechanical methods and introduce a permutation invariant part the wave function. the so-called Jastrow factor (Jastrow 1955 Brown
et al. 2007 Herrman et al. 2020). We implement the Jastrow factor J R R as permutation
invariant function on the final electron embeddings h LWw);
((r) exp (J (n(ww))
We det o²
J(h) E MLP(h.)
(14)
Avoiding dead neurons. In recent years, the flow of information in neural networks has been heavily studied (Brock et al. 2022; 2021 Hendrycks & Gimpel 2020). While such works have shown great success classical deep learning tasks such as computer vision, they have not yet been applied to neural wave functions. We implement several such improvements: careful parameter initialization, rescaling after activation functions (Brock et al. 2022). and the SiLU activation function (Hendrycks) & Cimpel| 2020). As all features are learned based on Euclidean distances, we still use tanh after the first layer to limit the magnitude of feature embeddings. We observed significant improvements in performance and trace it back significantly reduced number of dead neurons, i.e., the number of neurons whose value is independent of the input. We illustrate these observations in Appendix B
Coordinate transform. In machine learning, one typically draws batch of samples independently from some data distribution. Unfortunately, we cannot do the same for molccular geometries. Recall from Section 3
MCMC sample the electron positions, we have keep track of the last electron positions for each geometry to sample new ones. Thus, instead of i.i.d. sampling, Gao & Günnemann (022) obtain the next batch of geometries by applying small perturbations to the previous one. While performing sufficiently well for low-dimensional energy surfaces, such small changes do scale to gigeee-dimensional problems where the search space grows exponentially. To close the gap to the i.i.d. setting, we increase the magnitude of perturbations. However, this comes with the danger of decorrelating the electron positions from the nuclei. We minimize such decorrelation effects by transforming each electron in accord with its closest nucleus. This displacement avoids electron configurations from reaching dead regions of the probability distribution by guaranteeing that the distance between electron and its closest nucleus may only decrease. Formally, if R' Mx3
the new nuclei positions we obtain the new electron positions r' via A R x IR M Mx3 x gM + g
X(r[R', R) = r + (R' R.):
(15)
arg min Rm 1122-
LIMITATIONS
Firstly, while PlaNet enables nigh-resolution multi multi-diimensonal energy surfaces, scaling to higher- dimensional energy surfaces remains a challenge for all computational methods. Due the curse of dimensionality, high dimensional energy surfaces will only be sparsely sampled, e.g., in Appendix G.5 we find PlaNet exceeding the threshold of chemical accuracy once modeling the full six-dimensional energy surface of H2-HF To lessen this issue, could look into systematically reusing previous labels, extending the training time for higher dimensional energy surfaces, or adopting equivariant surrogate architectures these have shown be more sample efficient (Batzner et al. 2022).
Secondly, DimeNet*+ cannot distinguish some molecules. As this work seeks to present the general PlaNet framework, the choice of surrogate architecture is not essential part of this work. We encourage the adoption of universal models later works (Thomas et al. 2018, Gasteiger et al. 2021.

---

Published as a conference paper at ICLR 2023
Table I: MAE in mEn of
Table 2: MAEI in mći between
0.40
VMC energy surfaces com-
VMC and PlaNet over several energy surfaces. SE indicates the standard error of VMC ener- gies mEh, Brackets indicate the standard deviation at the last digit(s) across trainings.
with transform
pared to experimental results on N2 (Le Roy et al., 2006). PESNet+ reduces the error by
w/o transform
0.30
0.25
74 %.
0.15
MAE +
20000
40000
60000
MAE 4
SE
Iteration (
PESNet
5,36
SiLU
4.90
H4
0.0089(5)
0.004
Figure 2: Energy standard de- viation T00 during optimiza-
Li2
0.051(11)
0.019
+ zero bias
4.23
H10
0.152(25)
0.018
Jastrow
3.29
tion with and without coordinate transform. For each setting
N2
0.26(5)
0.186
+ Dense orbitals
2.60
H2-HF
0.324(14)
0.14
Restricted
1.39
C2HOOH
0.298(26)
0.33
rums are plotted.
Lastly, the restricted neural wave function formulation only suitable for molecules where the ground state a singlet state. For molecules with multiple unpaired electrons, for instance the O2 ground-state, or atomic systems, restricted neural wave functions cannot properly describe the ground state and, thus, lead to higher energies.
6 EXPERIMENTS
We split our experiments into two parts. Firstly, we present exhaustive ablation studies the contributions from Sections4ande Secondly, we analyze PlaNet's scaling by applying several potential energy surfaces, including the challenging nitrogen dimer, two-dimensional energy surface, and a high-resolution energy surface the larger ethanol molecule.
We denote the sampling-free inference energy estimates as PlaNet and the VMC energies by PES- Net++. Note that VMC energies are guaranteed to be upper bounds to the true energy, while no such bound exists for the PlaNet energies. When evaluating PlaNet energies, PESNet+ should be seen target since PlaNet should surrogate the VMC integration. Appendix Dependii detaaiss the setup, and Appendix the optimization. Timings are given in Nvidia A A100 GPU hours.
6.1 ABLATION STUDIES
PESNet++. Firstly, we evaluate the proposed architectural changes of PESNet++ (Section 5) on the nitrogen molecule in Table We picked the nitrogen molecule as it is known for strong electron correlation effects that result relatively high errors for most computational methods (Gdanitz) 1998 Pfau et al. 2020 Gao & Günnemann 2022). Evidently, all improvements gradually lower the MAE to the experimental results leading to a total error reduction of 74 %. Interestingly, restricting the wave function to doubly occupied orbitals results in the largest improvements, both in percentage (46 %) and absolute value (1.21 En). See Appendix for a training dynamics comparison.
Coordinate transform. To identify whether the in Section proposed proposed coordinate transform benefits training on multi-dimensional energy surfaces, we train multiple PESNet-+ with and without the coordinate transform, on the full six-dimensional energy surface of Hz-HF and observe the convergence of the standard deviation of the energy (t)) As the wave functions converge expect (+) to approach O. Figure 2 shows the convergence of several PESNet++ with and without the proposed coordinate transform. We find our transformation aiding in consistency and quality of convergence. Noticeably, without our coordinate transforms 3 out of 5 runs diverge during training. Compared to the converged runs, final standard deviations are on average 17 % lower, reducing the number of samples to obtain the same statistical enror by 31 %.
PlaNet energies. Lastly, we analyze PlaNet's reconstruction of VMC energies across a variety of energy surfaces, Appendix for the list of used geometries. The low errors in Table suggest very close fit between the PlaNet and VMC energies as all deltas are significantly below the threshold for chemical accuracy of 1.6 F.. Further, for larger systems, the fit is indistinguishable from
7

---

Published as a conference paper at ICLR 2023
12.5
PlaNet PRSNat
Gerard et al.
1 -Erpstionn is ma..
10.0
DMC
5.0
7.3
FarmiNer
Experiment
3.0
1.0
PlaNet CCSD(TT
0.0
-1.0
5
120
Bond length dp
Tossion angle
Figure 3: Potential energy of N₂ Pfau et al.
Figure 4: Potential energy curve of the CHş tor- sion angle of Trans-ethanol (Nandi al. (022). Error bars indicate the standard error.
2020
Gao & Günnemann 2022 Le Roy et al.
2006
Gerard al. 2022 Ren al. 2022).
CCSD(T)
PlaNet
ECCED
Brasa
5.0
4.5
10
20
0.0
.20
207
0.2
0.4
3.0
45
90
135
1800
90
135
180.0
90
1.55
180
6g
0g
9g
Figure 5: Two-dimensional slice of the potential energy surface of H2-HF. The remaining four dimensions fixed the equilibrium structure. Energies zeroed at the equilibrium structure and in m F. PlaNet agrees well ( 0.28(9) m Fg) with the CCSD(T)-based fit (Yang et al. 2018).
the 'ground-truth' VMC energies the VMC error per geometry increases. Appendix presents additional ablation studies on PlaNet hyperparameters and Appendix E an analysis of relative errors.
6.2 POTENTIAL ENERGY SURFACES
Evaluating ab-initio methods remains complicated issue as true energy values are rarely known. Furthermore, the threshold for chemical accuracy at kcal mol xs 1.6 is often relatively small compared the absolute magnitude of energies. However, as energy differences in chemistry
much more important than their absolute value, we follow the standard procedure (Pfau et al. 2020
Hermann al. 2020 Gerard et al. 2022 Gao & Günnemann 2022) and compare relative energy surfaces. As reférences, we use CCSD(T) or experimental results. Additional experiments on the lithium dimer (G.I), the hydrogen chain G.2
cyclobutadiene GG3). different ethanol states (G.4).
and higher-dingg dimensional energy surfaces of HF G.5)) are available in Appendix G
Heavily correlated systems. The nitrogen dimer is known be challenging system due to signifi- cant electron correlation effects (Le Roy et al. (2006), causing even the "gold-standard" CCSD(T) method to fail (Lyakh al. 2012). Nonetheless, FermiNet presented very accurate results comparable to the factorially scaling r 12-MR-ACPP method (Pfau et al. 2020 Gdanitz 1998). The comparison in Figure shows Figure that PESNet+ estimates the lowest variational energies. As for the relative error
A Fmaa
A Emin. PESNet+ outperforms Gerard et al. (2022)'s FermiNet extension and ex expensive DMC methods with relative error of m compared to 3.2 m F,, for DMC (Ren et al. 2022) and
Gerard et al. (2022), despite training only single model for fraction of the time. Finally, PlaNet's
cquilibrium distance 2.0747 ag is very close to the experimental results at 2.0743 ap
(Le Roy et al.
2006). For details finding the equilibrium geometries Appendix K
Scaling to larger systems. To explore the generalization towards larger systems, we investigate the potential energy surface of Trans-ethand along the CH3 torsion angle. In Figure 4 we compare PESNet+ and PlaNet estimates with accurate CCSD(T) results from Nandi et al. (2022). PlaNet agrees with the 99% confidence interval of all PESNet+ energies. Further, PlaNet almost perfectly
8

---

Published as a conference paper at ICLR 2023
recovers the CCSD(T) results with maximum disagreement of m.E, while only requiring s 0.11 ms per inference after training compared to the = 2.7 h per inference required for VMC integration. Additionally. the PlaNet energies are more robust and less noisy than the VMC energies obtained via numerical integration. While this discrepancy can be closed by increasing the number of samples during Monte-Carlo integration, reducing the VMC error bars by order magnitude requires ~ 270 h per geometry. This consistency is an encouraging sign that PlaNet-based interpolation has significant advantages in aulli-geometry VMC over traditional post hoc interpolation methods (Nandi et al. 2022, Yang et al., 2018) Jiang & Guo, 2013). In Appendix G.4, we explore the same CH3 torsion angle energy surfaces for different ethanol states.
Two-dimensional energy surface. Finally, we are interested in PlaNet's scaling to munti-diimensioall energy surfaces. We chose two-dimensional slice of the H2-HF interaction energy surface with CCSD(T) reference calculations (Yang et al. 2018) as described in Appendix [F] Figure 5 depicts a comparison between the two energy surfaces. PlaNet accurately reproduces this high-resolution two-dimensional energy surface with an MAE of x 0. (.28(9) mEn, leading to T-shaped equilibrium structure at R = 5.34 By and 0g = 90º consistent with previous results (Berhholtt et al., 1986 Yang et al., 2018; Guillon et al., 2008). Note that such high-resolution energy surface would not have been obtainable with previous neural wave function methods, see Appendix 6.3] In Appendix G.5 analyze the fitting of PlaNet to higher-dimensional slices of the energy surface.
6.3 TRAINING AND INFERENCE TIMES
In the previous experiments, we have shown that PlaNet approximates the underlying Monte Carlo integration well with little to no additional errors. In this section, we put the training and inference times of PlaNet's 's surrogate in perspec- tive to the standard VMC procedure. In Table we list the required times for training + (VMC training), performing Monte Carlo inte- gration (VMC inference), training the surrogate, and performing surrogate inference. While the training times refer training the model to whole energy surface, the inference times are given per geometry. Note that we exclude pre- training and initial MCMC thermalization in training since these contribute only small frac- tion of the total time. Evidently, training the surrogate is equivalent to between <1 and 3 VMC inferences worth of time, while enabling to 9 orders of magnitude faster inference. Further, surro- gate training accounts for <1 % of the total training time. These time advantages are growing with system size, e.g., for ethanol the surrogate training accounts for <0.2 % of the total training time.
Table 3: VMC and surrogate training and inference times per system. Inferences times are given per geometry. The training of the surrogate accounts for <1 % total training times while accelerating inference by 6 to 9 orders of magnitude.
VMC
Surrogate
Training
Inference
Training
Inference
H₂ Liz
32 h
min
20 min
4.7 us
40
13.6 min
18.5 min
U.75
Hoo
81 h
26.8 min
47 min
188
11-IIFF
107 h
38.4min
20 min
4.7
No
123
52.7
18.5 min
U.75
C2H,OH
478h
2.7h
43 min
118
CONCLUSION
We presented the PlaNet framework for sampling-free inference of potential energy surface networks by training additional surrogate on intermediate noisy variables from the VMC optimization. We addressed the changing energies during training by viewing the problem as an online learning task. Moreover, we identified two training regimes depending on PlaNet's accuracy and energy distribution. These regimes enabled us to perform efficient temporal data aggregation, reducing the impact of noisy labels. In addition, we explored doubly-occupied orbitals for neural wave functions by introducing PESNet++. In our evaluation, we observed that: (1) PESNet+ results in new Ktate-ffthhe-ne VMC results on several molecules, and (2) PlaNet suffers no loss in accuracy while decreasing inference times from hours to milliseconds. Thanks to these improvements in inference and accuracy, accurate high-resolution muulti-dimensional energy surfaces with neural wave functions obtainable. These are promising indicators for the future application of neural wave function methods for tasks that typically require hundreds of thousands or millions of energy evaluations, like analyzing multi-dimensional energy surfaces molecular dynamics simulations.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank Daniel Zügner, Filippo Guerranti, and Jan Schuchardt for their invaluable feedback on the manuscript, and Marten Lienen for our discussion on visualizing torsion angles.
Funded by the Federal Ministry of Education and Research (BMBF) and the Free State of Bavaria under the Excellence Strategy of the Federal Government and the Länder.
REPRODUCIBILITY
Our source is publicly available |licensed under the Hippocratic license (Ehmke 2019 Further, Ap- pendix pendix[A]detailee detaiils the PESNet++ architecture, Appendix [te optimization algorithm, and Appendix the of the experimental setup including all hyperparameters.
REFERENCES
Albert P. Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B, 87(18):184115, May 2013. ISSN 1098-0121, 1550-235X. doi: 10.1103/hhyseeV..87 184115.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022.
ISSN 2041-1723, doi: 10. 10.1038//41147---222999999
Jörg Behler. Atom-centeren symmetry functions for constructing high-dimensional neural network potentials. The Journal Chemical Physics, 134(7):074106, February 2011. ISSN 0021-9606, 1089-7690. doi:
David E. Bernholdt, Shi-yi Liu, and Clifford E. Dykstra. theoretical study of the structure, bonding, and vibrational frequency shifts of the H -HF complex. The Journal of Chemical Physics, 85(9):
5120-5127, November 1986. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.451700
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnn and Qíao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
Andrew Brock, Soham De, and Samuel L. Smith. Characterizing signal propagation to close the per- formance gap in unnormalized ResNets. In International Conference on Learning Representations, February 2022.
Andy Brock, Soham De. Samuel L. Smith, and Karen Simonyan. High-Performance Large-Scale Image Recognition Without Normalization. In Proceedings of the 38th International Conference on Machine Learning, pp. 1059-1071. PMLR, July 2021.
M. D. Brown, J. R. Trail, P. López Ríos, and R. J. Needs. Energies of the first row atoms from quantum Monte Carlo. The Journal of Chemical Physics, 126(22):224110, June 2007. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.2743972.
Giuseppe Carleo and Matthias Troyer. Solving the Quantum Many-Body Problem with Artificial Neural Networks. Science, 35566325)60026666, February 2017. ISSN 0036-8075, 1095-9203. doi: 10.1 126/scienceeaag2302..
G. Cassella, H. Sutterud, Azadi, N. Drummond, D. Pfau, J. S. Spencer, and W. M. C. Foulkes. Discovering Quantum Phase Transitions with Fermionic Neural Networks. arXi::2202.05183 [cond-mat, physics:physics), February 2022.
D. Ceperley, G. V. Chester, and M. H. Kalos. Monte Carlo simulation of many- study. Physical Review B, 1609),,81-3309, October 1977. doi: 10.1103/hhyseev..66.3000
http://wwww...ccrrru.ddddd/////s
10

---

Published as a conference paper at ICLR 2023
Kenny Choo, Antonio Mezzacapo, and Giuseppe Carleo. Fermionic neural-network states for ab- initio electronic structure. Nature Communications, 11(1):2368, December 2020. ISSN 2041-1723. doi: 10.1038/541467-020-B53 15724-9.
Anders S. Christensen, Lars A. Bratholm, Felix A. Faber, and O. Anatole von Lilienfeld. FCHL revisited: Faster and more accurate quantum machine learning. The Journal of chemical physics, 152(4):044107, 2020.
Coraline Ada Ehmke. The Hippocratic License 2.1: An Ethical License for Open Source. http://firsdonnhhmm.de,, 2019.
R. P. Feynman and Michael Cohen. Energy Spectrum of the Excitations in Liquid Helium. Physical Review, 189-)144,, June 1956. ISSN 0031-899X. doi: D..100//hysRevv102.11 102.1189.
Nicholas Gao and Stephan Günnemann. Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions. In International Conference Learming Representations, April 2022.
Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional Message Passing for Molecu- lar Graphs. In International Conference on Learning Representations, September 2019.
Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and Uncertainty-Awe -ware Directional Message Passing for Kon-Equulibrium Molecules. In NeurlPS-W. December 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal Directional Graph Neural Networks for Molecules. In Advances in Neural Information Processing Systems, May 2021.
Robert J. Gdanitz. Accurately solving the electronic Schrödinger equation of atoms and molecules using explicitly correlated (r12-)MR-CII The ground state potential energy curve of N2. Chemical Physics Letters, 28365),253-261, February 1998. ISSN 0009-2614. doi: 10.1016/50009-2614077 01392-4.
Leon Gerard, Michael Scherbela, Philipp Marquetand, and Philipp Grohs. Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need? Advances in Neural Information Processing Systems, May 2022.
G. Guillon, T. Stoecklin, A. Voronin, and Ph. Halvick. Rotational relaxation of HF by collision with ortho- and para-H2 molecules. The Journal of Chemical Physics, 129(10):104308, September 2008. ISSN 0021-9606, 1089-7690. doi: 10.1063/1 19975199
Thomas A. Halgren. Merck molecular force field. Basis, form, scope, parameterization, and performance of MMFF94. Journal of Computational Chemistry, 17(5-6):400-51, 1996. ISSN 1096-987X. doi:
Jiequn Han, Linfeng Zhang, and Weinan Solving many-electron Schrödinger equation using deep neural networks. Journal of Computational Physics, 19999, December 2019. ISSN 0021-9991. doi: j0.10166j.cc.2211...e 108929.
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs). arXiv:1606.0841 [cs], July 2020.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-networkk solution of the electronic Schrödinger equation. Nature Chemistry, 12(10)899-897, October 2020. ISSN 1755-4330, 1755-4349. doi: 10. 10.10088/441557-0000554444
Johannes Hoja, Leonardo Medrano Sandonas, Brian Ersst, Alvaro Vazquez-Mayagoitta. Robert DiStasio Jr., and Alexandre Tkatchenko. QM7-X, a comprehensive dataset of quantun-mechanical properties spanning the chemical space of small organic molecules. Scientific Data, 8:43, February 2021. ISSN 2052-4463. doi: 10. 1911038//1155-02--0000
Robert Jastrow. Many-body problem with strong forces. Physical Review, 98(5):1479, 1955.
Jan H. Jensen. Molecular Modeling Basics. CRC Press, 2010.
11

---

Published as a conference paper at ICLR 2023
Bin Jiang and Hua Guo. Permutation invariant polynomial neural network approach to fitting potential energy surfaces. The Journal of Chemical Physics, 139(5)005412., August 2013. ISSN 1921-9606. doi: 10.1063//.481711
Jan Kessler, Francesco Calcavecchia, and Thomas D. Kühne. Artificial Neural Networks as Trial Wave Functions for Quantum Monte Carlo. Advanced Theory and Simulations, 4(4):2000269, 2021. ISSN 2513-0390, doi: 10. we..1002add.2000000
James Kirkpatrick, Brendan McMorrow, David H. Turban, Alexander L.. Gaunt, James S. Spencer, Alexander G. D. G. Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, Lara Román Castellanos, Stig Petersen, Alexander W. R. Nelson, Pushmeet Kohli, Paula Mori- Sánchez, Demis Hassabis, and Aron J. Cohen. Pushing the frontiers of density functionals by solving the fractional electron problem. Science, 374(6573): 1a.11///cienc..ab65555.
December 2021. doi:
Robert J. Le Roy, Yiye Huang, and Calvin Jary. An accurate analytic potential function for ground- state N2 from direcc--otential-it analysis of spectroscopic data. The Journal of Chemical Physics, 125(16):164310, October 2006. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.2354502
Jun Li, Bin Jiang, and Hua Guo. Permutation invariant polynomial neural network approach to fitting potential energy surfaces. II. Four-atom systems. The Journal of Chemical Physics, 139(20); 204103, November 2013. ISSN 0021-9606. doi: 10. 10.101 48832997.
Xiang Li, Cunwei Fan, Weiluo Ren, and Ji Chen. Fermionic neural network with effective potential. Physical Review Research, 4(1):013021, January 2022a. ISSN 2643-1564. doi: 10.1 10t.PPhysleeeeeeech...000.0.0
Xiang Li, Zhe Li, and Ji Chen. Ab initio calculation of real solids via neural network ansatz. arXiv:2203.15472 [cond-mat, physics:physicss, March 2022b.
Jeffmin Lin, Gil Goldshlager, and Lin Lin. Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation. /2X04:21/2.03491 [physics], December 2021.
Dmitry I. Lyakh, Monika Musial, Victor F. Lotrich, and Rodney J. Bartlett. Multireference Nature of Chemistry: The Coupled- Cluustrr View. Chemical Reviews, 112(1)) January 2012. ISSN 0009-2665, 1520-6890.
doi: 10.1021/e2201441..
Angelo M. Maniero and Paulo H. Acioli. Full configuration interaction pseudopotential determination of the ground-state potential energy curves of Li2 and LiH. International Journal of Quantum Chemistry, 103(5):111-717, 2005. ISSN 1097-461X. doi: 10. 2055.
W. L. McMillan, Ground State of Liquid He4. Physical Review, 138(2A):A442 A4511, April 1965.
doi: 10.
Mario Motta, David M. Ceperley, Garnet Kin-Lic Chan, John A. Gomez, Emanuel Gull, Sheng Guo, Carlos A. Jiménez-Hoyonn Tran Nguyen Lan, Jia Li, Fengjie Ma. Andrew J. Millis, Nikolay V. Prokof ev, Ushmish Ray, Gustavo Scuseria, Sandro Sorella, Edwin M. Stoudenmire, Qiming Sun, Igor S. Tupitsyn. Steven R. White, Dominika Zgid, Shiwei Zhang, and Simons Collaboration the Many-Electron Problem. Towards the Solution of the Many- Many-Electron Problem in Real Materials: Equation of State of the Hydrogen Chain with State-of-the-AAe Many-Bode Methods. Physical Review X, 7(3):03 1059, September 2017. ISSN 2160-3308. doi: 10. 2..1100/yyeeexx...0331100 7.0.3.,,
Apurba Nandi, Riccardo Conte, Chen Qu, Paul L. Houston, Qi Yu, and Joel M. Bowman. Quantum Calculations on a New CCSD(T) Machine-Learned Potential Energy Surface Reveal the Leaky Nature of Gas-Phase Trans and Gauche Ethanol Conformers. Journal of Chemical Theory and Computation, pp. acs.jctc.2c00766 August 2022. ISSN 1549-9618, 1549-9626. doi: jc0e.2000760.
Eric Neuscamman, C. J. Umrigar, and Gamet Kin-Lic Chan. Optimizing large parameter sets in variational quantum Monte Carlo. Physical Review B, 85(4):045103, January 2012. ISSN 1098-0121, 1550-235X. doi: 10. 10.
12

---

Published as a conference paper at ICLR 2023
David Pfau, James S. Spencer. Alexander G. D. G. Matthews, and W. M. C. Foulkes. Ab initio solution of the mlny-electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, September 2020. doi: 10.1 00.103/Phyy.eereeeccc..2..2.3.
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller III. OrbNet: Deep Learning for Quantum Chemistry Using Symmetry-Adapted Atomic- Orbital Features. The Journal of Chemical Physics, 153(12):124111 September 2020. ISSN
0021-9606, 1089-7690. doi:
Zhuoran Qiao, Anders S. Christensen, Frederick R. Manby, Matthew Welborn, Anima Anandkumar, and Thomas F. Miller III. UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry. 1rxXiv:2105.4666 [physics], May 2021.
Weiluo Ren, Weizhong Fu, and Ji Chen. Towards the ground state of molecules via diffusion Monte Carlo on neural networks. arXiv:2204. [physics], April 2022.
Michael Scherbela, Rafael Reisenhofer, Leon Gerard, Philipp Marquetand, and Philipp Grohs. Solving the electronic Schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks. Nature Computational Science, 2(5):331 1-411, May 2022. ISSN 2662-8457. doi: 10.1088/43388--22-0022222 00228-x.
K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24): 241722, June 2018. ISSN 0021-9606, 1089-7690. doi: 10. (/06//.55099777
Kristof Schütt, Oliver Unke, and Michael Gastegger, Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Proceedings of the 38th International Conference on Machine Learning, pp. 9377-9388. MMLR,July 2021.
J. C. Slater. The Theory of Complex Spectra. Physical Review, 34(10)::293-1322, November 1929. doi: 10.1103/hhy.Re.344.1999
John C. Snyder, Matthias Rupp. Katja Hansen, Klaus-Robert Müller, and Kieron Burke. Finding Density Functionals with Machine Learning. Physical Review Letters, 108(25)225302,, June 2012.
ISSN 0031-9007, 1079-7114, 14. doi: 10.1 10.1103/Phssevvl.et 10..23300..
James S. Spencer, David Pfau, Aleksandar Botev, and W. M. C. Foulkes, Better, Faster Fermionic Neural Networks. 3rd NeurIPS Workshop on Machine Learning and Physical Science, November 2020.
Attila Szabo and Neil S. Ostlund. Modern Quantum Chemistry; Introduction to Advanced Electronic Structure Theory. Courier Corporation, 2012.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and -qquuation- neural networks for 3D point clouds. arXiv:1802.08211 [cs], May 2018.
Simon Wengert, Gábor Csányi, Karsten Reuter, and Johannes T. Margraf. Data-efficient machine learning for molecular crystal structure prediction. Chemical Science, pp. 10. 10.1039..DD DOSC05566,, 2021. ISSN 2041-6520, 0041-653,. doi: 10. 10.1039/DSCC05663.
Max Wilson, Nicholas Gao, Filip Wudarski, Eleanor Rieffel, and Norm M. Tubman. Simula- tions of state-of-the-an fermionic neural network wave functions with diffusion Monte Carlo. arXiv:2/03.15 12570 [physics, physics:quant-pl), March 2021.
Max Wilson, Saverio Moroni, Markus Holzmann, Nicholas Gao, Filip Wudarski. Tejs Vegge, and Arghya Bhowmik. Wave function Ansatz (but Periodic) Networks and the Homogeneous Electron Gas. 22.222222222222 [cond-mat, physics:physics, physics:quant-aaa
February 2022.
Dongzheng Yang, Jing Huang, Junxiang Zuo, Xixi Hu, and Daiqian Xie. A full-dimensional potential energy surface and quantum dynamics of inclastic collision process for H2- 12.HF, The Journal of Chemical Physics, 148(18):18433) May 2018 ISSN 0021 0021-9606. doi: 10.1003/1.500038.
13

---

Published as a conference paper at ICLR 2023
Smot jo %
SIOIIaI 3o %
40
40
20
20
10
10
20k
2014
15k
16k
10k
Hogran
10k
6k
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1,5
(a) tanh + normal bias
(b) SiLU + zero bias
Figure 6: Distribution of standard deviation of neurons during training.
A PESNET+ ARCHITECTURE
+ consists of three key ingredients, the MetaGNN, the equivariant coordinate system, and the wave function model (WFModel). The WFModel is used within the VMC framework to generate gradients and find the ground-state wave function. The MetaGNN's goal is to adapt the WFModel to the geometry hand and the equivariant coordinate system enforces the physical symmetries of the energy. We directly adopt the MetaGNN and the equivariant coordinate system from (ao Günnemann (2022) and would like to redirect the reader the original work for more information. All architectural improvements PESNet+ affect the WFModel. The WFModel acting on an electron
configuration first constructs single-clectron h and pair-wise g features in permutation cquivariant way. We adopt the same construction as in (Gao Günnemann 2022
M Σ
h
MLP [[r; - RR) E, ||r R&ll + 2m), m=1
(16)
11) 9
((r: ((:-rr)) r3) E, ||: rl)l)
(17)
where E' 3x3 is our equivariant coordinate system and km are nuclei embeddings outputted by the MetaGNN. These features are then iteratively updated with our new update rules from Equation (TOP
and Equation 11 After Twr many layers, we use the final electron embeddings 1, (2(Z))) construct
the orbital matrices á with Equation (13). Finally, compute the final amplitude with the Jastrow
factor with Equation [4]
B DEAD NEURON ANALYSIS
To illustrate the impact of tanh an activation function and the normally distributed biases, Figure plots the standard deviation of PESNet's neurons throughout training. One may that the number of dead neurons increases during training to >40 %, reducing improvements reduce the number of dead neurons to <10 % as seen in Figure 6b) Note that still keep one tanh activation function in the embedding block to limit the magnitude of our embeddings if distances increase.
effective network size. Our
C OPTIMIZATION
Each PlaNet optimization step consists of two parts, the first is VMC optimization step with the additional coordinate transform, and the second the optimization of the surrogate.
In each VMC step, we first perturb the geometry from the previous iteration followed by the coordinate transformation from Equation (15). Next, we sample for each geometry e new electron positions via Eettrpoois-Hasstin and calculate the local energy Ec, for each electron configuration Tz,i (McMillan 1965). Given these energies, we compute the gradients and construct our updates via natural gradient descent with the conjugate gradient (CG)-method (Neuscamman et al. 2012). For
14

---

Published as a conference paper at ICLR 2023
Algorithm 1/ 1th optimization step
Input: R (*) r((0 ,0)0 (x))
Output: R. r.r. e.
R ~ (=(RR(
r = A(r(0 R.
b Equation
15
for for (1........
do
~~~ 1/2 6
MCMC
Ec. (0t)) (Ne,i)- H ogo (Teas)
D Pfaau al. 2020
Hermann et al.
2020
Sc.: Be,t ( Σi= Ec,s)
end for
V een C = En, 1d V ecce log Iiec (Fe,i)]
θ ( =n V
CG-method
(Gao Günnemann
2022
for i {l..... Nurr } do
X' + Adam W step on Laur
> Equation 66
end for Compute
Equation
X (x
(4) + (1--))x
further details the VMC step, would like refer the reader Pfau al. 2020) and Gao & Günnemann (022)). Next, we fit our surrogate model over Namr many steps to the local energies and, finally, temporally average the surrogate parameters via the moving average described in Section 4
The complete optimization algorithm is given in Algorithm]] I Note that we dropped the dependence on the i step and explicitly wrote down the loop over all geometries. Though, in practice one implements those as an additional batch dimension and performs these operations in parallel. To
reduce the dependence on the current batch, we smooth ((t) and D(*) when evaluating Equation with E.Mss.
D EXPERIMENTAL SETUP
All default hyperparameters are reported in T To avoid exhaustive hyperparameter searches, we the default hyperparameters for PESNet (Gao & Günnemann 2022) and DimeNer (Gasteiger ett al. 2020). Exceptions are the nitrogen dimer for which we increased the number of determinants 32. the H2-HF energy surfaces for which we increased the determinants to 32, and the number of geometry random walkers to 64. In contrast to PESNet (Gao & Günnemann 2022), we do not deploy any early stopping criteria for the CG- CG-method as we found the convergence benefit slightly if we always do the full 100 steps. Note that this has barely any impact on the optimization time as the early stopping criteria were rarely met.
We implemented PlaNet and PESNet+ + on top of the official JAX (Bradbury al. 2018 implementa- tion PESNet (Gao Günnemann] 2022) which distributed under the Hippocratic license (Ehmke) 2019). We ran all experiments machine with 16 AMD EPYC 7742 cores and single Nvidia A100 GPU.
E RELATIVE PLANET ERRORS
While we analyzed the quality of the potential fit the main body, the MAE is not perfect metric, as one is often interested in relative energies rather than their absolute value. To account for this, we present extended results in Table where we also list the relative MAE and the MAD at the last iteration. We define the relative MAE as
rel. MAE
1 N
Σ
(E
D') -EE -
(18)
15

---

Published as a conference paper at ICLR 2023
Table 4: Default hyperparameters.
Parameter
Value
Optimization
Local energy clipping
5.0
Optimization
Batch size
4096
Optimization
Iterations
60000
Optimization
#Geometry random walker
16
Optimization
Learning rate r
(1++/1000)
Natural Gradient
Damping
10-1o(0)
Natural Gradient
CG steps
100
WFModel
Nuclei embedding dim
64
WFModel
Single-stream width
256
WFModel
Double-stream width
32
WFModel
#U #Update layers #Determinants
4
WFModel
16
WFModel
#Jastrow layers
3
MetaGNN
#Message passings
2
MetaGNN
Embedding dim
64
MetaGNN
Message dim
32
MetaGNN
NsBE
7
MetaGNN
NRBF
6
MetaGNN
MLP depth
2
MCMC
Init proposal step size
0.02
MCMC
Steps between updates
40
Pretraining
Iterations
2000
Pretraining
Learning rate
0.003
Pretraining
Method
RHF
Pretraining
Basis set
STO-6G
Evaluation
#Samples
106
Evaluation
MCMC Steps
400
DimeNet*
Cutoff re
10.0ag
DimeNet
Basis embedding size
8
DimeNet
Interaction embedding size
64
DimeNet
Out embedding size
256
DimeNet +
#Layer after skip
2
DimeNet*
#Layer before skip
1
DimeNet
#Blocks
4
DimeNet
#Layer out
3
DimeNet
NRBF
6
DimeNert+
NsBr
7
PlaNet Optimization
Taasc
0.99
PlaNet Optimization
Thigh
0.0099
PlaNet Optimization
Learning rate
10
[1+4/10000)
PlaNet Optimization
Nuur
5
PlaNet Optimization
S
1 1.05
PlaNet Optimization
Optimizer
AdamW W
PlaNet Optimization
(t) and (()) EMA decay
0.999
16

---

Published as a conference paper at ICLR 2023
Table 5: Extended version of Table 2 energies. 'SE' indicates the standard error of the mean for VMC target energies. The 'MAD' row refers the MAD of the energy distribution at the end of training. Numbers in brackets indicate the standard deviation at the last digit(s) across 5 trainings.
MAE in méh between VMC and PlaNet absolute and relative
H
H
Li2
1
N2
H2-HF
C2HçOH
MAD
0,43
0.3
1.5
1.4
9.0
18.8
15.0
SE
0.004
0.004
0.019
0.018
0.186
0.14
0.33
abs. MAEL
1,04)
0.0089(5)
0.05I(11)
0.152(25)
0.26(5)
0.324(14)
0.298(26)
rel. MAE]
0.0101(4)
0.00442(22)
0.049(11)
0.13(4)
0.25(32)
0.219(32)
0.274(5)
H
H T2
0
R
02
T1
F
Figure 7: Internal coordinate description of the Hę- system.
where E, and E refer the VMC and the PlaNet estimate of the energy the ith configuration, and É and E' to their respective means.
An interesting result is the comparatively low relative error for systems with heavy nuclei, i.e., non-hydrogen nuclei. We suspect due to the exponential moving average, PlaNet is biased towards higher energies if the VMC energies converge slowly.
F H2-HF SYSTEM
Figure 7 describes the internal coordinates of the H2-HF system. We follow (Yang et al. 2018) in the definiiio of the equilibrium structure and it to r1 1.4051, 72 1.73496, Φ 01 90, 82- 180 and R = 5.4. For the two-dimensional slices, we fix all but 0 and R at the equilibrium structure.
G ADDITIONAL ENERGY SURFACES
2.0
10.0
PlaNet
RCIPP
PlaNes
Eiverrment ma
1.6
PRSNea+
CCSD/T) CBS
8.0
PESNoE*4
PRSSNe
Experimmnt
PESNNt
6.0
FermiNet
FermiNet
PauliNet
0.5
MIC1+++F112 CBS
2.0
0.5
2.0
12
14
1.0
1.5
2.0 length
255
3.0
3.5
Rond lmngths dig
Figure Potential energy curve of Lig (Pfau et al., 2020 Maniero & Acioli 2005). Lightly shaded regions indicate standard deviation across 5 surrogate fits. The inset shows the absolute energy surface. Compared to PESNet, PESNet+ reduces the maximum error by 0.9 m
Figure 9: Potential energy curve the hydrogen
chain (Pfau et al. 2020 Hermann et al. 2020
Gao Günneeann 2022 Motta et al. 2017) 1)
Lightly shaded regions represent the standard de- viation across 5 surrogate fits. PESNet+ im- proves upon PESNet by 0.27 m.2
17

---

Published as a conference paper at ICLR 2023
Ground
Transition
Barricr
154.58
20
154.60
Transition barrier in local
+
Total energy in En.
15
PESNet
-254.62
FermiNet
10
PauliNet
-154.64
CCSD(T)
MR-CC
1544.66
Experiment
-154.68
Comparison of + PESNet, FermiNet, and PauliNet on estimating the transition barrier of cyclobutadiene. The first two plots show the total energy of the ground and transition state while the right plot shows estimates of the transition barrier. All neural wave function methods estimate similar transition barriers the upper end of the experimental range. In terms of total energy. PESNet+ outperforms PauliNet by 80.5 m Eb, PESNet by 23.6 m Ep and FermiNet by 7.6 m
Figure 10:
G.1 LITHIUM DIMER
typical task in chemistry is identifying equilibrium structures. To quantify PlaNet's suitability for such tasks apply it to the lithium dimer. One see in Figure [ that PlaNet agrees well with experimental results and PESNet++. Further, PlaNet virtually indistinguishable from the VMC estimates. When searching for the equilibrium structure with PlaNet, find the minimum 5.041 Pg. i.e. only 0.01 a apart from the experimental results (Maniero & Acioli 2005). PESNet+ reduces the maximum error from 1.4 E. 0.5 mE (64 %)..
G.2 HYDROGEN CHAIN
The hydrogen chain is a classic benchmark geometry with rich history (Motta al. 2017). Thus. we are interested in how our PlaNet and + compare to other neural wave functions (Pfau 2020 Gao Günnemann) 2022). The potential energy curve plotted in Figure We find our PESNet+ to produce the, to date, lowest energies with improvement of 0.3 m En over PESNet and 0.6 mb over FermiNet. Further, the PlaNet energies only differ by E, from the VMC estimates.
al. 2020, Hermann
G.3 TRANSITION BARRIER OF CYCLOBUTADIENE
To compare the architectural improvements of PESNet+ + to existing neural wave functions, report total energies and transition barrier estimates in Figure 10 As in FermiNet (Spencer et al.) 2020) and PESNet (Gao & Günnemann 2022), increased the hidden dimension of the single electron features to 512 and use 32 instead of 16 determinants. We resolved the convergence difficulties reported by Gao & Günnemann (022)) by training on the continuous energy surface obtained by linearly interpolating between the ground and transition state instead of only training on the ground and transition state. From the results in Figure TO it apparent that PESNet++ improves variational energies significantly by setting new ffte-of--heeant variational energies that are average 7.6 m better than FermiNet's. Though, should be noted that compared to PESNet more compute resources have been used due to the additional cost of the dense orbitals and restricted neural wave function, see Appendix H
G.4 DIFFERENT ETHANOL STATES
In the main body, we covered the CH,, torsion angle of Trans-ethanol which is in perfect agreement with the CCSD(T) reference (Nandi al. 2022) calculations. Here, additionally explore the
18

---

Published as a conference paper at ICLR 2023
6.0
4.0
3.0
1.0
PlaNes PESNet
CUSD(T)
-1.0
60
120
160
240
300
960
(a) Gauche
6.0
5.0
3.0
1.0
PlaNes
PESNet
A.PES
-1.0
60
120
180
240
300
390
Tursion
(b) TS1
8- Bain in mdh
4.0
3.0
Net PESNeT* & PES
-1.0
60
120
180
240
300
360
Tomiion angle
(c) TS2
Figure 11: Potential energy curves of the CH₃ torsion angle of ethanol for different OH torsion angle arrangements.
19

---

Published as a conference paper at ICLR 2023
CCSD(T)
PlaNet 2D
4.5
R
4.0
30
100
80
3.5
3.0
PlaNet 3D
PlaNet 6D
10
4.5
10
20
4.0
40
30
100
3.5
3.0
45
135
180
45
90
135
180
0g in
01 km
Figure 12: Comparison between different PlaNets trained on two-dimensional slide (top right), diceensnnns slice (bottom left), and the full six-dimensional energy surface of H2-HF compared to the reference CCSD(T) energy surface (Yang et al. 2018).
Gauche conformer as well as the two transition states TS1 and TS2. All energy surfaces are plotted in Figure II] Note that while the reference calculations for Gauche in Figure Ta are CCSD(T) based, for both transition states we use reference data from an A-ML approach (Nandi et al. 2022). While all surfaces generally agree on the approximate shape of the energy surface, we notably identify differences between our PlaNet estimates and the reference data for the Gauche and TS1 states. For the Gauche conformer, we locate slightly different minimum and for TSI we find the 120 periodicity to be broken in the given geometries. Note that the broken symmetry is due to working on partially relaxed structures (Nandi et al. 2022).
In terms of relative error, observe maximum discrepancy of 0.63 mEh. 0.55 m/h and 0.21 m m for Gauche, TSI, and TS2, respectively.
G.5 HIGHER DIMENSIONAL ENERGY SURFACES
While the main body already contains results the two-dimensional energy surface of H2-HF, are often interested the full-dimensional energy surface. Though, obtaining sufficient samples from higher dimensional surfaces is challenge due to the curse of dimensionality. To how well our PlaNet framework generalizes to more dimensions, we train additional models on a three-dimensional slice and the full six-dimensional energy surface of H3-HF, For the three-dimensional training, we additionally include the 0₂ angle. Albeit training on additional dimensions, we stick to the comparison of the same vioddimnsiooal slice in the main body ensure comparability.
Figure 12 shows comparison between the CCSD(T) baseline, and the three differently trained PlaNet models. While the general shape similar between all models, additional dimensions cause the energy surface to deform at angles of 0 = 90" increasing the MAE from 0.28 mEn to 1.02(34)) mE, and 2.39(25) mE respectively. We hypothesize that this is due to the curse of dimensionality and the consequent scarcity of data. Note that while the CCSD(T) baseline required first the selection and evaluation of 35,000 individual data points and an additional function fit, PlaNet framework is able to approximately capture the energy surface with single training. Although we do not know the exact time requirements for the CCSD(T) calculations, we expect them to be in the range of few minutes per geometry modern CPU, i.e., if we take
noticeable that the
20

---

Published as a conference paper at ICLR 2023
1.0
PESNet
0.8
Zero bias
0.6
SILU
+ Jastzow
Dense orhitals
Restricted
0.2
20000
40000
60000
25
50
Training sleep
Time in
Figure 13: Training plots on the nitrogen dimer with different enhancements to PESNet. On the left, the standard deviation of the local energy (lower is better) is plotted on the y-axis, and the training steps on the x-axis. The right plot shares the same y-axis but uses the training time ax x-axis. is apparent that all enhancements of PESNet++ strictly enhance the efficiency of PESNet they provide lower errors in less time and steps.
Table 6: Ablation study the number of surrogate training steps Nam-
N
H
H₂
Li2
H10
N2
H2-HF
0.0124(4)
0.0084(8)
0.0819(13)
0.14(4)
0,57(21)
0.338(31)
5
0.0117(4)
0.0072(7)
0.0786(9)
0.14(7)
0.33(7)
0.333(35)
10
0.0115(4)
0.00551(20)
0.0758(12)
0.19(8)
0.32(6)
0.33(4)
optimistic guess of minute per evaluation, the total compute time for the energy surface least 24 days. In contrast, our PlaNet model has been trained in only 4.5 days with single Nvidia A A100 GPU.
H CONVERGENCE
To analyze the efficiency our proposed enhancements to PESNet, we investigate the convergence rate based terms of the number of steps well as time. The resulting training curves for the nitrogen molecule are plotted Figure [3] it can be seen that all enhancements strictly improve convergence and result in significantly lower errors. While especially the dense orbitals add non-trivial amount of additional computational requirement, they also result in significantly better energies. The convergence rate per time and per step is strictly improving for all of our proposed improvements.
1 ABLATION STUDIES
To evaluate the value of our training heuristics, we perform ablation studies on the number of surrogate training steps Nsur and the exponential moving average decay rate 7. For the first experiment, alter the N hyperparameter between 1, 5, and 10 and run each experiment 5 times. The results are listed in Table can be seen that performing more than one step typically improves the results across all systems. This effect is more pronounced for larger systems.
Table 7: Ablation study with fixed values for the decay 7 of the exponential moving average.
7
H
H4
Liz
Hio
N2
H2-HF
0.99
0.0117(4)
0.0076(6)
0.0786(9)
0.112(21)
0.57(15)
0.39(4)
0.9999
0.019(4)
0.167(30)
0.087(4)
0.55(7)
0.46(14)
0.26(5)
adaptive
0.0117(4)
0.0072(7)
0.0786(9)
0.14(7)
0.33(7)
0.333(35)
21

---

Published as a conference paper at ICLR 2023
1.0954 10
-0.00075
-0.00000
-0.00126
PlaNet
E in
-..01150
Expeciment
Equilibrium distance
0.001775
0.0012000
-0.00225
2.05
2,06
2.07
2,08
2.00
2.10
Bond length in
(a) Nitrogen dimer
1. 199
10
-0.00188
PlaNet
Experiment
-0.0019
Equiibriuu distunce
E in PA
0.0050
-0.0051
4.950
4.975
5,000
5.025
5,050
5,075
5.100
5.125
5.150
Bond longth in
(b) Lithium dimer
Figure 14: High-resolution energy surface around the equilibrium structure.
the second experiment, we fix the exponential moving average parameter 7 instead adapting The results in Table show that initial high decay factor results in significantly higher final errors due to the large influence the initial bad labels. While a low y works reasonably across all systems, able to obtain better results larger systems with adaptive scheduling.
In general, we find our PlaNet framework to be robust against different hyperparameter settings.
SYSTEMS
In the following, we list all geometries used throughout this paper:
Hydrogen rectangle H4: geometries taken from Pfau al. 2020).
Lithium dimer Lig: 32 evenly distributed distances between 3.5 ag and 14 ag-
Hydrogen chain Hin: geometries taken from Motta et al. (017
Nitrogen dimer N2: geometries taken from Pfau et al. (2020).
H2-HF: 64 regular grid points of the N-dimensional energy surface. The boundaries are chosen as: T1, T2 [1.2 ag, 1.8 1.) Rt E [3.0 ag, 8.0 0,, 02 ø [0", 180%]. Ethanol CHH 64 evenly distributed torsion angles between 0 and 360".
K EQUILIBRRUM STRUCTURES
To find equilibrium structures, we performed high-resolution energy surface scans to find the global minimum. As many structures result in identical energies when performing inference with float32 precision, switch float64 precision. Note that all models were still trained with float32 and just converted for finding the equilibrium structure. We visualize these high dimensional energy surfaces around the equilibrium geometries for the nitrogen and lithium dimer in Figure 14
22


Reviews:
[{"review_id": "ouadBKOtyk", "reviewer": "Reviewer_oynR", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for speeding up the learning of potential energy surfaces by improving PESNet (a neural wavefunction model parameterized by a graph neural network and trained by variational Monte Carlo). In each step, previous methods run Monte Carlo simulations to estimate the energy, and the paper proposes to use these estimated energies to train an energy surrogate GNN (a Dimenet++ model). In addition, several changes to the PESNet neural wave function are presented.", "summary_of_the_review": "An interesting and well written paper that takes a small but important step towards solving the Schrodinger equations in a more efficiently and precisely via ML.", "strength_and_weaknesses": "The paper is well written and nicely presented.\n\nThe idea of using an energy surrogate as part of the VMC is promising.\n\nThe proposed improvements to PESNet seem relevant and improve performance.\n\nWeaknesses:\nThe training of the surrogate energy model is done via samples (targets) which only asymptotically are correct, which is handled by training only on the newest samples. Because of label noise, a heuristic exponential moving average is also used. It could be interesting with some further theoretical insight into the convergence properties of this approach.\n\nI expect that most of the paper might be difficult to follow for the general ML audience.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "i8azb6VtBL", "reviewer": "Reviewer_CnTt", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to simulate the potential energy surfaces of atomic systems by accelerating the variational Monte-Carlo (VMC) approach. It builds on the PESNet by Gao and Gunnemann and contributes two major innovations: 1) it introduces PlaNet which trains a surrogate model to avoid the expensive Monte-Carlo integration; 2) it introduces PESNet++ which improves multiple aspects of the PESNet. It shows 7 orders of magnitude acceleration and reduces the error up to 74%.", "summary_of_the_review": "In summary, the paper makes two contributions: 1) a PlaNet that learns a PES from the noisy data of VMC; 2) a PESNet++ that improves PESNet and achieves SOTA in several molecules. I think the motivation of 1) needs to be further clarified and additional baseline may be required. Contribution 2) alone makes significant progress in the ab-initio modeling of molecules. The paper may be accepted based on contribution 2) but the paper needs to be significantly revised. I give a score of 5 for now, but I am open to increase my score after satisfying responses from the authors.", "strength_and_weaknesses": "Strength:\n- PlatNet introduces an online learning scheme that learns a PES using the noisy labels from the VMC optimization. This innovation is significant because it enables the use of the very noisy energy data from VMC.\n- Multiple innovations are introduced to improve the PESNet++, which results in a new SOTA in VMC for several molecules.\n\nWeakness:\n- The motivation of learning a surrogate PlaNet from VMC optimization is unclear. From my understanding, the authors first use VMC to explore the entire PES and use the noisy data generated in the process to train a surrogate PlatNet. Since the entire PES has already been explored with VMC, why do we need a surrogate PatNet to predict the PES again? One argument may be that PatNet can generalize to new configurations that are not simulated by VMC. In that case, the authors probably should compare with a baseline DimeNet++ that is trained on the clean PES data from VMC to demonstrate the advantage of learning from noisy VMC energy.\n- The generalization performance of PlaNet is not studied in this paper. In Figure 3, 4, the PESNet++ simulated configurations are extremely dense and PlatNet only interpolates between nearby points. \nThe discussion on the results in Table 3 is slightly misleading. It is unfair to compare the training/inference time of the surrogate model to VMC, because the latter requires the former to generate data. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "djUf3xvAsP", "reviewer": "Reviewer_jX7P", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed two improvements over the baseline model PESNet. The first one, PlaNet, use the local energy computed in each PESNET training step to train a surrogate model, which takes in the geometry of nucleus and output final energy, namely no dependence on electrons. At inference time, only the surrogate model is used to do prediction. Hence, the orders of magnitude for inference is faster than PESNet. The second one, PES++, includes several improvements to achieve better accuracy than PES with similar cost.", "summary_of_the_review": "This paper proposed two improvements over PESNet, which makes the inference of PESNet faster and more accurate. The paper is written in a good manner.", "strength_and_weaknesses": "Strengths: 1. PESNET++ can produce highly accurate energy results (compared to FermiNet).\n                  2. PlaNet can approximate PESNET++ results well, but avoid the QMC inference phase, and thus can produce PES results much more efficiently.\n\nWeakness: The model is still limited, not applicable for real-world PES, for instance, molecular dynamics.\n\nSome questions to the authors: 1. Does PlaNet still have the variational property? 2. The neural Jastrow factor is not considering cusp condition as the traditional ones. Is it by design or just for convenience?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2a4rfBgxxr", "reviewer": "Reviewer_7yt6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper builds on the PESNet model of NN-based wave functions for small polyatomic molecules [this reviewers recommends avoiding re-using names for model architectures to avoid compromising anonymity). The works proposes two innovations. One is learning a surrogate model of the VMC energies to avoid O(N^4) scaling in the MC integration of the VMC integration. The other is a series of updates based on common approximations and tricks in quantum mechanics ", "summary_of_the_review": "This paper reports exciting practical improvements to an ML-based model for generatic wave functions and energies of small molecules. The method seems to provide desired improvements in speed while retaining accuracy, but the innovations do not seem to be arising from ML insights and are mostly through use of a surrogate function for energies, which is a common strategy in chemistry and from inclusion of quantum mechanical approximations that are well stablished.", "strength_and_weaknesses": "Strengths \n\nThe approach seems fundamentally sound and has good performance. The surrogate model is rather accurate and accelerates inference with little loss of accuracy, whereas the approximations and model updates based on traditional quantum chemistry heuristics are gainful and results in good cost/accuracy tradeoffs.\n\nWeaknesses.\n\nNeither of the innovation of performance improvements are particularly innovative as far as ML is concerned and this seems more fit for a quantum chemistry scientific journal. The use of a surrogate force-field-like potential seems relatively trivial and does not present a large methodological or task-specific innovation. All the quantum chemistry tricks are effective, but again, not that innovative on the ML side. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster