Variational Monte Carlo on a Budget — Fine-tuning pre-trained Neural Wavefunctions
Authors: Michael Scherbela
Leon Gerard
Philipp Grohs
Venue: NeurIPS 2023 poster (2023)
Keywords: Computational Physics
Machine Learning for Science
Quantum Monte Carlo
Fermionic Neural Networks
TL;DR: We obtain high-accuracy solutions to the Schrödinger equation cheaply, by fine-tuning a wavefunction pre-trained on many molecules.

Abstract: Obtaining accurate solutions to the Schrödinger equation is the key challenge in computational quantum chemistry. 
Deep-learning-based Variational Monte Carlo (DL-VMC) has recently outperformed conventional approaches in terms of accuracy, but only at large computational cost.
Whereas in many domains models are trained once and subsequently applied for inference, accurate DL-VMC so far requires a full optimization for every new problem instance, consuming thousands of GPUhs even for small molecules.
We instead propose a DL-VMC model which has been pre-trained using self-supervised wavefunction optimization on a large and chemically diverse set of molecules. 
Applying this model to new molecules without any optimization, yields wavefunctions and absolute energies that outperform established methods such as CCSD(T)-2Z.
To obtain accurate relative energies, only few fine-tuning steps of this base model are required.
We accomplish this with a fully end-to-end machine-learned model, consisting of an improved geometry embedding architecture and an existing SE(3)-equivariant model to represent molecular orbitals. 
Combining this architecture with continuous sampling of geometries, we improve zero-shot accuracy by two orders of magnitude compared to the state of the art.
We extensively evaluate the accuracy, scalability and limitations of our base model on a wide variety of test systems.

Variational Monte Carlo on a Budget
Fine-tuning pre-trained Neural Wavefunctions
Michael Scherbela" University of Vienna
Leon Gerard
University of Vienna
mi michael. schheel.schorbblladmmviiee...
leon. w.w.. gean.gerardddniii...oo ac.
Philipp Grohs
University of Vienna
philipp grohs@univie..c.commm
Abstract
Obtaining accurate solutions the Schrödinger equation is the key challenge in computational quantum chemistry. Deep-learning based Variational Monte Carlo (DI--VMC) has recently outperformed conventional approaches in terms of accuracy, but only at large computational cost. Whereas in many domains models are trained once and subsequently applied for inference, accurate DL-VMC far requires a full optimization for every new problem instance, consuming thousands GPUhs even for small molecules. We instead propose a DL-VMC model which has been pre-trained using self-supervised wavefunction optimization on a large and chemically diverse set of molecules. Applying this model to new molecules without any optimization, yields wavefunctions and absolute energies that outperform established methods such as CCSD(T)-22 To obtain accurate relative energies, only few fine-tuning steps of this base model are required. We accomplish this with a fully end-to-end machine-learned model, consisting of improved geometry embedding architecture and existing SE(3)-equivariant model to represent molecular orbitals. Combining this architecture with continuous sampling of geometries, we improve zero-shot accuracy by two orders of magnitude compared to the state of the art. We extensively evaluate the accuracy, scalability and limitations of our base model wide variety of test systems.
Introduction
Solving the Schrödinger equation is of utmost importance for the prediction of quantum chemical properties in chemistry. The lime-independent Schrödinger equation in the Born- Oppenheimer approximation [T] for molecule with Name nuclei and na electrons an eigenvalue problem with Hamiltonian H:
Z1Zs
Z7
H = EV,
H -=-
(1)
Tsj
RIJ
T
(>)
TSJ
1.J
By R (R....., RN M. and Z (Zi..... I...... Z N
NN we denote the nuclear
positions and charges. The electron positions are denoted by (r1...., T . Tna) E max with n, spin-up electrons and n, down electrons. The inter-particle difference and distance vectors are written as R = R R,, RIJ RyJ,, Tu r R,, Tu |raz]. ráz and rig |r,jl with
T 1, L..... Nanc and i, 1, ...... Net. The eigenvalues E of Eq. represent
"quual contribution, author order random
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

the energy states of a molecule, whereas a special interest lies in finding the smallest eigenvalue Eg. called the ground-state energy. The corresponding dighdiimessooall wavefunction &! : R - R
can be found via the Rayleigh-Ritz principle, by minimizing
H (t) (p((r)
[((b)) ((())
E
(2)
Due to electrons being fermions, the solution must fulfill the anti-symmetry property, stating that the sign of the wavefunction must change for any permutation P of two electrons of the same spin: ()() = -b Having access to the solution 1ú, allows in principle a complete description of the considered molecule. Unfortunately, only for one electron systems there exists analytical solution and due to the curse of dimensionality, with increasing number of particles, obtaining an accurate approximation of the wavefunction becomes intractable already for medium-sized molecules. This is because many high-accuracy approximation methods scale poorly with nea For example CCSD(T) coupled cluster with its single. double-, and perturbative gold-standard reference in computational chemistry, but its computational cost scales as ((na²) (2). Deep-lcarning-bens based Variational Monte Carlo (DL-VMC) has emerged as a promising alternative solution. A single step scales as (naii)) and it has surpassed the accuracy of many conventional methods such as CCSD(T), when applied to small molecules (3). In DL- VM,, the wavefunction is represented by a neural network with trainable parameters 0 and optimized via Eq. ussing gradient based optimization. Since the expectation value of Eq. cannot be computed analytically, it is approximated by sampling the electron positions during optimization and evaluation with Monte Carlo methods like
apple-excitationss variant- is considered the
Related work FermiNet by Pfau et al. (3 and its variants have emerged as the leading architecture for DL-VMC in first quantization. It can reach highly accurate energies, but typically requires tens of thousands of optimization steps for convergence. Many improvements have been proposed to further increase accuracy (5 and accelerate convergence (8). Furthermore, DL-VMC has been extended to properties beyond energies [9-[]] and systems beyond molecules (2-16). Despite the favorable scaling of DL-VMC, computational cost is still high, even for small molecules, often requiring thousands of GPUhs to find & for a single small molecule (17]. This is because unlike typical machine learning applications which train an expensive model once, and subsequently achieve cheap inference- - in DL-V DL the minimization of Eq. 2 is typically done from scratch for every new system.
A promising line of research to scale-up expensive ab-initio solvers such as DL-VMC or CCSD(T) has been develop proxy methods, which can be trained on outputs of ab-initio methods to directly predict molecular properties [18 19] or wavefunctions (20(2) from the molecular geometry. While these proxy methods can often reproduce the underlying ab-initio method with high fidelity and scale to millions of atoms (22), they are fundamentally limited by the accuracy of their reference method and need many high-accuracy samples for training, reiterating the need for scaleable, high-accuracy reference methods.
To enable DL-VMC methods to efficiently compute wavefinncioos for many molecules, several methods have been proposed to amortize the cost of optimization, by learning a single wavefunction across multiple systems. This has been demonstrated to work for different geometries of a single molecule [10 23 24). and recently two approaches have been proposed to learn wavefunctions across entirely different molecules, each with their own limitations. Gao et al. [17] reparameterized the heuristics to determine orbital positions,
orbitals of a wavefunction, by using chemistry-inspin managing efficiently generalize wavefunctions across different geometries of a single molecule. However, when learning single wavefunction across different molecules, their results deteriorated and transfer to new molecules proved difficult. Scherbela et al. (25) do not require heuristic orbital positions, but instead use orbital descriptors of low-accuracy conventional method to parameterize DL-VMC orbitals. However, while their wavefunction ansatz successfully transfers new molecules, their method requires separate, iterative Hartree-Fock (HF) calculation for every new geometry.
This work presents the first end-to-end machine learning approach, which successfully learns a single wavefunction across many different molecules with high accuracy. Our contributions are:
Our contribution

---

A transferable neural wavefunctions, which requires neither heuristic orbital positions, nor iterative HF calculations. We achieve this by building on the architecture by Scherbela et al. [5 and an orbital prediction model by Unke et al. (20)
A simplified and improved electron embedding architecture, leveraging expressive nuclear features from our orbital prediction model and message passing step between nuclei.
A chemically diverse dataset with up to 100 molecules based on QM7-X [26). a data augmentation method based on normal mode distortions, and successful training of a neural wavefunction on these with continuously sampled geometries. Additionally, we improve the initialization of electrons around the molecule, reducing the computational overhead.
As a final result, an accurate neural wavefunction, which shows for the first time zero-shot capabilities, i.e. high-accuracy energy predictions without additional optimization steps, on new systems. In particular it achieves better absolute energies than well established, high-accurace gold-standard reference methods, such as CCSD(T) -Z on unseen systems without any finetuning (cf. Fig. 22).
Overview of the paper Sec. we outline our method and the procedure to optimize a trans- ferable wavefunction across molecules. In Sec. B we thoroughly test the accuracy of our obtained wavefunction, by analyzing absolute energies (Sec. 3.1.. relative energies (Sec. (.2). and the impact of design choices in ablation study (Sec. 3.3). Throughout this work, we compare against other high-aceuracy methods, in particular results obtained by state-of-the-art DL-VME methods and
CCSD(T). Finally we analyze the scalability and limitations of our base model, by applying it large-scale dataset in Sec. 3.4 before a discussion and outlook for future research in Sec. 4
Methods
Our approach is divided into two parts (ef. Fig. D: containing electron embedding, orbital embedding and Slater determinant. On the other hand method for geometry sampling based on cormal-mod distortions.
On the one hand, a wavefunction ansatz,
2.1 Our wavefunction ansatz
electron model
wavefunction model
nuclear
electron emb. net
GNN
nuclcar embedding
electron embeddings
Slater Det.
Yp
Hamil- tonian
E
22
nuc. feat.
orbital model
S.S.
orbital emb.
Forb seb
PhisNet
FC
SCc
orbital coefis
orbital embeddings
Z
Pylies
8
Hphasst
normal distribution
AR
trainable
energy
OR,OR,
hessian estimate
protrainad, then frozon
Figure 1: Overview of our approach: Wavefunction ansatz (top) and geometry sampling (bottom)
single forward-pase -aass for our wavefunction model
Nasi
Nus
det (a)()
(x
(oo eoth
f raagha i,
Rel
again be divided into three blocks: An electron model acting nuclear and electron coordinates, generating electron embeddings æ! an SE(()--quivariant orbital model acting only on nuclear coordinates, generating orbital embeddings nd and farb, a Slater determinant combining electron- and orbital-embeddiings and ensuring anti-symmetry of the wavefunction.

---

Message passing neural network Throughout this work we use message passing neural networks (MPNN), to operate on the graph of particles which are connected by edges containing information about their relative positions. The electron-electron electron-nuclent and nuclear-nuclear edges are embedded with a multi-layer perceptron (MIP)),
chol MLP((ri5,..) riyl) C
ett =
eIJ
mac-mue = MMLP((R,,,R)))) (4)
by using a concatenation ((-) of distance and difference vectors, and separate weights for each MLP. A single message passing step is decomposed into the following operations
a MessagePassing(a "" {aa (a {e{e))
(5)
o Linear(a;")) + Linear
Σ Linear(aj) gend)
o Linear(ess))
(6)
for a receiving particle a and the set of sending particles {agmdd, connected via their edges { By denote the non-linear activation and with the element-wise multiplication along the feature dimension. An MPNN is obtained by stacking MessagePassing layers
(a,,( { (eij]) = MessagePassingg MessagePassing (az), (ei)))
(7)
In the following we use these message passing steps to model all inter-particn interactions.
-Euuevvquvaariant orbital model The orbital model is simplified version of PhisNet 20 neural network predicting the overlap matrix S and the Fock matrix F, via nuclear embeddings a mmo x phisnet, RR Z)
SI = (0(
puc x²)
FIJ = So (x²
muc
(8)
Here xu E R. NR. ((Lx)x)(+xxNNxa))) of the predicted orbitals denoted by Nasis and the feature dimension of the nuclear embeddings by Nehannels. The full overlap- and Fock-matrices are assembled from the corresponding blocks T Each layer of PhisNet is SE33i_-quivarian,, ensuring that any 3D-rotation or inversion of the input coordinates R, leads to an equivalent rotation of its outputs. This done by splitting any feature vector into representations of varying harmonic degree | = 0, ..... L, each with components m -l... I. A detailed description of SE3-equivariant networks general, as well as PhisNet in particular can be found in (20). list of changes and simplifications we made to PhisNet can be found in Appendix C
and SIJ and FIJ are each in RR x Nusis The basis-set size
SJ and FIJ, leading to matrices of shape NNusss
The orbital embeddings (corresponding orbital expansion coefficients conventional quantum chemistry calculation) are obtained by solving the generalized eigenvalue problem:
FC = SC&ck
arb reshape Ck. [Nuae NNbssis)
It has been shown empirically that it is beneficial to obtain the orbital coefficients TiK solutions to this generalized eigenvalue problem, rather than predicting them directly [27] as functions of R and Z. This is because the orbital coefficients are neither unique, nor do they share the molecule's symmetry. The matrices and on the other hand are unique and transform equivariantly under Eisnsrransfoonnaiion of the molecule, leading to well defined learning problem.
oeb
we do not use the orbital energies c and obtain the backflow factors and exponents g. by first localizing the resulting orbitals using the Foster-Boys localization (28) (cf. Appendix B and subsequently using an MPNN and MLP acting on the orbital embeddings.
Following [25]
Non
orb lk
Σ
kn ar
orb Ins
orb 1
MPNN (K
oob Jk), {e7s (u
(10)
f porb MLP (2))))
forb I. New Not Nax Neab
(11)
91k MLP((.7)))
g pch E TR Nu x N
(12)
The electron embedding message passing neural network. To incorporate the geometric information of the molecule considered, we leverage the equivariant prediction of the PhisNet nuclear embeddings nue by first performing a message passing step between the nuclear embeddings
Electron model
nBC 1
= MLP(x:
nUC 1
MessagePassinggaat (i) { (e)))-))))))

---

and then using these features to initialize the electron embeddings
e1,0 MessagePassingg(O.. { " ), { (e))
by using zero vector for the initial receiving electrons. This differentiates our electron model from previously proposed methods (3 8), leading to a better generalization when optimized across molecules (cf. Sec. B.3). The final step the embedding a multi-iteration message passing between electron embeddings to capture the necessary electron-electron interaction
MPNN (a;;") (x::) { (ef)
resulting in Naan- -imb-dimmnioonal representation for each electron at E R-Venh, i i= .....,, Tel.. Tel-
Overall -quuiarrrree Like existing approaches (3 9 17 our overall wavefunction does not enforce E(3)-symmetry. This is because the wavefunction have lower symmetry than the molecule [33]. for example in the case of the excited states a hydrogen atom. We therfore choose all parts of the network that purely on nuclear coordinates (i.e. the PhisNet model and the energy/hessian estimate) to be equivariant under ((3)=rransoormations. All parts of the network acting on electron coordinates (in particular the electron model) break this symmetry by depending explicitly on the cartesian coordinates of the electrons. The architecture is thus only invariant under translations, but under rotations or inversions. We bias the model towards approximately invariant energies using data augmentation as discussed in Sec. 2.3
2.2 Sampling
Markov Chain Monte Carlo (MCMC) sampling of electron positions We use MCMC to draw samples from the probability distribution (rr)) to evaluate the expectation value of Eq. 2 One notable difference compared to other works our initialization r of the Markov Chain. the limit of infinite steps, the samples are distributed according to v/², but for finite number of steps the obtained samples strongly depend on r This issue is typically addressed by a "burn-in", where MCMC run for a fixed number of steps (without using the resulting samples) to ensure that has diffused to state of high probability. Previous work has initialized rd using a Gaussian distribution of the electrons around the nuclei. We find that this initialization is far from the desired distribution 2 and thus requires xs 10 MCMC steps to reach the equilibrium distribution. We instead initialize r by samples drawn from an exponential distribution around the nuclei, which much better approximates the correct distribution and thus equilibrates substantially faster (cf. Appendix (A). We find that exponential initialization reduces the required number of burn-in steps by ca. 50%, reducing the computational cost of a 500-step zero- shot evaluation by ca. 5%.
Normal mode sampling of geometries Since DL-VMC is an ab-initio method, we do not require labeled dataset of reference energies, but to obtain transferable wavefunction, which generalized well to new systems, diverse dataset of molecular geometries set of geometries R², update R on the fly, by perturbing each geometry every 20 optimization steps by adding random noise A R to the nuclear coordinates. Using uncorrelated, isotropic random noise for AR would yield many non- con-physccal geometries R', since the stiffness of different degrees of freedom can vary by orders of magnitude. Intuitively we want to make large perturbations along directions in which the energy changes slowly, and vice-versa. We achieve this by sampling AR from a correlated normal distribution
required. Starting with an initial
AR ~ N(a(R 3H phis
R'=R+AR.
(13)
The bias term ox( R R) ensures that geometries stay sufficiently close to their starting point R. The covariance matrix chosen proportional the pseudo-inverse of the hessian of the energy Ephis which is predicted from the scalar component of the nuclear embeddings using a pre-trained MLP. By using H phis as covariance matrix, we take large steps along soft directions and small steps along stiff directions, thus avoiding unphysical geometries with very high energies.
Nas MLP nuc
Ephis
Ehhis
Σ
H phis
I.J = 1 Nau 5.5= = 1...3 (14)
IC.JE
OR
5

---

After distorting the nuclear coordinates R. we also adjust the electron positions r, using the space- warp coordinate transform described in (29). which effectively shifts the electrons by weighted average of the shift of their neighbouring nuclei. In addition to this distortion of the molecule, we also apply a random global rotation to all coordinates, to obtain more diverse dataset.
2.3 Optimization
To obtain orbital descriptors (and energies calculate the hessian of the energy) we pre-train PhisNet against the Fock matrix, the overlap matrix, the energy and the forces of Hartree-Fock calculations in minimal basis set across 47k molecules. Further details of the loss function, dataset, and the adaptions to PhisNet can be found in Appendix(C] For all subsequent experiments, we freeze the parameters of PhisNet. full DL-VMC calculation obtain a ground-state energy prediction can be divided into three consecutive steps:
1. Supervised optimization using PhisNet: Initially, the neural-network orbitals (cf. Eq. are optimized to minimize the residual against orbitals obtained from PhisNet. It ensures that the initial wavefunction roughly resembles the true ground-state and is omitted when fine-tuning an already optimized base model.
2. Variational optimization:
Minimization of the energy (ef. Eq. 2 by drawing samples from 2 using MCMC and updating the wavefunction parameters using the KFAC optimizer (30).
3. Evaluation: For inference of the ground-state energy, we sample electron positions using
MCMC, and evaluate the energy using Eq. 2 without updating 0.
To train a multi- guntti-ge transferable wavefunction we further divide the variational optimization of a neural wavefunction into two steps:
1. Pre-training: A single wavefunction model is trained across many molecules and geometries. In every gradient step we only consider single geometry per batch. The next geometry to optimize is chosen based on the energy variance as proposed by [10]. To sample continuously the space of molecular geometries
distort each geometry every 20 optimization steps as described in Sec. 2.2 We refer to evaluations of this model on new systems as "zero-shot".
2. Fine-tuning: A small number of additional variational optimization steps is done using geometries of interest, starting from the weights of pre-trained base model. This procedure yields a model that is specialized to the molecule at hand and typically yields more accurate energies on the specific problem than the raw pre-trained model.
3 Results
We pre-train our wavefunction model on a dataset of 98 molecules (699 conformers) for 256k optimization steps using the architecture and training procedure outlined in Sec. 2 Below demonstrate the performance of this model, for zero-shot evaluations and after subsequent fine- tuning.
3.1 Accuracy of pre-trained model for absolute energies
To test the transfer capabilities of the model, evaluate on test-sets, which each contain randomly chosen and perturbed molecules, grouped by molecule size (measured as the number of non-Hydrogen atoms). To avoid train/test leakage, we excluded all molecules that are part of these test sets from the training set (cf. Appendix DD. Although the model has only been trained on molecules containing up to 4 heavy atoms, we evaluate its performance across the full range up to heavy atoms. We find that for molecules containing up to 6 heavy atoms our method outperforms CCSD(T) with 2Z basis set and outperforms CCSD(T) with a 4Z basis set after only 4k fine-tuning steps. This is a large improvement over the state of the art: Zero-shot evaluations by Gao et al. [17] did not manage to outperform Hartree-Fock Fock baseline, even on the toy system of Hydrogen-chains. Similarly, Scherbela et al. (25) achieve high accuracy after fine-tuning, but result that are worse than Hartree-Foct Fock a zero-shot setting. contrast, our improvements their method increase zero-shot accuracy by more than 2 orders of magnitude.

---

Zero-shot
b After 4k fine-tuning steps
Test set with heavy atoms
105
Our work (no pre-training)
Our work (pre-trained)
10
330x
Scherbela (202 3)
Gerard et (2022)
HF-CB5
HF-CBS
EccsOTT ccss
10a
22
CC:22
Paiorrmr
HF-C85
CC-22
02
0C:32
10a
0C:47
CC.32
CC.42
104
0
CC-C8S
CC.CBS
cC.c85
3
4
5
6
7
1k
2K 4k 8k 16k 32k
Nr. of heavy atoms
Nr. of heavy atoms
Optimization steps
Figure 2: Absolute energies: Energies relative to CCSD(T)-CBS (complete basis set limit) when re- using the pre-trained model on molecules of varying size without optimization (a) and after fine-tuning (b). (c) depicts energy for the test set containing 3 heavy atoms as function of optimization steps and compares against SOTA methods. Solid lines are with pre-training, dashed lines without. Gray lines correspond to conventional methods: Hartree-Fock in the complete basis set limit (HF-CBS), and CCSD(T) with correlation consistent basis basis sets of double to quadruple valence (CC-nZ).
We furthermore find that fine-tuning pre-trained wavefunctions can be more cost effective than well established conventional methods. For typical molecule containing 5 heavy atoms, we require 9.5 node-hours for 4k fine-tuning steps, achieving accuracy that surpasse CCSD(T)-442 In contrast, CCSD(T)-4Z requires 10.5 node-hours and 16k steps of PsiFormer (achieving the same accuracy), require 53 node-hours. All run-times are listed in Tab. 3 Appendix
To further test the accuracy of our method, evaluate for heavy atoms the performance with increasing fine-tuning steps (cf. Fig. ( We compare our work (with and without fine-tuning) against CCSD(T) and reference calculations done with state-of-the-art DL-VMC methods (6] 8 For to 16k optimization steps, our pre-trained model yields energy errors that are 1-2 orders of magnitude lower than other reference methods. After longer optimization the method by Gerard et al. [8] and PsiFormer (6) surpass our predictions. We hypothesize that this is not to blame on pre-training but that our orbital prediction framework, which allows us optimize across molecule, yields less expressive wavefunctions than fully trainable backflow. This is demonstrated by the fact that a pre-trained and non-pre- con-pre-trainee trained model converge the same energy in the limit of long optimization.
Like the absolute energies, also the variance of the local energies (another measure of wavefunction accuracy) is substantially improved by our method and results are depicted in Appendix E
3.2 Accuracy for relative energies
While Sec. 3.1 demonstrates high accuracy for absolute energies, we find that relative energies (being the small difference between two large absolute energies) can be unsatisfactory in the zero-shot regime, and small number of fine-tuning steps is required to reach quantitatively correct relative energies. Fig. demonstrates this issue on 4 distinct systems, each highlighting a different challenge for our model. For each system, we evaluate our pre-trained base model without any system specific optimization (zero-shot), and after 4000 fine-tuning steps. The fine-tuning optimization is done separately for each of the 4 systems, analogously to Sec. 2.3 yielding 4 distinct wavefunctions that each represent the ground-state wavefunctions of all considered geometries per system. Results after more fine-tuning steps can be found in Appendix
Bicyclobutane conformers Fig. 3 depicts the energies of 5 conformers of bicyclobutane relative to the energy of its initial structure. The system of interest, because CCSD(T) severely underestimates the energy of the dis-TS conformer by ~ 60 mHa [3.. While our zero-shot results yield the correct sign for the relative energies, they are quantitatively far off from the gold-standard DMC reference calculation (31). in particular for the dis-TS geometry, where overestimate the relative energy
7

---

bicyclobutane
Fecelitttuns /mHa
150
100
Our work: 0-shot
FermiNet: 200k
work: steps per geom.
FarmiNet: steps per goom.
50
DMC (Kinal et all)
CCSDIT) (Kinal et al.)
-50
TS
con_TS
N₁ dissocition
Propadiene global rotation
Propadiene torsion
10
Our work (0-shol)
Our work (0-shot)
Our work (0 (0-shol)
fine-tuning)
(0-shet avg.)
100
fine-tuning)
500
Gerard et al. (2022)
Our work (4k fine (nnntuning)
CCSDIT//CB
Ey mka
E- Eo niHa
400
Ea
300
40
100
0
1000
200
20 40
60
l bohr
angle
torsion deg
Figure 3: Challenging relative energies: Relative energies obtained with and without fine-tuning on 4 distinct, challenging systems, compared against high-aceuracy reference methods. a) Relative energy of bicyclobutane conformers the energy of bicyclobutane b) Potential energy surface (PES) of N2, c) global rotation of propadiene d) relative energy of twisted vs. untwisted propadiene.
by 90 mHa. However when fine-tuning our model for only 700 steps per geometry (4k total), we obtain relative energies that are in close agreement with DMC (max. deviation 2.1 mHa). Comparing to FermiNet (5) we find that our results are more accurate than a FermiNet calculation after 10k steps (requiring twice batch size; max. deviation 7.5 mHa), and slightly less accurate than FermiNet calculation optimized for 200k steps (max. deviation 1.4 mHa). As opposed to CCSD(T) model does not suffer from systematic errors, for the challenging dis-TS geometry. A table of all relative energies can be found in Appendix F
Nitrogen dissociation When evaluating the energy of an N2 molecule at various bond-lengths we obtain high zero-shot accuracy near the equilibrium geometry (d = 2.1 bohr; in training set), but substantially lower accuracy at smaller large bond lengths (ef. Fig. 3p). This is due the lack of dissociated atoms in the pre-training dataset and again mostly remedied by finetuning, In the most challenging regime around = 3.7, even with fine-tuning we obtain energies that are mHa above high-accuracy results obtained by Gerard et al. 8 indicating the need for longer fine-tuning.
Global rotation of propadiene Energies of molecules are invariant under global translation or rotation of all particle positions. The wavefunction however is not invariant and neither is our wavefunction ansatz, which can lead to different energies for rotated copies of a molecule. When evaluating the energy of our model on 20 copies of propadiene (C2H4) rotated around random axis, we find typical energy variations of mHa, but also a individual outliers, deviating by up to mHa. This highlights a dilemma facing all existing DL-VMC models: On the one hand, constraining the wavefunctions to be fully invariant under rotation (or even just invariant under symmetries of the Hartree-Fock orbitals) is too restrictive to express arbitrary ground-state wavefunctions [23] One the other hand, our approach of biasing the model towards augmentation, appears to be helpful but not to be fully sufficient. When evaluating earlier checkpoint of our base model (trained for 170k epochs instead of 256k), we observe mean energy variations of +3 mHa and outliers of up 30 mHa, indicating the positive impact of prolonged training with data augmentation. We again find 4000 fine-tuning steps split across all geometries are sufficient reduce errors below chemical accuracy of 1.6 mHa
energies by data
In addition augmenting data during training time, we can also augment the data during inference time, to obtain fully rotation- innvariant energies, from approximately rotation invariant model. We

---

achieve this by randomly rotating the molecule at every inference step and averaging across all rotated geometries. Since the Monte- Carlo estimate is already an average across many different samples anyways, this comes at no additional computational cost. Fig. 3 shows that with rotation averaging we obtain energies that are fully invariant, up to Monte-Carke noise.
Twisted propadiene Twisting one of the C=C bonds of propadiene leads to a transition state with an energy difference of 110 mHa. Evaluating the energy without fine-tuning on an equidistant grid of torsion angles, we obtain the correct barrier height (112 mHa), but also deviations of up to 40 mHa. During pre-training we purposefully sampled twisted molecules, but only included equilibrium geometries, transition geometries, and one intermediate twist (cf. Appendix Dj. This seems sufficient correct zero-shot barrier heights, but insufficient for high accuracy along the full path. Short fine-tuning (4k steps distributed across all 10 geometries) yields excellent agreement with CCSD(T): 0.2 mHa discrepancy for the barrier height and maximum deviation of 2 mHa along the path.
3.3
Ablation studies
Zero-shat accuracy
Accuracy after 4000 fine fine-tuning tutiinnnn steps
Scherbela et al. 128k pre-training
5620
improved
4250
architecture
(-76%)
(-35
nomal-mode distortions
(+10%)
(-12%)
-1470
-0.9 (-996]
training set
(-98%)
determinants
(+5%)
-0.7
258k pretraining steps
.1..5
-1.8
(375))
(-21%) %.
work
21.5
20
100
200
500
1000
2000
5000
15
20
Eccsomy / mHa
Eccsor
mHa
Figure 4: Ablation study: Breakdown of absolute energy difference between Scherbela et al. 25 and our work, when evaluating the respective base models on the test set consisting of heavy atoms. Panel shows accuracy gains for zero-shot evaluation, panel b the accuracy after 4k fine-tuning steps. First and last row depict their and our absolute energy respectively, intermediate bars show the subsequent improvements of various changes in mHa (and %). Note that panel a. is plotted on logarithmic scale due to the large energy difference.
To analyze the relative importance of our changes, we break down the accuracy gap between our work and the prior work by Scherbela et al. (25) in Fig. 4 We start with their model checkpoint trained for 128k epochs on their dataset consisting of 18 different compounds. Next, we train a model using our improved architecture with their dataset and methodology, already reducing the zero-shot error by 76% and fine-tuning error by 35%. The next model additionally uses normal-mode distortions to augment the training dataset, decreasing the fine-tuning error by another 12%. Additionally increasing the training set size to our dataset containing 98 molecules and the corresponding torsional conformers, reduces the zero-shot error by 98% and yields modest improvements in the fine-tuning case. Increasing the number of determinants from 4 to and increasing the number of pre-training steps from 128k to 256k improve fine-ture tine-tuning accuracy by 8% and 21% respectively. The two largest contributions overall are the improved architecture yielding substantial gains both in zero-shot well the larger training set, which crucial for high zero-shot accuracy.
and fine-tuning This is consistent with (25). which found that both model size and training set size do improve performance, while pre-training duration shows diminishing returns after 256k steps. Since pre-train the PhisNet model against Hartree-Focd references and do not update its parameters during optimization, the PhisNet model by itself contributes no substantial accuracy improvement. It impacts energy predictions indirectly, by providing pre- pre-trained nuclear emebddings " and enabling fast geometry distortions and rotations during training and inference, which in turn improve accuracy.
3.4 Large-scale experiment
To showcase the scalability of our approach, we evaluate zero-shot predictions of the absolute energies for subset of 250 molecules from the QM7 dataset 332. containing molecules with 14-58 electrons.

---

Since CCSD(T) calculations would be prohibitively expensive for a dataset of this size, in Fig. 54 we compare against density functional theory (DFT) reference calculations (PBE0+MBD, (26). Fig. 5p breaks down the energy differences compared to DFT, by molecule size. For molecules with less than heavy atoms we obtain energies that are lower than DFT by 10-120 mHa. Because our ansatz is variational (in contrast to DFT), our lower energies translate to more accurate prediction of the true ground-state energy, confirming again our high zero-shot accuracy. With increasing system size the accuracy deteriorates, due to increasing extrapolation and out-of-distribution predictions, consistent with our results in Sec. 3.T One reason for the large energy differences larger molecules is the increasing inter-atomic distance within the molecules with increasing number of atoms (cf. Fig. 5p). While the largest inter-atomic distance observed in the training set is 11 bohr, the evaluation set contains distances up to 17 bohr. This issue could be overcome by employing distance cutoff, as it is already applied in supervised machine-learned potential energy surface predictions [19 33
or has just recently been incorporated into a neural wavefunction (17] via an exponential decay with increasing inter-particle distance. Another potential solution is to include larger molecules or separated molecule fragments in the pre-training molecule dataset, reducing the extrapolation regime.
b
-100
125
Nr. of heavy atoms
EreEo MBD /HH
10"
mak R
100
200
Count
75
60
10
-300
25
0
330
-200
-100
500
1000
12 16
EFBEO MBD HH
E EFeco MBD mHa
max Ry / bohr
Figure 5: Zero-shot on QM7: a) Our zero-shot energies vs. DFT 26
Histogram of energy
residuals (truncated at 1.25 Ha for clarity) c) Residuals vs. largest inter-molecular distance Rt..
4
Discussion
We have presented to our knowledge the first ab-initio wavefunction model, which achieves high- accuracy zero-shot energies new systems (Sec. and See. 3.4). Our pre-trained wavefunction yields more accurate total energies than CCSD(T)-2Z across all molecule sizes and outperforms CCSD(T)-3Z on molecules containing up to 5 heavy atoms, despite having been trained only on molecules containing up to 4 heavy atoms. We find that relative energies of our model are qualitatively correct without -unintn but need on the order of 4000 fine-tuning steps to reach chemical accuracy of 1.6 mHa (Sec. (.2).. This is a substantial improvement over previous work, which so far has fallen in two categories: anch-accuracy ansätze (such as (3 6) that cannot generalize across molecules and thus need ca. 10x more compute to reach the same accuracy, or methods that can generalize ([77) 25), but yield orders of magnitude lower accuracy the zero- or few-shot regime. We demonstrate in Sec. 3.3 that these improvements are primarily driven by an improved architecture (containing a more expressive electron embedding and an ML-based orbital model), improved geometry sampling. and larger training dataset.
While results encouraging in the zero- and few-shot regime, open questions for further research abound. The most pressing issues are currently limited zero-shot accuracy for relative energies, and potentially limited expressiveness of the ansatz in the regime of very long optimization. Zero-shot accuracy could be further improved by training on an even larger dataset, further improved geometry sampling (in particular of torsion angles), and an interaction cut-off to avoid previously unseen particle pairs for new large molecules. Furthermore SE(3)-symmeet of the wavefunction should be explored further, since currently only the orbital part of our architecture is -quivaaquvrr We experimented with fully equivariant architecture, but found the resulting wavefunctions to not be expressive enough. To improve overall accuracy, attention based embeddings (6. 34) could be pursued. Additionally, we currently freeze the weights of the orbital embedding to simplify the architecture and avoid back proppaaapan through the iterative orbital localization procedure. Optimizing these weights in addition to the electron embedding will lead to more expressive ansatz.
10

---

Acknowledgements
We gratefully acknowledge financial support from the following grants: Austrian Science Fund FWF Project 3403 (P...). WWTF-ICTI9-04 (L.G.). The computational results have been achieved using the Vienna Scientific Cluster (VSC). The funders had no role in study design, data collection and analysis, decision to publish or preparation of the manuscript.

---

References
[1] M. Born and R. Oppenheimer. "Zur Quantentheoric der Molekeln". In: Annalen der Physik 1927389022 [2] László Gyevi-Nagy, Mihály Kállay, and Péter R. Nagy. "Accurate Reduced-Coss CCSD(T) Energies: Parallel Implementation, Benchmarks, and Large-Scall Applications". In: Journal of
389.20 (1927), pp. 457-484. DOT: https: /tttp:://do..orr/000 10. 1002/
Chemical Theory and Compulation 17.2 (Feb. 2021), pp. 860-878. DOI: 10. 1021/acs jctc
0c01077
[3] David Pfau et al. "Ab Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks". In: Phys. Rev. Res. 2.3 (Sept. 2020), 033429. DOL: 10. 1103/
PhysRevResearch. 2...2.033334 033429
[4] W. K. Hastings. "Monte Carlo Sampling Methods Using Markov Chains and Their Applica- tions". In: Biometrika 57.1 (Apr. 1970), pp. 97-109. DO1: 10. 1009////i/ee/////////
[5] James S. Spencer al. "Better, Faster Fermionic Neural Networks", . In: (2020).
[6] Ingrid Glehn, James S. Spencer, and David Pfau. A Self-Altention Ansalz, Ab-initio Quantum Chemistry. Nov. 2022. DOI: [0. 400.48850/// arxiv 2211. 1. 13672
[7] Weiluo Ren et al. "Towards the ground state of molecules via diffusion Monte Carlo on neural networks". In: Nature Communications 14.1 (Apr. 2023), p. 1860. DOI: 10. 00008/s4446
023 376-37609-3
[8]
Leon Gerard et al. Soond-Standard Solutions to the Schrödinger Equation Using Deep Learning: How Much Physics Do We Need?" In: Advances Neural Information Processing Systems. Oct. 2022.
[9] M. T. Entwistle et al. "Electronic excited states in deep variational Monte Carlo", In: Nature Communications 14.1 (Jan. 2023), p. 274. DOI: 10. 1008/ 9117/7444677777 022-111777-2222233 022-35534-5)
[10] Michael Scherbela et al. "Solving the Electronic Schrödinger Equation for Multiple Nuclear Geometries with Weight-Sharing Deep Neural Networks". In: Nature Computational Science
2.5 (May 2022), pp. 331-341. DO1: 10. 1038/s43588 543588 022-00222222222
[11] Yubing Qian et al. "Interatomic Force from Neural Network Based Variational Quantum Monte Carlo". In: The Journal of Chemical Physics 157.16 (Oet. 2022), p. 164104. DOL:
10. 00.1063/5.0112344
[12] Gino Cassella et al. "Discovering Quantum Phase Transitions with Fermionic Neural Networks". In: Physical Review Letters 130.3 (Jan. 2023), p. 036401. DOI: 10 1103 /
PhysRevLett. 130. 130...640033611
[13] Xiang Li, Zhe Li, and Ji Chen. "Ab initio calculation of real solids via neural network ansatz". In: Nature Communications 13.1 (Dec. 2022). p. 7895, DO1: 10. 1038/ 0111///44666-70235533335
[14] Giuseppe Carleo and Matthias Troyer. "Solving the Quantum Many-Body Problem with Artificial Neural Networks". In: Science 355.6325 (Feb. 2017), pp. 602-606. ISSN: 0036-8075, 1095-9203. DOT: 10. 1 1e21111iiiiccci aag2302
[15] Giuseppe Carleo et al. "NetKet: A Machine Learning Toolkit for Many-Body Quantum Systems". In: SoftwareX 10 (July 2019), P. 100311. DOL: 10. www./........ softx. 2019. 100311 [16] Yuan-Hang Zhang and Massimiliano Di Ventra. "Transformer Quantum State: A Multipur- pose Model for Quantum Many-Body Problems". In: Physical Review B 107.7 (Feb. 2023),
p. 075147. DO1: 10. 1103//hyysee....... 10.. 075147
[17] Nicholas Gao and Stephan Günnemann. Generalizing Neural Wave Functions. Feb. 2023. DOI:
10. 48550/arXiv 2302. 04168
[18] Simon Batzner al. Fquiiaaaqa Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials". In: Nature Communications 13.1 (May 2022), p. 2453. DOI: 10. 1038/
111467-022-299399
[19] Ilyes Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields". In: Advances in Neural Information Processing Systems. 2022.
[20] Oliver Unke et al. Equivaiaan Prediction of Molecular Wavefunctions and Electronic Densities", In: Advances in Neural Information Processing Systems. 2021.
[21] M. Gastegger et al. "A Deep Neural Network for Molecular Wave Functions in Quasi-Atomic Minimal Basis Representation". In: The Journal of Chemical Physics 153.4 (July 2020),
p. 044123. DO1: 10. 1063/0000000 0012911
12

---

[22] Albert Musaelian et al. "Scaling the Leading Accuracy of Deep Equivariant Models to Biomolecular Simulations of Realistic Size". In: (Apr. 2023). DO1: 10. 48550/ war.://rrrvvv... 2304
10061 Xiv: 2304. 10061
[23]
Nicholas Gao and Stephan Günnemann. "Ab-Initio Potential Energy Surfaces by Pairing
GNNs with Neural Wave Functions". In: arXiv:2110.05064 [physics] (Nov. 2021). arXiv:
2110. 05064 [physics]
[24]
Nicholas Gao and Stephan Günnemann. "Sampling-free Inference for Ab-Initio Potential
Energy Surface Networks". In: (2022). DO1: 000.4855/ABRA 2205 14962
[25] Michael Scherbela, Leon Gerard, and Philipp Grohs. "Towards a Foundation Model for Neural
Network Wavefunctions" : In: arXiv.org (Mar. 2023).
[26] Johannes Hoja et al. "QM7-X, a comprehensive dataset of quantum-mechanical properties spanning the chemical space of small organic molecules". In: Scientific Data 8.1 (Feb. 2021),
p. 43. DOI: 10. 94177-02119-------
[27] Ruard van Workum, Joao Malhado, and Philipp Marquetand. CASNet: Learning Complete Active Space Orbitals Using Message Passing Neural Networks. June 2023. DOT: 26434/
chemrxiv-2023-11e 18
[28] János Pipek and Paul G. Mezey. "A Fast Intrinsic Localization Procedure Applicable for Ab Initio and Semiempirical Linear Combination of Atomic Orbital Wave Functions". In: The Journal of Chemical Physics 90.9 (May 1989), pp. pp. 491664996.. DOT: 10. 1063/10011111 0ttp//1456883 Claudia Filippi and C. Umrigar. "Correlated Sampling in Quantum Monte Carlo: Route to Forces". In: Physical Review B 61.24 (June 2000),
[29]
DOI: 10. 1103/
[30] James Martens and Roger Grosse. "Optimizing neural networks with kronecker-factored approximate curvature". In: International conference machine learning. PMLR. 2015,
pp.
[31] Armagan Kinal and Piotr Piecuch. "Computational Investigation of the Conrotatory and Disrotatory Isomerization Channels of Bicyclo[1. (.c.looo11.....oo. I.OJbutane Buta-I.3-diene: A Completely Renormalized Coupled-Cluster Study". In: The Journal of Physical Chemistry A 111.4 (Feb. 2007). Publisher: American Chemical Society. pp. 734-742,
ISSN: 1089-5639, DO1: [0. 1021/
jp065721k URL: hattps://dot org 10. 1021/ p065721k
[32] Matthias Rupp al. "Fast and Accurate Modeling of Molecular Atomization Energies with Machine Learning". In: Phys. Rev. Lett. 108 (5 Jan. 2012), p. 058301. Do1: 10 1103/ /
PhysRevLett. 108. 058301
[33] Kristof T. Schütt et al. "SchNet: A Continuous-
Convolutional Neural Network for Modeling Quantum Interactions". In: (June 2017). DOL: 10. 485500/arXiv 10..17066688800
[34]
Gabriel Pescia et al. Message-Pussing Neural Quantum States for the Homogeneous Electron Gas. May 2023. DO1: 10. 48550/arXiv... 2305. 07240
[35] Qiming Sun et al. "Recent Developments the PySCF Program Package". In: The Journal of Chemical Physics 153.2 (July 2020), p. 024109. DOT: 10. 1063/5 5 0006074
[36]
Mario Geiger and Tess Smidt. E3nn: Euclidean Neural Networks. July 2022. arXiv: 2207
09453 [cs]]
[37] Diederik P. Kingma and Jimmy Ba. "Adam: A Method for Stochastic Optimization", In:
arXiv: 1412.9900 (Jan. 2017). DO1: 10. 48550/ aet0//rXiiv..o 1412. 6980 arXiv: 1412. 6980 [cs]
[38] Greg Landrum. RDKit: Open-source cheminformatics,
https github.. com / /rdkit/
rdkit GitHub repository. 2009.
[39] Frank Neese al. "The ORCA Quantum Chemistry Program Package" In: The Journal of Chemical Physics 152.22 (June 2020), p. 224108. DOI: 10. 1063/0.601000 0004608
[40] David Pfau James S. Spencer and FermiNet Contributors. FermiNet. 2020.
URL: http
//github. cww/ wwweeeeie/eeeminneeee
[41] Aleksandar Botev and James Martens. KFAC-JAX. Version 0.0.1. 2022. URL: Setp:///gthin
com/ (eepminddkkacc-a
13

---

Supplementary material for Variational Monte Carlo on a Budget Fine-tuning pre-trained Neural Wavefunction
A Electron MCMC initialization
To investigate the impact of the initial distribution of electron positions on the equilibration of the Markov Chain, we run two evaluations for glycine molecule, using a pre-trained wavefunction. We perform no imitial burn-in and use every 50th sample for energy evaluation. If the chain was perfectly cquilibrated right after initialization, all sampled energies would fluctuate around the mean energy. However as Fig. [6b shows, it takes several thousand steps for the sampled energies to converge to the correct mean. This particularly pronounced with Gaussian initialization of electron positions.
which is the default in state-of of the-ar art DL-VMC codes such as FermiNet [3]. Using an exponential distribution of the initial electron positions much more closely resembles the correct electron density yl/ (cf. Fig. (6)) and thus reaches equilibrium substantially faster.
b
10a
-280
Gaussian init.
Gaussian init,
Exponential init.
Exponential init.
Electron density, a.u.
Ground truth
-281
Ground truth
10s
Energy Ha
104
-283
-284
285
0.2
0.6
1.0
1.4
4
distance from nucleus bohr
MCMC steps
Figure 6: Effeet of electron initialization: Initializing the electron positions using an exponential distribution instead of Gaussian, better fits the actual density (a), and thus leads faster equilibration of observables during evaluation (b).
B Orbital localization
Our model uses orbital embeddings 20t as inputs parameterize the backflows furb, and exponents nhth of the orbitals. These orbital embeddings were introduced by Scherbela et al. (25) in the form of
molecular orbital expansion coefficients, obtained from a self consistent Hartree-Food Fock calculation. In this setting, the coefficients arh are not uniquely defined, but only to a linear transformation U with determinant
Noh
orb lk
U
orb
U E TERNNNR Non det U = +1.
(15)
This stems from the fact that the corresponding Hartree-Fock wavefunction invariant under such a transformation. Consequently there is free choice, which linear combination of embeddings orb to choose from without any loss of information. We follow the approach of (25) by choosing U such that the corresponding Hartree-Foo -Fck orbitals are maximally localized according to the Foster-Boys

---

metric, i.e. minimize the spatial variance L:
Nuw Musis
Φx (r, U)
bru (r))kn
(16)
1=1
C(U) - Σ Le (f
(17)
Here (r) denotes 4-thh basis function of the Hartree-Fock expansion, centered the /-th nucleus. In practice the integrals of Eq. d7] not have to be evaluated explicitly, but can instead be computed via the overlap matrix S. The minimization of L is typically done iteratively, requires on the order of 10 steps, and is readily implemented in many open open-source quantum chemistry codes such as pySCF
35
C Adaption of PhisNet
We heavily rely PhisNet by Unke al. [20] to obtain orbital descriptors without the need for separate SCF calculation. Compared to their original work, we made several simplifications, which motivated by the fact that we do not predict final high-accurace orbitals in a large basis set, but only use PhisNet as a feature extractor by predicting orbitals in a minimal basis-set:
Layer Norm We found deep variants of PhisNet to be unstable to train and mitigated the issue by adding (equivariant) layer norm after each PhisNet module.
Simplified Fock matrix prediction The original PhisNet implementation uses final interaction between the node embeddings, before predicting the elements of the Fock matrix. We found this interaction to be superfluous for our purposes and left it out for simplicity.
Separate energy head The original PhisNet computes energies via the eigenvalues obtained by diagonalization of the Fock matrix. We instead predict energies using a separate head on top of the scalar features of the node embeddings.
Smaller network We changed the hyperparameters to obtain a smaller and faster version of PhisNet which obtained sufficient accuracy for our purposes. We used 2 layers (instead of 5) and Lmax = 2 (instead of 4). This reduces the number of parameters from 17M to 3M.
Diverse training set While the original work optimized separate models for each molecule (e.g. by training on different geometries of a molecular dynamics simulation), we optimize single model to predict F, S, E, and across a dataset of 47k geometries sampled from QM7-X [26].
JAX re-implementation We re-implemented PhisNet in JAX, using the e3nn library 36 to construct the -qquivrraa operations.
We train the Phiseet-model dataset of 47k molecules from QM7X 26 using the Adam optimizer 37 on the following loss
Ephis R Z Fref.n Σ
(18)
Σ
Ephis (R" 2
-O1s
+
(19)
dRY
nI
+
Σ
(Ephis (R", 2²m
pact.
11
+
(20)
+
Σ
SPhis (R. 2
stef.n 11pp
2
(21)
nIJue
Here E denotes energies, G denotes gradients of energies, F Fock matrices, and S overlap matrices. The indices 7, J run over nuclei, the indices μl, v over basis functions, and the index n over samples in batch.

---

D Molecule datasets
Bicyclobutane For the Bicyclobutane to
1,3 1.3-butadiene transition we use the geometries from Kinal et al. 31 and compare against the reference energies stated in Spencer et al. [5].
N₂ For the N₂ potential energy surface with various bond-lengths we used the geometries including reference calculations from [
The global rotation of 360 degrees for propadiene is performed on the geometry which is part of the test set for 3 heavy atoms. For the torsion experiment we used the equilibrium geometry and rotated the torsion angle by 90" degrees in steps of 10 degrees.
Propadiene
Zero-shot and fine-tuning dataset The results zero-shot and few-shot predictions for increasing number of heavy atoms are performed on random subsets of molecules. For 5-7 heavy atoms we sample 4 unique and distorted molecules from QM7-X [26]. For 4 heavy atoms we use all geometries from the Bicyclobutane dataset. For heavy atoms we use the ablation dataset.
Ablation dataset For the ablation study, distribution test set from Scherbela et al. [25]. leading to set of four distinct molecules. We ensure that these molecules are not part of the training set.
geometry per molecule from the out-of-
Large scale experiment For the large scale experiment we used a stratified random sample of 250 molecules from QM7 (32). It contains all molecules with up to 4 heavy atoms, and additionally 65 randomly chosen molecules for 5, 6 and 7 heavy atoms each.
Pre-training dataset for transferable neural wavefunctions To train our pre-trained wavefunc- tions we use two datasets, consisting of 18 and 98 disparate molecules. For part of the ablation we use the dataset proposed in (25) and extended version with 80 additional molecules. The additional compounds are a combination of all valid SMILES generated with RDKit (38) with 3 heavy atoms, allowing only Nitrogen, Oxygen and Carbon with single-, double- or triple-bonds, and all molecules up to four heavy atoms from QM7-X (26) (excluding molecule containing Fluorine). To prevent train-test leakage, we remove Bicyclobutane (including all conformations) and the four molecules from the ablations dataset. Since the cormal-mode-distee
by design do not generate strongly distorted geometries, we augment the -8-molecule-aataset with rotated dihedral angles. To generate subset of all possible dihedral angles for a heavy-atom bond we first generate samples with equidistant angles for all possible dihedral angles and compute Hartree-Fock energies with a minimal basis-set. set. We include the equilibrium geometry and all extrema of the potential energy surface with respect to the rotation of a single dihedral angle if the energy of the extrema is significantly different to already included geometries of the same molecule. Additionally, we include the transition geometry towards the respective extrema and again only include energetic diverse states. Finally, to make sure that certain molecules not underrepresented in the dataset make sure that all molecules have at least 5 geometries that get distorted during pre-training by adding copies of the equilibrium geometry. Overall this yields 699 initial geometries R for pre-training.
E
Energy variance
Fig. 7 depicts the energy variance obtained by various models. Analogously to the the energies presented in Fig. of the main text, we also find that our pre-trained model achieves substantially lower energy variance compared previous generalizing wavefunctions. Notably on the test set with 3 heavy atoms, the variance of our zero-shot model is on par with the variance of a PsiFormer model trained for 16k steps.
F Relative energies
To better depict the accuracy of the models' relative energies for the test systems in Fig. of of the main text, we depict the difference of relative energies between our model and a references method in Fig. Furthermore we list numerical values vor all energies of Fig, 3l in Tab. I We find that on these
3

---

Zero-shot
After 4k fine-tuning steps
C
Teat set with heavy atoms
Our work (pn-trainec)
Our work (no pro-imening)
(2023)
(pre trainee))
Scherbela et al. (2523)
Gerard (2022)
610x
PaiFormer
Ever Ha?
10
10
10
Nr. of heavy atoms
Nr. of heavy atoms
Optimization steps
Figure 7: Variance of energies: Variance when re-using the pre-trained model on molccules of varying size without optimization (a) and after fine- tuning (b). For the test set containing 3 heavy atoms a function optimization steps, comparing against current state-of-the-ar methods. Solid lines correspond to pre-trained models, dashed lines to models without pre-training (c).
challenging systems our model yields relative energy errors of up 15 mHa when only fine-tuned for 400 steps per geometry. When fine-tuning our pre-trained model for 3200 steps per geometry, we obtain relative energies that accurate within 2-3 mHa.
a)
dissociation
b)
Propadiene torsion
25
work (4k fine-tuning steps)
20.0
Our
20
Our work (16k fine-tuning steps)
PsiForer
Our work (32k fine-tuning steps)
17.5
CCSD(T) / CBS
AE AErer mHa
15
Gerard et al. (2022)
AF - Arrer mHH
15.0
12.5
10
10.0
7.5
5.0
2.5
3 bond length bohr
4
5
8
12
16
20
24
28
32
Optimization steps
Figure 8: Challenging relative energies: The difference between relative energies of reference method work for challenging systems. a) Nitrogen dimer: Difference of relative energies to the equilibrium geometry for our approach to reference calculation from Gerard et al. for N2, a) Propadiene: Difference of relative energies the transition barrier of twisted vs. untwisted propadiene. As reference method we choose PsiFormer with 64k optimization steps. This figure complements Fig. 3b.d of the main text with results for additional fine-tuning steps.
G Reference energies
CCSD(T) All CCSD(T) energies except explicitly stated otherwise were obtained using ORCA 39 starting from a restricted Hartree-Fodd caleulation. We use correlation consistent basis sets the cc-pCVXZ family, with X in (2, 3, 4). To extrapolate to the complete basis set limit (CBS), use

---

(CSD(T)
DMC
FermiNet 200k
FermiNet 10k
Our work zero-shot
Our work
structure
[31]
31
5
700 per geom
con con_TS
64.4
64.4
64.1
63.9
94.0
66.6
dis _S
34.7
93.4
92.0
87.1
183.8
94.5
g-but -40.0
-40.2
-40.3
-44.9
-48.5
-40.4
gt-TS -35.5
-35.4
-35,9
-42.9
-46.9
-36.7
t-but -44.6
-44.5
-45.3
-47.5
-51.8
-43.2
Table I: Energies relative to the energy of bicyclobutane in mHa, including the zero-poinn vibrational
energy correction from Kinal et al. [3I.. This data complements Fig. 3a in the main text.
the approach outlined in and fit the following functions with free parameters CBS PCBS a,
HF CBB + ae
-bX
EX E
ECCSDIT
CCB
c.
ELCSDCT) CBS
EBS + EBBS
We stress that although CCSD(T) are often considered as "gold-standard", they do not necessarily represent the actual ground-state energy. There many cases, where CCSD(T) either overestimates the true ground-state energy, or even underestimates it, because CCSD(T) does not yield upper bounds to the ground-state energy.
PsiFormer For Fig. 2 we used the open-source FermiNet codebase 440] The codebase didn't allow for inference calculation, therefore a slight fix was applied. All calculations were performed with the small settings proposed in von Glehn et al. [6]
H Hyperparameters
A detailed description of the hyperparameter used in this work can be found below (ef. Tab. 2 For the mapping the orbital descriptors to the electron embeddings to build the orbitals we rely on the hyperparameter from (25). For optimization we rely on the second-order method KFAC (30) and use their Python implementation [41]. During the continuous sampling of the geometries we allow each geometry perform maximum of 20 steps of normal-mode distortion from the initial geometry and reset to the original one once the threshold is reached.
I Runtime and computational resources
Tab. listss the run-times for our method and a com- parable conventional reference method (CCSD(T)- 4Z) for a typical molecule with 5 heavy atoms. We find that our fine-tuning our model for 4k steps yields lower absolute energies than CCSD(T)-4Z (cf. Fig. 2) similar computational cost. This is in contrast to many earlier works that achieve highly accurate energies, but computational cost that far surpasses the cost of conventional methods for small molecules.
Method
Runtime /
Our work, zero-shot
1.0
Our work, 4k fine-tuning
9.5
PsiFormer, 16k steps
53.0
CCSD(T), 4Z basis
10.5
Table 3:
Runtime of various quantum chem- istry methods for C,OH,. All timings are in node hours, being A A100-hours for our work and PsiFormer, and 256 CPU-hours for CCSD(T).
Overall we used x 5k GPUhs (A100) for devel- opment and training of our base models, and an- other 5k GPUhs (A40) on evaluations and fine- tuning. Additionally we required ss 20k CPUhs for CCSD(T) reference calculations.

---

J Code and data availability
All code, configuration files, geometries, datasets and obtained energies are available on GitHub Model weights are available on figshare
com wwtp:///w/wwwe//eeeeeeeeeeeeeoe under https: org/ 6084/m9 figshare..... 23585358.v1
under
Hidden dimension Nemb No iterations
256
Electron Embedding
Hidden dimension ""ue Ne layer MLP
64
Nuclear Embedding
Activation function
SiLU
No layer edge embedding
3
Message passing
Dimension edge embedding
Dimension linear layer
32
N: walkers
2048
Markov Chain
Ne decorrelation steps
50
Monte Carlo
Target acceptance prob.
50%
Pre-trained against basis set
STO-6G
No iterations
2
Harmonic degree
PhisNet 20
No radial basis functions
128
Hidden dimension x nue
128
Distance cutoff (bohr)
30
Ne determinants Maat
No hidden layers forb Hidden dimension of foorh
256
No hidden layers grb Hidden dimension orb
2
Transferable atomic orbitals 25
128
No iterations MPNN
No radial basis functions Hidden edge embedding dimension
16
Hidden node embedding dimension
16
Activation function
SiLU
Optimizer
KFAC
Batch size
2048
Norm constraint
3 x 10-3
Initial damping do
Variational pre-training
Minimal damping dmin
0.001
Damping rate decay
d(l) = dn exp( -1/20000)
Initial learning rate Iro
0.1
Learning rate decay
Ir(() = Iro + t. //000))
Optimization steps
128, 208,000 256.000
Changes for fine-tuning
Learning rate decay Optimization steps
Ir(t) = 1no + t/6000) 0 8230000
Distortion energy /
0.005 Ha
Sampling geometries
Max age
20
Bias towards original geometry
0.2
Table 2: Hyperparameter settings used in this work


Reviews:
[{"review_id": "DQNSj4jwFr", "reviewer": "Reviewer_m8hG", "summary": "The paper proposes a neural network wave function model which generalizes to different molecules and geometries. The neural network includes electron model, orbital model, both with message passing, Slater determinants, and an energy estimate for geometry sampling. The network is pretrained on 98 molecules and results show lower energies with one-shot and fine-tuning, on new molecules, compared to other generalized wavefunctions and some computational chemistry approaches. The paper also includes an ablation study and prediction of absolute energies for 250 molecules.", "questions": "From Fig. 3c, it seems that the errors are somewhat dependent on the rotation angle. Is that actually happening? \n\nHow much of the improvement attributed to \"improved architecture\" may be due to Phisnet? ", "limitations": "The authors are upfront about the limitations: potential limited expressiveness with more optimization iterations, prediction is worse for increasing extrapolation and out-of-distribution predictions. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The idea is original since there are very few methods of generalizable neural networks for wavefunctions. \nThe normal mode of sampling geometries from an inverse hessian of energy seems to be original.\nThe experiments and figures are high quality. \nThe results are very good, it seems to be the first generalizable, ab initio quality wavefunction, that can be trained with quantum Monte Carlo. Generalizability can be computationally efficient if not every optimization needs to be started from scratch.\nThe problem of solving Schrodinger equation is an important and significant one. ", "weaknesses": "I'm not so familiar with CC-2Z, -3Z etc methods the paper compares against. Not sure how the proposed method compares with something like Slater-Jastrow variational Monte Carlo (which is quite fast and generalizable for an abinitio method).  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review! Regarding your specific questions and potential weaknesses:\n\n### Weaknesses\n\n>I'm not so familiar with CC-2Z, -3Z etc methods the paper compares against. Not sure how the \nproposed method compares with something like Slater-Jastrow variational Monte Carlo (which is quite fast and generalizable for an abinitio method).\n\nWhen we refer to CC-2Z, etc. we refer to CCSD(T), Coupled Clousters Single Doubles (Perturbative Triples), which is the workhorse of high-accuracy calculations in computational chemistry. A key choice in these calculations is the basis size, for which typically correlation consistent basis sets are being used. Their size is typically denoted by double-zeta, triple-zeta, etc. and thus often abbreviated as 2Z, 3Z, etc. Thus when we refer to CC-4Z, we refer to CCSD(T) with the basis set ppCV4Z. We will explain this in more detail in the updated manuscript.\nWhen comparing CCSD(T) against Slater-Jastrow Monte Carlo, CCSD(T) is typically vastly superior in energies, but also computationally much more expensive. A systematic comparison of SJ-VMC and SJ-DMC vs. neural wavefunctions can for example be found in Fig. 7 of [1], showing that SJ-VMC is typically ~100 mHa higher in energy and SJ-DMC ~10 mHa higher in energy.\nAs simple test, we took a random molecule (H2CO) from a SJ-DMC benchmark [2] and compared energies obtained from SJ-DMC, CCSD(T) and our 0-shot wavefunction. We also find CCSD(T) to be superior to SJ-DMC, only outperformed by neural wavefunctions:\n| Method                | Energy / mHa |\n|-----------------------|--------------|\n| SJ-DMC                | -114.480     |\n| CCSD(T)-4Z            | -114.485     |\n| Our 0-shot evaluation | **-114.488**     |\n\n### Questions\n\n>From Fig. 3c, it seems that the errors are somewhat dependent on the rotation angle. Is that actually happening?\n\nYes, because our wavefunction is not fully SE(3)-equivariant, the obtained energies are not rotation invariant, and thus depend on the chosen global rotation angle. Currently no architecture is known that allows SE(3)-equivariance of the wavefunction, efficient antisymmetrization via the determinant and smooth dependency on the geometry at the same time. We will add further details on this point in the appendix of the revised manuscript.\nHaving said that, our energies are approximately rotation invariant, because we randomly perturb and rotate geometries during variational pre-training, leading to similar energies for any global rotation angle.\nFrom such an approximately rotation-invariant model we can even obtain fully rotational invariant energies, as depicted in the rebuttal Fig. 3. By also applying random rotations during evaluation (which requires no substantial extra cost), we can obtain fully rotational invariant energies (up to Monte Carlo noise).\n\n> How much of the improvement attributed to \"improved architecture\" may be due to Phisnet?\n\nPhisNets orbital descriptors do not contribute to the accuracy directly. PhisNet does have an implicit effect on the accuracy, by being able to efficiently predict the Hessian matrix and therefore being able to enhance sampling of the geometries it contributes to improved accuracy.  \n"}, {"review_id": "db3bMukUny", "reviewer": "Reviewer_NURi", "summary": "This paper tackles the problem of solving quantum Monte Carlo optimization on multiple molecules. A PhisNet model is used to encode geometric information from nuclei and other GNN models are then used to embed electrons. Additionally, the authors propose a method to generate more diverse nuclear geometries by a specially designed distortion method. Experiments show that the method can achieve good accuracy when trained smaller molecules and test on larger molecules. After some fine-tuning, the method can achieve improved results.\n", "questions": "- Among the many experimented datasets, which one(s) most decisively show the benefit of the proposed method?", "limitations": "None", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The authors provide detailed description of the architecture.\n-  The paper uses diverse datasets to evaluate their method.", "weaknesses": "- The improvements of the proposed method mainly comes from an improved architecture. To my understanding, the novelty seems to be limited.\n- The proposed method does not yield improved results when more optimization steps are used.\n- The presentation of the training and evaluation steps is not clear. Specifically, how is pre-training and fine-tuning done.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and the effort to improve our work. We hope we can address in the following rebuttal all raised questions and comments. \n\n### Weaknesses\n\n> The improvements of the proposed method mainly comes from an improved architecture. To my understanding, the novelty seems to be limited.\n\nThe architecture yields roughly 50% of the improvements, which we explicitly demonstrate in the ablations (Fig. 4). Other important contributions to the performance are for example the initialization of Monte Carlo walker, significantly reducing the number of necessary burn-in steps as highlighted in Appendix B (Fig. 6), which is a large part of the computational cost when performing zero-shot evaluations or the continuous sampling of geometries with normal mode distortion during pre-training of a base model. \n\n> The proposed method does not yield improved results when more optimization steps are used.\n\nOur method does yield improved results with more optimization steps, both during pre-training and during fine-tuning: For pre-training we show in Fig. 4 that increasing the number of pre-training steps from 128k to 256k yields large gains in accuracy of 37% for zero-shot and 21% for fine-tuning respectively. For fine-tuning Fig. 2c shows continuing improvements in energy, and our rebuttal Fig. 1 shows corresponding improvements in energy variance.\nIt is true that during fine-tuning we see diminishing returns with more fine-tuning, but this can be primarily ascribed to already starting from a very good wavefunction right after pre-training.\nFurthermore, as shown in Fig. 2, we outperform other generalized wavefunctions by two orders of magnitude in the zero-shot regime. We are able to reach the same performance of other generalized wavefunctions with 8x less optimization steps and surpass this accuracy with additional steps (cf. Fig 2c blue vs. orange). Additionally, we also outperform well established high-accuracy methods such as CCSD(T) with very few fine-tuning steps, again demonstrating a large improvement over the previous SOTA.\n\n> The presentation of the training and evaluation steps is not clear. Specifically, how is pre-training and fine-tuning done.\n\nThank you for the feedback. We will add an additional clarifying paragraph to the main text in section 2.3 and highlight the specific difference to pre-training and fine-tuning a neural wavefunction for the final version of the manuscript. \nTo summarize in the present manuscript we refer to pre-training as the stage, where a neural wavefunction is optimized extensively on a number of molecules (in our case up to 98 molecules) with varying number of geometries. This pre-trained wavefunction is used to initialize a wavefunction for a potentially completely new molecule for further fine-tuning steps. Normally, the fine-tuning stage is rather short with a couple of hundred optimization steps (vs. thousands of steps in a standard calculation). When no additional optimization steps are performed and the wavefunction is directly evaluated we speak of zero-shot predictions. \n\n\n### Questions\n\n> Among the many experimented datasets, which one(s) most decisively show the benefit of the proposed method?\n\nThe most direct comparison is against the work by Scherbela et al.. We are able to significantly improve upon their zero-shot results by two orders of magnitude (Fig. 2a). We also outperform CCSD(T) calculations for systems significantly larger than the ones present in our training set. The fine-tuning accuracy is improved by a factor of 5x in the extrapolation regime of 5 heavy atoms (Fig 2b). Additionally, we reach the same accuracy as Scherbela et al. for 3 heavy atoms 8x faster with a significantly lower final ground-state energy prediction (cf. Fig 2c). \nMoreover, our approach demonstrates competitive performance even against DL-VMC methods tailored for specific systems, requiring an order of magnitude fewer optimization steps, as evidenced in the bicyclobutane experiment (cf. Fig 3a).\n"}, {"review_id": "spc9J5mUod", "reviewer": "Reviewer_37BZ", "summary": "The paper proposes a meta-learning model for neural wavefunctions. The neural model is first pre-trained by solving the variational optimization problem on many molecules. To adapt the model for a new problem (such as a new molecule), relatively few fine-tuning steps are needed to arrive at an accurate solution. The method consists of an orbital model based on PhisNet (with SE-3 equivariance), an improved electron embedding using message passing on the nuclear embeddings from PhisNet , and a Slater determinant to ensure antisymmetric property of the wavefunciton. The model is trained on a dataset of 98 small molecules with up to 4 heavy atoms. The resulting \"base model\" shows stronger zero-shot transfer across different number of heavy atoms than recent meta-learning wavefunction models. With a small amount of fine-tuning (4k steps), the model achieves very accurate energies, outperforming CCSD(T)-4Z on molecules up to 5 heavy atoms. It also achieves excellent accuracy on challenging systems like conformer energies, dissociation curves, and relative energies of transition states.", "questions": "- During sampling of molecular geometries, are they non-physical geometries discarded?\n- Are the electron embeddings SE-3 equivariant? (I assume the MPNN used does not preserve equivariance.)", "limitations": "N.A.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Use of the improved architecture embeddings based on PhisNet that preserves SE-3 equivariance.\n- The sampling of initial electron positions to run MCMC for VMC are improved by using an exponential distribution from the nuclei.\n- For sampling geometries of the molecules, the covariance of the Gaussian is rectified by inverse Hessian of the energy, yielding better geometries that make more sense physically.", "weaknesses": "- Relative energies, which depend on small energy differences, are not very accurate without fine-tuning on that specific system. The transferability is not as good for relative energies.\n- The overall end-to-end training is built upon previous works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the review and feedback regarding our submission. \n\n### Weaknesses \n\n> Relative energies, which depend on small energy differences, are not very accurate without fine-tuning on that specific system. The transferability is not as good for relative energies.\n\nPredicting accurate relative energies remains a challenging task for DL-VMC methods. Previous studies have highlighted a noticeable discrepancy between relative energies predicted by DL-VMC and conventional high-accuracy quantum chemistry approaches. For instance, in the case of the studied Nitrogen dimer, several high-accuracy DL-VMC methods predict an energy barrier for bond breaking that is several mHa higher than what is observed in experiments and conventional approaches [3]. Additionally, for the Benzene dimer, despite investing several thousands of GPU hours, DL-VMC fails to accurately predict the sign of the relative energies [2, 4].\nAnother relevant work, which also optimizes a wavefunction across multiple molecules simultaneously, overestimated the energy barrier for Ethene by several 100mHa's [1, 5]. While DL-VMC generally achieves state-of-the-art accuracy in absolute energies, it falls short in providing accurate predictions for specific systems' relative energies. Nonetheless, we openly and transparently acknowledge this limitation in our limitations section and anticipate that future research will address and improve upon this aspect.\n\n> The overall end-to-end training is built upon previous works.\n\nFirst of all, we believe that building on good previous work is admirable rather than deplorable. However, we do add several novelties that yield our substantially improved results: \n\n- The original PhisNet was only trained on MD simulations of a single molecule, which is insufficient in our case. We therefore pre-train PhisNet across a variety of chemically diverse systems simultaneously. To enable this we made several simplifications to PhisNet to increase stability as listed in Appendix F. \n\n- We improved the architecture of the electron embedding, lowering energy errors by 50 %.\n\n- We added a sophisticated sampling strategy to obtain a diverse but physically plausible set of molecule geometries.\n\n- We successfully scale training to this more diverse pre-training dataset, allowing us to reach two orders of magnitude better zero-shot performance to previous work. \n\n\n### Questions\n\n> During sampling of molecular geometries, are they non-physical geometries discarded?\n\nWe avoid non-physical geometries by distorting the geometries with respect to the normal modes of the Hessian of the energy. Furthermore, we bias each distortion towards the original geometry and reset the distortion trajectory every 20 epochs to the origin. Apart from that we do not reject any geometries and don’t observe instabilities with respect to unphysical geometries.   \n\n> Are the electron embeddings SE-3 equivariant? (I assume the MPNN used does not preserve equivariance.)\n\nThe electron embeddings are indeed not SE(3)-equivariant. We deliberately break this symmetry by using cartesian coordinates as input features, similar to all other neural wavefunctions such as FermiNet, Psiformer, DeepErwin or PESNet.\nThis symmetry breaking is required, because the orbitals used to build the wavefunction must be symmetry-broken in order to be sufficiently expressive. If the orbitals are built from SE(3)-equivariant features without symmetry breaking, the resulting wavefunction is generally invariant to rotations [6, appendix B] and thus insufficiently expressive.\nTo obtain approximately rotation invariant energies from this symmetry-broken wavefunction, we use data augmentation during pre-training of our wavefunction, by randomly rotating geometries every 20 steps as described in Sec 2.2 (l. 174-175).\nTo obtain truly rotation invariant energies from this approximately invariant model, we can also apply random rotations during inference. Because the evaluation energy is a Monte Carlo average over many different electron samples, we can choose different global rotations for these samples at no additional computational cost, to obtain rotationally averaged (and thus invariant) energies.\nWe demonstrate this technique on the example of a globally rotated propadiene molecule in the rebuttal document, Fig. 3.\n\n[1] Scherbela et al., arXiv:2303.09949, 2023.\n\n[2] Von Glehn et al., arXiv:2211.13672, ICLR (2023)\n\n[3] Gerard et al., arXiv:2205.09438, NeurIPS (2022).\n\n[4] Ren et al., doi.org/10.1038/s41467-023-37609-3, Nature Communications (2023).\n\n[5] Gao et al., arXiv:2302.04168, ICML (2023).\n\n[6] Gao et al., arXiv:2110.05064, ICLR (2022)\n"}, {"review_id": "LUN1T7KvDr", "reviewer": "Reviewer_R7Vm", "summary": "This paper concentrates on applying deep learning techniques to electronic structure computations. Unlike the strategy adopted in [1], the authors introduce a pretraining process to enhance the few-shot or zero-shot performance of Deep Learning Variational Monte Carlo (DL-VMC) methods. As their initial step, they design a novel architecture for the conventional model. Drawing inspiration from [2], the paper proposes a learnable orbital model for adjusting to molecular structures with a disparate number of electrons. The orbital model employs a neural network to predict the Fock and overlap matrices in the initial stages, and subsequently, it resolves a generalized eigenvalue problem to derive the orbital coefficients. Additionally, the network within the orbital model predicts the energy and selects new configurations for training. With the improved architecture, the authors pretrain the entire model using a dataset containing 98 distinct molecules. The pretrained model is benchmarked under various metrics, including absolute energy and relative energy. Furthermore, the authors conduct extensive experiments, comparing their zero-shot performance with Density Functional Theory (DFT) methods on the QM7 dataset. The empirical results demonstrate the effectiveness of their methodologies.\n\n[1] David Pfau et al., Ab Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks\n\n[2] Nicholas Gao and Stephan Günnemann, Generalizing Neural Wave Functions\n", "questions": "Questions about experiments have been shown in the weakness.\n\nQuestions about methods:\n1. Why do the authors use the Hessian matrix of predicted energy to guide the sampling of configuration?\n2. How to train the network to predict $E_{phis}$?\n3. Why did the authors choose to solve the generalized eigenvalue problem to produce orbital coefficients? How is the gradient computed through this block? ", "limitations": "Compared to other ab-initio methods, the proposed method requires pretraining on a large molecule dataset. Therefore, its performance highly depends on the similarity between pretraining and evaluation datasets, which limits the application of the method. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Improved architecture. Compared to the [1, the authors employ a refined architecture to extract embedding features from electrons and nuclei. Rather than using orbital coefficients derived from the Hartree-Fock method, they directly predict the Fock matrix and overlap matrix, subsequently solving the generalized eigenvalue problem. This strategy eliminates the reliance on the Hartree-Fock method, an inaccurate solution for electronic structure. Ablation studies reveal that this enhanced structure significantly improves absolute energy performance in both zero-shot and few-shot settings.\n2. Enhanced Sampling Methods. Despite the insignificance of the MCMC burn-in step in the standard FermiNet-like training process, it can be time-consuming in few-shot/zero-shot settings. To counteract this issue, the authors refine the initialization of MCMC walkers, achieving equilibrium within a shorter MCMC burn-in stage. Beyond this initialization enhancement, they also identify and address the burn-in requirements during the pretraining stage by adjusting the movement of electrons relative to the shifts in nuclei position. Both these strategies improve the convergence of MCMC walkers, leading to superior training performance.\n3. Larger Pretraining Dataset. In addition to the technical enhancements, the authors also construct a more substantial dataset for pretraining, a development that could prove beneficial for future research in this area.\n4. Extensive Empirical Study. The authors execute extensive experiments to showcase the effectiveness of their methods, including comparisons with the Density Functional Theory (DFT) method on a large-scale dataset.\n\n[1] Michael Scherbela, Leon Gerard, and Philipp Grohs, Towards a Foundation Model for Neural Network Wavefunctions.", "weaknesses": "1. The relative energy experiments may not be convincing. Unlike the absolute energy, if the relative energy is not converged, the discrepancy between CCSD(T) and the proposed method could magnify with additional finetuning steps. Although model with 4k finetuning steps achieves relative energy nearly identical result as CCSD(T),  the reviewer suggests the authors to supply the relative energy result of model with an increased number of finetuning steps (e.g., 8k and 16k) to demonstrate energy convergence.\n2. The comparison with Density Functional Theory (DFT) methods should also encompass speed comparison. Typically, Deep Learning Variational Monte Carlo (DL-VMC) methods should benchmark their energy results against post-HF methods and Quantum Monte Carlo (QMC) results, given that the total computation cost of these methods is comparable. Despite this paper's focus on zero-shot performance, the computational cost of evaluation of DL-VMC remains high. Therefore, the reviewer suggests that the authors provide speed comparisons to elucidate the effectiveness of their method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback!  Regarding your specific questions and potential weaknesses:\n\n### Weaknesses \n\n**W1**\n\nWe ran additional experiments, increasing the nr of fine-tuning steps up to 32k, and depict the result in Fig. 2 of the rebuttal.\nWe find that with increasing nr of fine-tuning steps, our relative energies converge towards the high-accuracy reference calculations. Furthermore, we wish to emphasize that we attain comparable results to the reference methods, utilizing a significantly lower number of fine-tuning steps per geometry (100k vs 3k for N2 and 64k vs. 3k for propadiene).\n\n**W2**\n\nUnfortunately, the reference we used for the DFT energies does not include run-times for the DFT calculations. Generally, DFT, being a mean-field solution, is typically significantly faster than Quantum Monte Carlo (including deep-learning-based methods), but also often much less accurate. \nThe essence of this experiment was to provide a glimpse into the potential of our proposed approach, enabling rapid and precise computation of extensive datasets. Additionally, inclusion of this experiment within the manuscript serves to underscore the remaining challenge of extrapolating to substantially larger molecules, i.e. going from 4 heavy atoms during pre-training to 7 heavy atoms in the experiment. It's noteworthy that in the case of 4 heavy atoms, our method outperforms DFT calculations in terms of accuracy (and even CCSD(T) as seen in Fig. 2). However, when going to larger system sizes the zero-shot accuracy of our model decreases, as is clearly acknowledged as a current limitation of this model.\nTo give a better sense of computational cost, we present a comparison of computational costs against CCSD(T) (cf. Table 1 of rebuttal pdf) and achieve similar accuracy with a similar amount of computational resources.\n\n### Questions\n\n**Q1**\n\nDifferent degrees of freedom (DoF) in a molecule can have a strongly varying impact on the total energy. Some DoF (such as bond lengths) are typically “stiff”, i.e. small changes lead to large changes in energies. Others (such as twisting a single bond) are typically “soft”, i.e. have little impact on the energy.\nMolecules will therefore typically see large changes along soft DoFs, and only small changes for stiff DoFs. Preconditioning the random perturbations with the inverse of the Hessian achieves exactly that.\nIntuitively, instead of making perturbations of a given magnitude in geometry, this methods leads to perturbations of a given energy scale.\nWe will expand Sec. 2.2 of the revised manuscript accordingly.\n\n**Q2**\n\nWe train PhisNet to optimize for multiple objectives, one objective is directly to minimize the energy of a reference method, which is in our case a Hartree-Fock calculation. As it is common in supervised ML potential we predict the energy per atom from the node features and use as an estimator for the energy the sum over all atom contributions. \nWe will add a clarifying sentence in the main text under section 2.3. and have a more detailed description of the loss function in Appendix F. \n\n**Q3**\n\nThis is an established “trick” used e.g. by architectures such as SchNOrb or PhisNet, to obtain a well defined learning problem.\nMean field orbital coefficients arise as solutions to an eigenvalue problem and are thus not uniquely defined. Eigenvectors are only defined up to their sign and for degenerate eigenvalues only up to an arbitrary rotation in the degenerate subspace. Furthermore, when looking for the eigenvectors corresponding to the lowest n eigenvalues (e.g. corresponding to the occupied orbitals), infinitesimal changes in the geometry can lead to different eigenvalues being included. The set of n lowest eigenvalues thus changes discontinuously as a function of geometry.\nThese problems make learning orbital coefficients directly an ill-posed problem. The matrices F and S on the other hand are well defined and continuous as a function of input geometry. All ambiguities described above only exist in the diagonalization procedure and thus do not influence the learning process of the network itself. A recent preprint comparing direct learning of orbitals vs. learning of the matrices F and S [6] finds the latter approach to be vastly superior.\nDifferentiation through this linear algebra routine is well supported in deep learning frameworks such as JAX (similar to differentiation through the determinant). In our specific application, we do not differentiate through the solver, but keep the weights of the pre-trained PhisNet model fixed during wavefunction optimization to avoid complexity.\n\n### Limitations\n\nAlthough, the mentioned claim about the similarity between pre-training and evaluation dataset holds for any machine-learning-based method, we want to highlight that the pre-training dataset already contains a diverse range of chemical structures, encompassing equilibrium geometries, various chemical conformers, broken bonds, and different types of bond orders. This aspect sets our approach apart from previous methods like PESNet [3, 4], Globe [2], DeepErwin [5], and Transferable Atomic Orbitals [1]. We have managed to enhance the accuracy of the first properly generalizing wavefunction, which had poor zero-shot accuracy in Scherbela et al.'s work [1], by two orders of magnitude.\nFurthermore, our research demonstrates that significant improvements can be achieved by increasing the size of the dataset. Based on this finding, we hypothesize that further enhancements in generalizability can be attained by considering an even larger number of molecules and employing a more advanced scheme for sampling torsional angles.\n\n\n[1] Scherbela et al., arXiv:2303.09949, (2023).\n\n[2] Gao et al., arXiv:2302.04168, ICML (2023).\n\n[3] Gao et al., arXiv:2110.05064, ICLR (2022). \n\n[4] Gao et al., arXiv:2205.14962, ICLR (2023).\n\n[5] Scherbela et al., doi.org/10.1038/s43588-022-00228-x, Nat. Comp. Sci. (2022). \n\n[6] Van Workum et al, 10.26434/chemrxiv-2023-lwj87, (2023)."}, {"review_id": "NYIXllEg22", "reviewer": "Reviewer_BduG", "summary": "The authors propose a pre-trained deep learning wavefunction ansatz for variational Monte Carlo (VMC) solutions of the Schrödinger equation. The ansatz is trained on a chemically diverse set of (different conformations of) molecules. The zero-shot accuracy of the proposed method is a significant improvement over the previous state-of-the-art of pre-trained wavefunction ansätze, but additional fine-tuning steps are still necessary to obtain accurate relative energies. Due to the relatively simple model architecture, the achievable accuracy (in the limit of unlimited fine-tuning steps) is lower compared to more complex model architectures like PsiFormer, which however, need to be trained from scratch for every system.", "questions": "- Looking at the accuracy vs. number of heavy atoms curves shown in Figure 2, I wonder how \"practical\" the current method is compared to running coupled cluster calculations. For example, how expensive is the calculation of a system with six heavy atoms at the CC-4Z level compared to running 4k fine-tuning steps with the presented method? Are we at a stage already where VMC methods are comparable in cost, or even cheaper than conventional quantum chemistry methods?\n\n- The authors show in Figure 3 and discuss in the text that the energies obtained with their method are not invariant w.r.t. global rotations. It is correct that the wavefunction is not invariant w.r.t. rotations, but it should be equivariant. The authors mention that they experimented with fully equivariant architectures, but observed worse results compared to their (non-equivariant) method. I assume this is because some ground state wavefunctions require \"symmetry-breaking\", e.g. the ground state of a carbon atom is not spherically symmetric. Still, once a \"symmetry-broken\" wavefunction is obtained, it should be equivariant under rotations. Have the authors already looked into the distinction between equivariance/symmetry-breaking? In my opinion, any insights regarding this problem are valuable and should probably be discussed in more detail (in the appendix/supplement).\n\n- It is well known that the energy variance in VMC methods allows to assess how close the found solution is to the true ground state wavefunction, which would have a variance of 0. The authors show many plots with absolute/relative energies, but how do energy variances compare between different methods? Related to this: How large are energy variances when comparing the zero-shot to the few-shot setting?\n\n- I believe there is a small mistake on line 122: Based on the description in the following lines ($l=0,\\dots,L$ and $m=-l,\\dots,l$), it should be  $x^{\\mathrm{nuc}} \\in \\mathbb{R}^{N_\\mathrm{nuc}\\times(L+1)^2\\times N_{\\mathrm{channels}}}$ (instead of $x^{\\mathrm{nuc}} \\in \\mathbb{R}^{N_\\mathrm{nuc}\\times(2L+1)^2\\times N_{\\mathrm{channels}}}$). Can you please check this? ", "limitations": "The authors fairly discuss limitations of their work. While its potential for negative societal impact is not discussed, I think this is acceptable considering that this work is about the solution to a fundamental problem in quantum chemistry and a discussion of negative societal impacts would probably appear contrived.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper adresses the important problem of reducing the computational overhead of VMC methods using a deep learning wavefunction ansatz. The description of the proposed method is clear and well-written. Further, the authors include their code for reproducibility and the method is tested on a large number of systems.", "weaknesses": "While the paper compares results with previous methods, these comparisons consider only metrics like the number of optimization steps, but not the raw computational cost in terms of CPU/GPU time, which is arguably more important in practice. It is difficult to assess how \"practical\" the presented ansatz is compared to more traditional quantum chemistry methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! Regarding your specific comments:\n\n### Weaknesses\n\n> While the paper compares results with previous methods, these comparisons consider only metrics like the number of optimization steps, but not the raw computational cost in terms of CPU/GPU time, which is arguably more important in practice. It is difficult to assess how \"practical\" the presented ansatz is compared to more traditional quantum chemistry methods.\n\nLike all other deep-learning-based VMC (DL-VMC) methods, our method is computationally costly. Computing the energy of a single molecule with 5 heavy atoms requires 10.5 node-hours with CCSD(T)-4Z (1 node = 256 CPU cores) and Psiformer requires 53h (for 16k epochs on 2xA100). Using our method requires 1h for zero-shot evaluation or 9.5h for 4k fine-tuning steps + evaluation (on 2xA100). This is to the best of our knowledge the first study that achieves similar accuracy at computational cost comparable to established quantum chemistry methods such as CCSD(T)-4Z. We will include these findings in the appendix of the final manuscript to clarify how the different methods compare with respect to computational cost.\n\nFurthermore a very recent preprint by Li. et al [3] has demonstrated that DL-VMC can be sped up by more than an order of magnitude, via a faster evaluation of the Laplacian operator in loss calculation. These per-step speed-ups can in principle be directly combined with our reduction in the required number of steps to yield a method that is competitive with established methods.\n\nOverall, our approach might not yet be practical compared to decades of quantum chemistry developments, but is already more accurate and efficient on several interesting systems.\n\n### Questions\n\n> Looking at the accuracy vs. number of heavy atoms curves shown in Figure 2, I wonder how \"practical\" the current method is compared to running coupled cluster calculations. For example, how expensive is the calculation of a system with six heavy atoms at the CC-4Z level compared to running 4k fine-tuning steps with the presented method? Are we at a stage already where VMC methods are comparable in cost, or even cheaper than conventional quantum chemistry methods?\n\nIn Table 1 of the rebuttal PDF, we have incorporated the timings for CCSD(T)-4Z, our zero-shot approach, and our approach with 4k fine-tuning steps (see also the response to weaknesses and the comment to all reviewers). In essence, when considering a molecule with 5 heavy atoms, the computational cost for our work, specifically with 4k fine-tuning, are on par with those of CCSDT(T), yielding equivalent levels of accuracy and runtime. \n\n> The authors show in Figure 3 and discuss in the text that the energies obtained with their method are not invariant w.r.t. global rotations. It is correct that the wavefunction is not invariant w.r.t. rotations, but it should be equivariant. The authors mention that they experimented with fully equivariant architectures, but observed worse results compared to their (non-equivariant) method. I assume this is because some ground state wavefunctions require \"symmetry-breaking\" [...]. Have the authors already looked into the distinction between equivariance/symmetry-breaking? In my opinion, any insights regarding this problem are valuable and should probably be discussed in more detail (in the appendix/supplement).\n\nWe have added a comment on equivariance in the overall author rebuttal and will expand on this topic in the appendix of the revised manuscript.\nBeyond the challenges of equivariant wavefunctions described there, we would like to comment on the reviewer’s suggestion of obtaining a symmetry broken wavefunction that rotates with the molecule. This has been tried in the form of equivariant coordinate systems: Instead of using E(3)-equivarient networks, one canonicalizes the input to the network, by choosing a coordinate system which rotates with the molecule. This has been proposed by Gao et. al 2022 [1] (in the form of a global coordinate system, which rotates with the molecule) and Gerard et al. 2022 [2] (in the form of local atom centered coordinate systems, rotating with their local environments). However, despite ensuring equivariance, these methods have not been used in follow-up works by the respective groups. This is because these coordinate systems / symmetry-breakers change discontinuously as a function of the molecular geometry and can thereby introduce artificial discontinuities / kinks in the energy.\nWe are unfortunately not aware of a method that can obtain symmetry broken, antisymmetric wavefunction that change smoothly as a function of geometry.\n\n> It is well known that the energy variance in VMC methods allows to assess how close the found solution is to the true ground state wavefunction, which would have a variance of 0. The authors show many plots with absolute/relative energies, but how do energy variances compare between different methods? Related to this: How large are energy variances when comparing the zero-shot to the few-shot setting?\n\nFig. 1 of the rebuttal PDF, shows variances for all curves from Fig. 2 of the main text and we will include it in the final version. \nWe observe a similar behavior as for the ground-state energy, wherein our approach exhibits very low variance even without fine-tuning. In contrast, other DL-VMC methods display a much higher initial variance and eventually converge to a similar point compared to our work after 32k optimization steps.\n\n\n> I believe there is a small mistake on line 122: Based on the description in the following lines (l=0, ..., L  and m=-l, ..., l) it should be x^{nuc} R^{N_nuc x (Lx1)^2 x N_channels} instead of x^{nuc} R^{N_nuc x 2(Lx1)^2 x N_channels}. Can you please check this?\n\nThank you for pointing this out, we will fix it in the final manuscript. \n\n[1] Gao et al., arXiv:2110.05064, ICLR (2022)\n\n[2] Gerard at al., arXiv:2205.09438, NeurIPS (2022).\n\n[3] Li et al., arXiv:2307.08214, (2023)."}]

Author rebuttal:
We thank all reviewers for their constructive feedback and appreciate the overall positive reception of our work (“excellent accuracy on challenging systems”, “extensive Empirical Study”), as well as its presentation (“experiments and figures are high quality”, “clear and well-written”).

Several reviewers have asked for additional experiments, which can be found in the attached rebuttal pdf. We ran the following additional experiments, which we will include in the appendix of the revised manuscript:

- Fig 1: Analyzing the energy variance of our wavefunction instead of the mean energy (reviewer BduG) shows the same behavior: Our zero-shot evaluations obtain substantially lower variance than the previous zero-shot-SOTA.
- Fig 2: Fine-tuning of our model for more epochs (reviewer R7Vm), shows that our relative energies do indeed improve further with additional fine-tuning and converge to accurate values.
- Fig 3: Our model yields rotation invariant energies through averaging energies of randomly rotated geometries during inference without additional computational cost
- Tab 1: Run-time comparison. We find that our method is on par with CCSD(T) in terms of accuracy and speed, and can be ~5x faster than Psiformer.

We have answered each reviewer individually, and additionally comment on three recurring themes, which we will also highlight better in the final manuscript.

### Novelty

We believe our contributions to be substantial and novel in multiple regards: We obtain for the first time deep-learning-based VMC-energies for small molecules that are competitive with CCSD(T) in both accuracy as well as run-time. We improve energy errors for zero-shot evaluations by 300x (Fig. 2a) and improve energy errors for fine-tuned wavefunctions by 5x (Fig. 2b) compared to the previous SOTA. These large improvements are obtained through a plethora of contributions:
- Improved electron embedding architecture, reducing energy error by 50% vs Scherbela et al.
- Hessian-based sampling strategy, generating a chemically diverse but physically plausible pre-training dataset
- Iterative Hartree-Fock replaced by adapted version of PhisNet. While the original PhisNet was only trained on geometries of a single molecule, we manage to train on a diverse dataset through various stability improvements, detailed in Appendix F.
- Improved electron position initialization, to reduce burn-in time and computational cost
- Thorough analysis of strengths and limitations of the resulting model, establishing a strong benchmark for further research

### Computational cost 

Like all other deep-learning-based VMC (DL-VMC) methods, our method is computationally costly. However, to our knowledge this is the first study that achieves high accuracy at computational cost comparable to established quantum chemistry methods such as CCSD(T). Computing the energy of a single molecule with 5 heavy atoms requires 10.5 node-hours with CCSD(T)-4Z and Psiformer requires 53h (for 16k epochs). Using our method requires 1h for zero-shot evaluation or 9.5h for 4k fine-tuning steps + evaluation.
We thus reach the same accuracy as CCSD(T)-4Z in the same amount of time, and are 6x faster than Psiformer.
Furthermore a very recent preprint by Li. et al [1] has demonstrated that DL-VMC can be sped up by more than an order of magnitude, via a faster evaluation of the Laplacian operator in loss calculation. These per-step speed-ups could in principle be directly combined with our reduction in the required number of steps, yielding even higher speed-ups.

### E(3)-equivariance

Like all other neural wavefunction architectures known to us (incl. Psiformer, DeepErwin, GLOBE), also our architecture is not equivariant.
While E(3)-equivariant networks have been successful in neural network potentials (e.g. Nequip, MACE), they appear inapplicable to neural wavefunctions. This is because the orbitals used to construct the Slater determinant can have lower symmetry than the Hamiltonian, i.e. be symmetry-broken. By design, E(3)-equivariant networks output functions which are symmetry compatible with the input and are thus not suited to model symmetry broken orbitals. Constructing the wavefunction from E(3)-equivariant orbitals, leads to an invariant wavefunction as shown in Appendix B of Gao et al 2022 [2].

Furthermore, Pozdnyakov and Ceriotti demonstrate in a concurrent pre-print [3] that highly accurate neural network potentials, which outperform the E(3)-equivariant SOTA, can be obtained without explicit E(3)-equivariance. They achieve this with a non-equivariant architecture and data augmentation during training. While E(3)-equivariance is a very elegant concept, it might ultimately not even be required to obtain accurate wavefunctions.

Our approach is thus as follows: All parts of the network, which purely depend on nuclear coordinates (ie. PhisNet) are equivariant. All parts of the network which depend on electron coordinates (requiring symm. breaking) are not equivariant. We obtain approximately rotation invariant energies through data augmentation, by randomly rotating geometries during variational pre-training. While this procedure does not yield fully rotationally invariant energies, our experiment on rotated propadiene (Fig. 3c) demonstrates that the energies are almost invariant (max. zero-shot deviation 5mHa) and improve via fine-tuning. 

We can fully enforce energy invariance, by also using data augmentation during evaluation (rebuttal Fig 3). Instead of evaluating the energy for a fixed geometry, we can sample energies for randomly rotated molecules during inference. Since we require Monte Carlo sampling of many different electron configurations anyway, this comes at no additional computational cost and yields rotation invariant energies up to Monte Carlo noise.

[1] Li et al., arXiv:2307.08214, (2023).

[2] Gao et al., arXiv:2110.05064, ICLR (2022)

[3] Pozdnyakov et al., arXiv:2305.19302, (2023).


Meta-review: This paper considers an approach to deal with the large computational cost of deep learning-base variational Monte Carlo (DL-VMC) methods. To do so, the paper proposes a DL-VMC model that is pre-trained on a large set of molecules and show that a small number of fine-tuning steps can lead to accurate energies in a number of test systems. While the reviewers have pointed out a number of limitations in this work, overall reviewers appreciated the empirical results of the paper and came to a consensus to accept the paper. For the final version, we strongly encourage the authors to revise the manuscript according to the reviewers' feedback.
Decision: Accept (poster)