Title: Variational Monte Carlo on a Budget — Fine-tuning pre-trained Neural Wavefunctions
Venue: NeurIPS 2023
Accepted: yes
Keywords: Computational Physics

Machine Learning for Science

Quantum Monte Carlo

Fermionic Neural Networks
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

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.

1 Introduction

Solving the Schrödinger equation is of utmost importance for the prediction of quantum chemical properties in chemistry. The time-independent Schrödinger equation in the Born-Oppenheimer approximation [1] for a molecule with N nuc nuclei and n el electrons is an eigenvalue problem with Hamiltonian H :

By R = ( R 1 , . . . , R N nuc ) ∈ R N nuc × 3 and Z = ( Z 1 , . . . , Z N nuc ) ∈ N N nuc we denote the nuclear positions and charges. The electron positions are denoted by r = ( r 1 , . . . , r n ↑ , . . . , r n el ) ∈ R n el × 3 with n ↑ spin-up electrons and n ↓ spin-down electrons. The inter-particle difference and distance vectors are written as R IJ = R I -R J , R IJ = | R IJ | , r iI = r i -R I , r iI = | r iI | , r ij = r i -r j and r ij = | r ij | with I, J = 1 , . . . , N nuc and i, j = 1 , . . . , n el. The eigenvalues E of Eq. 1 represent

Leon Gerard *

the energy states of a molecule, whereas a special interest lies in finding the smallest eigenvalue E 0 , called the ground-state energy. The corresponding high-dimensional wavefunction ψ : R n el × 3 → R can be found via the Rayleigh-Ritz principle, by minimizing

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: ψ ( r ) = -ψ ( P r ) . Having access to the solution ψ , allows in principle a complete description of the considered molecule. Unfortunately, only for one electron systems there exists an 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 n el. For example CCSD(T) coupled cluster with its single-, double-, and perturbative triple-excitations variant - is considered the gold-standard reference in computational chemistry, but its computational cost scales as O ( n el 7 ) [2]. Deep-learning-based Variational Monte Carlo (DL-VMC) has emerged as a promising alternative solution. A single step scales as O ( n el 4 ) and it has surpassed the accuracy of many conventional methods such as CCSD(T), when applied to small molecules [3]. In DL-VMC, the wavefunction is represented by a neural network with trainable parameters θ and optimized via Eq. 2 using gradient based optimization. Since the expectation value of Eq. 2 cannot be computed analytically, it is approximated by sampling the electron positions during optimization and evaluation with Monte Carlo methods like Metropolis-Hastings [4].

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-7] and accelerate convergence [8]. Furthermore, DL-VMC has been extended to properties beyond energies [9-11] and systems beyond molecules [12-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-VMC 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 to develop proxy methods, which can be trained on outputs of ab-initio methods to directly predict molecular properties [18, 19] or wavefunctions [20, 21] 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 wavefunctions 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 orbitals of a wavefunction, by using chemistry-inspired heuristics to determine orbital positions, managing to efficiently generalize wavefunctions across different geometries of a single molecule. However, when learning a 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 a low-accuracy conventional method to parameterize DL-VMC orbitals. However, while their wavefunction ansatz successfully transfers to new molecules, their method requires a separate, iterative Hartree-Fock (HF) calculation for every new geometry.

Our contribution 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:

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. [25] 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 a 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-accuracy gold-standard reference methods, such as CCSD(T)-3Z on unseen systems without any finetuning (cf. Fig. 2a).

Overview of the paper In Sec. 2 we outline our method and the procedure to optimize a transferable wavefunction across molecules. In Sec. 3 we thoroughly test the accuracy of our obtained wavefunction, by analyzing absolute energies (Sec. 3.1), relative energies (Sec. 3.2), and the impact of design choices in an ablation study (Sec. 3.3). Throughout this work, we compare against other high-accuracy methods, in particular results obtained by state-of-the-art DL-VMC methods and CCSD(T). Finally we analyze the scalability and limitations of our base model, by applying it to a large-scale dataset in Sec. 3.4, before a discussion and outlook for future research in Sec. 4.

2 Methods

Our approach is divided into two parts (cf. Fig. 1): On the one hand, a wavefunction ansatz, containing an electron embedding, orbital embedding and a Slater determinant. On the other hand a method for geometry sampling based on normal-mode distortions.

2.1 Our wavefunction ansatz

PhisNet

R

Z

orbital

coeffs

electron

embeddings

orbital

embeddings

E

nuclear embedding

nuc. feat.

r

orbital

emb. net

nuclear

GNN

electron

emb. net

Hamil-

tonian

energy estimate

normal

distribution

hessian estimate

orbital model

wavefunction model

pretrained, then frozen

fixed

trainable

A single forward-pass for our wavefunction model

can again be divided into three blocks: An electron model acting on nuclear and electron coordinates, generating electron embeddings x el ; an SE(3)-equivariant orbital model acting only on nuclear coordinates, generating orbital embeddings g orb and f orb ; a Slater determinant combining electronand orbital-embeddings, 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-nuclear and nuclear-nuclear edges are embedded with a multi-layer perceptron (MLP),

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

for a receiving particle ˜ a rec i and the set of sending particles { a send j } , connected via their edges { e ij } . By σ we denote the non-linear activation and with ⊙ the element-wise multiplication along the feature dimension. An MPNN is obtained by stacking MessagePassing layers

In the following we use these message passing steps to model all inter-particle interactions.

SE(3)-equivariant orbital model The orbital model is a simplified version of PhisNet[20], a neural network predicting the overlap matrix S and the Fock matrix F , via nuclear embeddings x nuc :

Here x nuc ∈ R N nuc × ( L +1) 2 × N channels , and S IJ and F IJ are each in R N basis × N basis . The basis-set size of the predicted orbitals is denoted by N basis and the feature dimension of the nuclear embeddings by N channels . The full overlap- and Fock-matrices are assembled from the corresponding blocks S IJ and F IJ , leading to matrices of shape [ N nuc N basis × N nuc N basis ] . Each layer of PhisNet is SE(3)-equivariant, ensuring that any 3D-rotation or inversion of the input coordinates R , leads to an equivalent rotation of its outputs. This is done by splitting any feature vector into representations of varying harmonic degree l = 0 , . . . , L , each with components m = -l, . . . , l . A detailed description of SE3-equivariant networks in general, as well as PhisNet in particular can be found in [20]. A list of changes and simplifications we made to PhisNet can be found in Appendix C.

The orbital embeddings (corresponding to orbital expansion coefficients in a conventional quantum chemistry calculation) are obtained by solving the generalized eigenvalue problem:

It has been shown empirically that it is beneficial to obtain the orbital coefficients ˆ x orb Ik as 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 F and S on the other hand are unique and transform equivariantly under E(3)-transformations of the molecule, leading to a well defined learning problem.

Following [25], we do not use the orbital energies ϵ k and obtain the backflow factors f 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.

Electron model The electron embedding is a message passing neural network. To incorporate the geometric information of the molecule considered, we leverage the equivariant prediction of the PhisNet nuclear embeddings x nuc , by first performing a message passing step between the nuclear embeddings

and then using these features to initialize the electron embeddings

by using a zero vector 0 for the initial receiving electrons. This differentiates our electron model from previously proposed methods [3, 6, 8], leading to a better generalization when optimized across molecules (cf. Sec. 3.3). The final step of the embedding is a multi-iteration message passing between electron embeddings to capture the necessary electron-electron interaction

resulting in a N emb-dimensional representation for each electron x el i ∈ R N emb , i = 1 , . . . , n el .

Overall E(3)-equivariance Like existing approaches [3, 8, 17] our overall wavefunction does not enforce E(3)-symmetry. This is because the wavefunction can have lower symmetry than the molecule [23], for example in the case of the excited states of a hydrogen atom. We therfore choose all parts of the network that act purely on nuclear coordinates (i.e. the PhisNet model and the energy/hessian estimate) to be equivariant under E(3)-transformations. 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 not 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 r from the probability distribution ψ ( r ) 2 , to evaluate the expectation value of Eq. 2. One notable difference compared to other works is our initialization r 0 of the Markov Chain. In the limit of infinite steps, the samples are distributed according to ψ 2 , but for a finite number of steps the obtained samples strongly depend on r 0 . This issue is typically addressed by a "burn-in", where MCMC is run for a fixed number of steps (without using the resulting samples) to ensure that r has diffused to state of high probability. Previous work has initialized r 0 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 ≈ 10 5 MCMCsteps to reach the equilibrium distribution. We instead initialize r 0 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 a labeled dataset of reference energies, but to obtain a transferable wavefunction, which generalized well to new systems, a diverse dataset of molecular geometries R is required. Starting with an initial set of geometries R 0 , we update R on the fly, by perturbing each geometry every 20 optimization steps by adding random noise ∆ R to the nuclear coordinates. Using uncorrelated, isotropic random noise for ∆ R would yield many non-physical 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 ∆ R from a correlated normal distribution

The bias term α ( R 0 -R ) ensures that geometries stay sufficiently close to their starting point R 0 . The covariance matrix is chosen proportional to the pseudo-inverse of the hessian of the energy E phis , which is predicted from the scalar component of the nuclear embeddings using a pre-trained MLP. By using H -1 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.

After distorting the nuclear coordinates R , we also adjust the electron positions r , using the spacewarp coordinate transform described in [29], which effectively shifts the electrons by a 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 a more diverse dataset.

2.3 Optimization

To obtain orbital descriptors (and energies to 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 a 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. A full DL-VMC calculation to obtain a ground-state energy prediction can be divided into three consecutive steps:

Supervised optimization using PhisNet : Initially, the neural-network orbitals (cf. Eq. 3) 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.

Variational optimization : Minimization of the energy (cf. Eq. 2) by drawing samples from ψ 2 using MCMC and updating the wavefunction parameters using the KFAC optimizer [30].

Evaluation : For inference of the ground-state energy, we sample electron positions using MCMC, and evaluate the energy using Eq. 2 without updating θ .

To train a multi-geometry transferable wavefunction we further divide the variational optimization of a neural wavefunction into two steps:

Pre-training : Asingle wavefunction model is trained across many molecules and geometries. In every gradient step we only consider a 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 we 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".

Fine-tuning : A a small number of additional variational optimization steps is done using geometries of interest, starting from the weights of a 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 we demonstrate the performance of this model, for zero-shot evaluations and after subsequent finetuning.

3.1 Accuracy of pre-trained model for absolute energies

To test the transfer capabilities of the model, we evaluate it on test-sets, which each contain 4 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 D). 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 7 heavy atoms. We find that for molecules containing up to 6 heavy atoms our method outperforms CCSD(T) with a 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 a Hartree-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-Fock in a zero-shot setting. In contrast, our improvements to their method increase zero-shot accuracy by more than 2 orders of magnitude.

3

4

5

6

7

Nr. of heavy atoms

0

10 1

10 2

10 3

10 4

10 5

E

E

CCSD(T)

CBS / mHa

HF-CBS

CC-2Z

CC-3Z

CC-4Z

CC-CBS

a

Zero-shot

330x

3

4

5

6

7

Nr. of heavy atoms

HF-CBS

CC-2Z

CC-3Z

CC-4Z

CC-CBS

b

After 4k fine-tuning steps

5.0x

0

1k

2k

4k

8k

16k

32k

Optimization steps

HF-CBS

CC-2Z

CC-3Z

CC-4Z

CC-CBS

c

Test set with 3 heavy atoms

Our work (no pre-training)

Our work (pre-trained)

PsiFormer

We furthermore find that fine-tuning pre-trained wavefunctions can be more cost effective than well established conventional methods. For a typical molecule containing 5 heavy atoms, we require 9.5 node-hours for 4k fine-tuning steps, achieving accuracy that surpasse CCSD(T)-4Z. 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 of Appendix I.

To further test the accuracy of our method, we evaluate for 3 heavy atoms the performance with increasing fine-tuning steps (cf. Fig. 2c). 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 up 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 to optimize across molecule, yields less expressive wavefunctions than a fully trainable backflow. This is demonstrated by the fact that a pre-trained and non-pre-trained model converge to 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 a small number of fine-tuning steps is required to reach quantitatively correct relative energies. Fig. 3 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 F.

Bicyclobutane conformers Fig. 3a depicts the energies of 5 conformers of bicyclobutane relative to the energy of its initial structure. The system is of interest, because CCSD(T) severely underestimates the energy of the dis-TS conformer by ≈ 60 mHa [31]. 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 we overestimate the relative energy

dis_TS

con_TS

gt-TS

g-but

t-but

50

0

50

100

150

E

E

Bicyclobutane / mHa

a

Conformers of bicyclobutane

Our work: 0-shot

Our work: 700 steps per geom.

DMC (Kinal et al.)

FermiNet: 200k steps per geom.

FermiNet: 10k steps per geom.

CCSD(T) (Kinal et al.)

2

3

4

5

bond length / bohr

0

100

200

300

400

500

600

E

E

0 / mHa

b

N

2 dissociation

Our work (0-shot)

Our work (4k fine-tuning)

0

100

200

300

rotation angle / deg

2

0

2

4

6

8

10

E

E

0 / mHa

c

Our work (0-shot)

Our work (0-shot w/ rot. avg.)

Our work (4k fine-tuning)

exact

0

20

40

60

80

torsion angle / deg

0

20

40

60

80

100

E

E

0 / mHa

d

Our work (0-shot)

Our work (4k fine-tuning)

CCSD(T)/CBS

B.

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 our batch size; max. deviation 7.5 mHa), and slightly less accurate than a FermiNet calculation optimized for 200k steps (max. deviation 1.4 mHa). As opposed to CCSD(T) our model does not suffer from systematic errors, even 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 N 2 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 or large bond lengths (cf. Fig. 3b). This is due to the lack of dissociated atoms in the pre-training dataset and again mostly remedied by finetuning. In the most challenging regime around d = 3 . 7 , even with fine-tuning we obtain energies that are 12 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 (C 3 H 4 ) rotated around a random axis, we find typical energy variations of ± 1 mHa, but also a individual outliers, deviating by up to 5 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 rotation-invariant energies by data augmentation, appears to be helpful but not to be fully sufficient. When evaluating an earlier checkpoint of our base model (trained for 170k epochs instead of 256k), we observe mean energy variations of ± 3mHa and outliers of up to 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 to reduce errors below chemical accuracy of 1.6 mHa.

In addition to augmenting data during training time, we can also augment the data during inference time, to obtain fully rotation-invariant energies, from an 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. 3c shows that with rotation averaging we obtain energies that are fully invariant, up to Monte-Carlo 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 D). This seems sufficient for 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 a maximum deviation of 2 mHa along the path.

3.3 Ablation studies

10

20

50

100

200

500

1000 2000

5000

E

E

5620

-4260

(-76%)

134

(+10%)

-1470

(-98%)

1.5

(+5%)

-12.5

(-37%)

21.5

128k pre-training

+ improved

architecture

+ normal-mode

distortions

+ larger

training set

+ 8 determinants

+ 256k pretraining-

steps

Our work

a

Zero-shot accuracy

0

2

5

8

10

12

15

18

20

E

E

17.8

(-35%)

-1.4

(-12%)

-0.9

(-9%)

-0.7

(-8%)

-1.8

(-21%)

6.8

b

Accuracy after 4000 fine-tuning steps

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 8 and increasing the number of pre-training steps from 128k to 256k improve fine-tuning accuracy by 8% and 21% respectively. The two largest contributions overall are the improved architecture - yielding substantial gains both in zero-shot and fine-tuning - as well as the larger training set, which is crucial for high zero-shot accuracy. 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 we pre-train the PhisNet model against Hartree-Fock 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-trained nuclear emebddings x nuc 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 a subset of 250 molecules from the QM7 dataset [32], containing molecules with 14-58 electrons.

Since CCSD(T) calculations would be prohibitively expensive for a dataset of this size, in Fig. 5a we compare against density functional theory (DFT) reference calculations (PBE0+MBD, [26]). Fig. 5b breaks down the energy differences compared to DFT, by molecule size. For molecules with less than 5 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 a 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.1. One reason for the large energy differences in larger molecules is the increasing inter-atomic distance within the molecules with increasing number of atoms (cf. Fig. 5c). 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 a 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.

300

200

100

E

PBE0+MBD

/ Ha

300

200

100

E / Ha

a

0

500

1000

E -

E

PBE0+MBD

/ mHa

0

25

50

75

100

125

Count

b

Nr. of heavy atoms

2

3

4

5

6

7

8

12

16

RIJ

/ bohr

10

2

0

10

2

10

3

E -

E

PBE0+MBD

/ mHa

R train

IJ

c

4 Discussion

We have presented to our knowledge the first ab-initio wavefunction model, which achieves highaccuracy zero-shot energies on new systems (Sec. 3.1 and Sec. 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 fine-tuning, but need on the order of 4000 fine-tuning steps to reach chemical accuracy of 1.6 mHa (Sec. 3.2). This is a substantial improvement over previous work, which so far has fallen in two categories: High-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 ([17, 25]), but yield orders of magnitude lower accuracy in 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 a larger training dataset.

While results are 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)-symmetry of the wavefunction should be explored further, since currently only the orbital part of our architecture is SE(3)-equivariant. We experimented with a 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-propagation through the iterative orbital localization procedure. Optimizing these weights in addition to the electron embedding will lead to a more expressive ansatz.

References

M. Born and R. Oppenheimer. 'Zur Quantentheorie der Molekeln'. In: Annalen der Physik 389.20 (1927), pp. 457-484. DOI: https://doi.org/10.1002/andp.19273892002 .

László Gyevi-Nagy, Mihály Kállay, and Péter R. Nagy. 'Accurate Reduced-Cost CCSD(T) Energies: Parallel Implementation, Benchmarks, and Large-Scale Applications'. In: Journal of Chemical Theory and Computation 17.2 (Feb. 2021), pp. 860-878. DOI: 10.1021/acs.jctc. 0c01077 .

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), p. 033429. DOI: 10.1103/ PhysRevResearch.2.033429 .

W. K. Hastings. 'Monte Carlo Sampling Methods Using Markov Chains and Their Applications'. In: Biometrika 57.1 (Apr. 1970), pp. 97-109. DOI: 10.1093/biomet/57.1.97 .

James S. Spencer et al. 'Better, Faster Fermionic Neural Networks'. In: (2020).

Ingrid von Glehn, James S. Spencer, and David Pfau. A Self-Attention Ansatz for Ab-initio Quantum Chemistry . Nov. 2022. DOI: 10.48550/arXiv.2211.13672 .

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.1038/s41467023-37609-3 .

Leon Gerard et al. 'Gold-Standard Solutions to the Schrödinger Equation Using Deep Learning: How Much Physics Do We Need?' In: Advances in Neural Information Processing Systems . Oct. 2022.

M. T. Entwistle et al. 'Electronic excited states in deep variational Monte Carlo'. In: Nature Communications 14.1 (Jan. 2023), p. 274. DOI: 10.1038/s41467-022-35534-5 .

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. DOI: 10.1038/s43588-022-00228-x .

Yubing Qian et al. 'Interatomic Force from Neural Network Based Variational Quantum Monte Carlo'. In: The Journal of Chemical Physics 157.16 (Oct. 2022), p. 164104. DOI: 10.1063/5.0112344 .

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

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. DOI: 10.1038/s41467-022-356271 .

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. DOI: 10.1126/science.aag2302 .

Giuseppe Carleo et al. 'NetKet: A Machine Learning Toolkit for Many-Body Quantum Systems'. In: SoftwareX 10 (July 2019), p. 100311. DOI: 10.1016/j.softx.2019.100311 .

Yuan-Hang Zhang and Massimiliano Di Ventra. 'Transformer Quantum State: A Multipurpose Model for Quantum Many-Body Problems'. In: Physical Review B 107.7 (Feb. 2023), p. 075147. DOI: 10.1103/PhysRevB.107.075147 .

Nicholas Gao and Stephan Günnemann. Generalizing Neural Wave Functions . Feb. 2023. DOI: 10.48550/arXiv.2302.04168 .

Simon Batzner et al. 'E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials'. In: Nature Communications 13.1 (May 2022), p. 2453. DOI: 10.1038/ s41467-022-29939-5 .

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.

Oliver Unke et al. 'SE(3)-Equivariant Prediction of Molecular Wavefunctions and Electronic Densities'. In: Advances in Neural Information Processing Systems . 2021.

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. DOI: 10.1063/5.0012911 .

Albert Musaelian et al. 'Scaling the Leading Accuracy of Deep Equivariant Models to Biomolecular Simulations of Realistic Size'. In: (Apr. 2023). DOI: 10.48550/arXiv.2304. 10061 . arXiv: 2304.10061 .

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

Nicholas Gao and Stephan Günnemann. 'Sampling-free Inference for Ab-Initio Potential Energy Surface Networks'. In: (2022). DOI: 10.48550/ARXIV.2205.14962 .

Michael Scherbela, Leon Gerard, and Philipp Grohs. 'Towards a Foundation Model for Neural Network Wavefunctions'. In: arXiv.org (Mar. 2023).

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.1038/s41597-021-00812-2 .

Ruard van Workum, Joao Malhado, and Philipp Marquetand. CASNet: Learning Complete Active Space Orbitals Using Message Passing Neural Networks . June 2023. DOI: 10.26434/ chemrxiv-2023-lwj87 .

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. 4916-4926. DOI: 10.1063/1.456588 .

Claudia Filippi and C. J. Umrigar. 'Correlated Sampling in Quantum Monte Carlo: A Route to Forces'. In: Physical Review B 61.24 (June 2000), R16291-R16294. DOI: 10.1103/ PhysRevB.61.R16291 .

James Martens and Roger Grosse. 'Optimizing neural networks with kronecker-factored approximate curvature'. In: International conference on machine learning . PMLR. 2015, pp. 2408-2417.

Armaˇ gan Kinal and Piotr Piecuch. 'Computational Investigation of the Conrotatory and Disrotatory Isomerization Channels of Bicyclo[1.1.0]butane to Buta-1,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. DOI: 10.1021/ jp065721k . URL: https://doi.org/10.1021/jp065721k .

Matthias Rupp et al. 'Fast and Accurate Modeling of Molecular Atomization Energies with Machine Learning'. In: Phys. Rev. Lett. 108 (5 Jan. 2012), p. 058301. DOI: 10.1103/ PhysRevLett.108.058301 .

Kristof T. Schütt et al. 'SchNet: A Continuous-Filter Convolutional Neural Network for Modeling Quantum Interactions'. In: (June 2017). DOI: 10.48550/arXiv.1706.08566 .

Gabriel Pescia et al. Message-Passing Neural Quantum States for the Homogeneous Electron Gas . May 2023. DOI: 10.48550/arXiv.2305.07240 .

Qiming Sun et al. 'Recent Developments in the PySCF Program Package'. In: The Journal of Chemical Physics 153.2 (July 2020), p. 024109. DOI: 10.1063/5.0006074 .

Mario Geiger and Tess Smidt. E3nn: Euclidean Neural Networks . July 2022. arXiv: 2207. 09453 [cs] .

Diederik P. Kingma and Jimmy Ba. 'Adam: A Method for Stochastic Optimization'. In: arXiv:1412.6980 (Jan. 2017). DOI: 10.48550/arXiv.1412.6980 . arXiv: 1412.6980 [cs] .

Greg Landrum. RDKit: Open-source cheminformatics . https://github.com/rdkit/ rdkit . GitHub repository. 2009.

Frank Neese et al. 'The ORCA Quantum Chemistry Program Package'. In: The Journal of Chemical Physics 152.22 (June 2020), p. 224108. DOI: 10.1063/5.0004608 .

David Pfau James S. Spencer and FermiNet Contributors. FermiNet . 2020. URL: http: //github.com/deepmind/ferminet .

Aleksandar Botev and James Martens. KFAC-JAX . Version 0.0.1. 2022. URL: http://github. com/deepmind/kfac-jax .

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 a glycine molecule, using a pre-trained wavefunction. We perform no initial burn-in and use every 50th sample for energy evaluation. If the chain was perfectly equilibrated 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 is particularly pronounced with Gaussian initialization of electron positions, which is the default in state-of-the-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 ψ 2 (cf. Fig. 6a) and thus reaches equilibrium substantially faster.

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

distance from nucleus / bohr

10 3

10 4

10 5

10 6

Electron density, a.u.

a

Exponential init.

Ground truth

2

0

2

4

6

8

10

MCMC steps / k

285

284

283

282

281

280

Energy / Ha

b

Exponential init.

Ground truth

B Orbital localization

Our model uses orbital embeddings x orb as inputs to parameterize the backflows f orb , and exponents g orb 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-Fock calculation. In this setting, the coefficients x orb are not uniquely defined, but only up to a linear transformation U with determinant ± 1

This stems from the fact that the corresponding Hartree-Fock wavefunction is invariant under such a transformation. Consequently there is free choice, which linear combination of embeddings x orb to choose from without any loss of information. We follow the approach of [25], by choosing U such that the corresponding Hartree-Fock orbitals are maximally localized according to the Foster-Boys

metric, i.e. minimize the spatial variance L :

Here b Iµ ( r ) denotes µ -th basis function of the Hartree-Fock expansion, centered on the I -th nucleus. In practice the integrals of Eq. 17 do 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-source quantum chemistry codes such as pySCF [35].

C Adaption of PhisNet

We heavily rely on PhisNet by Unke et al. [20] to obtain orbital descriptors without the need for a separate SCF calculation. Compared to their original work, we made several simplifications, which are motivated by the fact that we do not predict final high-accuracy 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 an (equivariant) layer norm after each PhisNet module.

Simplified Fock matrix prediction The original PhisNet implementation uses a 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 L max = 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 a single model to predict F , S , E , and ∇ E 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 SE(3)-equivariant operations.

Wetrain the PhisNet-model on a dataset of 47k molecules from QM7X [26], using the Adam optimizer [37] on the following loss

Here E denotes energies, G denotes gradients of energies, F Fock matrices, and S overlap matrices. The indices I, J run over nuclei, the indices µ, ν over basis functions, and the index n over samples in a batch.

D Molecule datasets

Bicyclobutane For the Bicyclobutane to 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 2 For the N 2 potential energy surface with various bond-lengths we used the geometries including reference calculations from [8].

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

Zero-shot and fine-tuning dataset The results on 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 3 heavy atoms we use the ablation dataset.

Ablation dataset For the ablation study, we use one geometry per molecule from the out-ofdistribution test set from Scherbela et al. [25], leading to a set of four distinct molecules. We ensure that these molecules are not part of the training set.

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 wavefunctions we use two datasets, consisting of 18 and 98 disparate molecules. For part of the ablation we use the dataset proposed in [25] and an 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 a train-test leakage, we remove Bicyclobutane (including all conformations) and the four molecules from the ablations dataset. Since the normal-mode-distortions by design do not generate strongly distorted geometries, we augment the 98-molecule-dataset with rotated dihedral angles. To generate a 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. 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 are not underrepresented in the dataset we 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 0 for pre-training.

E Energy variance

Fig. 7 depicts the energy variance obtained by various models. Analogously to the the energies presented in Fig. 2 of the main text, we also find that our pre-trained model achieves substantially lower energy variance compared to 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. 3 of the main text, we depict the difference of relative energies between our model and a references method in Fig. 8. Furthermore we list numerical values vor all energies of Fig. 3a in Tab. 1. We find that on these

3

4

5

6

7

Nr. of heavy atoms

10

1

10

0

10

1

10

2

10

3

EVar

/

Ha

2

a

Zero-shot

610x

3

4

5

6

7

Nr. of heavy atoms

b

After 4k fine-tuning steps

Our work (pre-trained)

7.4x

0

1k

2k

4k

8k

16k

32k

Optimization steps

c

Test set with 3 heavy atoms

Our work (no pre-training)

Our work (pre-trained)

PsiFormer

challenging systems our model yields relative energy errors of up to 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 are accurate within 2-3 mHa.

2

3

4

5

bond length / bohr

5

0

5

10

15

20

25

E

/ mHa

a)

N

2 dissociation

Our work (4k fine-tuning steps)

Our work (16k fine-tuning steps)

Our work (32k fine-tuning steps)

4

8

12

16

20

24

28

32

Optimization steps / k

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

E

/ mHa

b)

Our work

PsiFormer

G Reference energies

CCSD(T) All CCSD(T) energies - except explicitly stated otherwise - were obtained using ORCA [39] starting from a restricted Hartree-Fock calculation. We use correlation consistent basis sets of the cc-pCVXZ family, with X in {2, 3, 4}. To extrapolate to the complete basis set limit (CBS), we use

the approach outlined in [3] and fit the following functions with free parameters E HF CBS , E corr CBS , a, b, c :

We stress that although CCSD(T)-energies are often considered as "gold-standard", they do not necessarily represent the actual ground-state energy. There are 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 true ground-state energy.

PsiFormer For Fig. 2 we used the open-source FermiNet codebase [40]. The codebase didn't allow for inference calculation, therefore a slight fix was applied. All calculations were performed with the small settings as proposed in von Glehn et al. [6].

H Hyperparameters

A detailed description of the hyperparameter used in this work can be found below (cf. Tab. 2). For the mapping of 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 to perform a 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. 3 lists the run-times for our method and a comparable 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) at similar computational cost. This is in contrast to many earlier works that achieve highly accurate energies, but at computational cost that far surpasses the cost of conventional methods for small molecules.

Overall we used ≈ 5kGPUhs (A100) for development and training of our base models, and another 5k GPUhs (A40) on evaluations and finetuning. Additionally we required ≈ 20kCPUhs for CCSD(T) reference calculations.

J Code and data availability

All code, configuration files, geometries, datasets and obtained energies are available on GitHub under https://github.com/mdsunivie/deeperwin . Model weights are available on figshare under https://doi.org/10.6084/m9.figshare.23585358.v1 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nWhile 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\nQUESTIONS:\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\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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nQuestions 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?\n\nLIMITATIONS:\nCompared 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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- The overall end-to-end training is built upon previous works.\n\nQUESTIONS:\n- 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.)\n\nLIMITATIONS:\nN.A.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}

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

STRENGTHS:
- The authors provide detailed description of the architecture.
-  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.
- The proposed method does not yield improved results when more optimization steps are used.
- The presentation of the training and evaluation steps is not clear. Specifically, how is pre-training and fine-tuning done.

QUESTIONS:
- Among the many experimented datasets, which one(s) most decisively show the benefit of the proposed method?

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI'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).\n\nQUESTIONS:\nFrom 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?\n\nLIMITATIONS:\nThe authors are upfront about the limitations: potential limited expressiveness with more optimization iterations, prediction is worse for increasing extrapolation and out-of-distribution predictions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-24"}