SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $47\mathrm{k}$  molecules. Further details of the loss function, dataset, and the adoptions 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:
1. 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.
2. Variational optimization: Minimization of the energy (cf. Eq. 2) by drawing samples from  $\psi^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  $\theta$ .
To train a multi-geometry 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 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".
2. Fine-tuning: 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.
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 fine-tuning.
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.
Figure 2: Absolute energies: Energies relative to CCSD(T)-CBS (complete basis set limit) when reusing 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 a 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 a typical molecule containing 5 heavy atoms, we require 9.5 node-hours for 4k fine-tuning steps, achieving accuracy that surpasses 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  $\approx 60\mathrm{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
Figure 3: Challenging relative energies: Relative energies obtained with and without fine-tuning on 4 distinct, challenging systems, compared against high-accuracy reference methods. a) Relative energy of bicyclobutane conformers vs. the energy of bicyclobutane b) Potential energy surface (PES) of  $\mathbf{N}_2$ , c) global rotation of propadiene d) relative energy of twisted vs. untwisted propadiene.
by  $90\mathrm{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\mathrm{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\mathrm{mHa}$ ), and slightly less accurate than a FermiNet calculation optimized for 200k steps (max. deviation  $1.4\mathrm{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\mathrm{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  $(\mathrm{C}_3\mathrm{H}_4)$  rotated around a random axis, we find typical energy variations of  $\pm 1\mathrm{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  $\pm 3\mathrm{mHa}$  and outliers of up to  $30\mathrm{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\mathrm{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  $\mathrm{C = C}$  bonds of propadiene leads to a transition state with an energy difference of  $110\mathrm{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):  $\sim 0.2$  mHa discrepancy for the barrier height and a maximum deviation of 2 mHa along the path.
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 3 heavy atoms. Panel a 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 a 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  $128\mathrm{k}$  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  $128\mathrm{k}$  to  $256\mathrm{k}$  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  $256\mathrm{k}$  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  $\boldsymbol{x}^{\mathrm{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\mathrm{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.
Figure 5: Zero-shot on QM7: a) Our zero-shot energies vs. DFT [26] b) Histogram of energy residuals (truncated at  $1.25\mathrm{Ha}$  for clarity) c) Residuals vs. largest inter-molecular distance  $R_{IJ}$ .
We have presented to our knowledge the first ab-initio wavefunction model, which achieves high-accuracy 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\mathrm{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.  $10\mathrm{x}$  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.
We gratefully acknowledge financial support from the following grants: Austrian Science Fund FWF Project I 3403 (P.G.), WWTF-ICT19-041 (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.
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  $\psi^2$  (cf. Fig. 6a) and thus reaches equilibrium substantially faster.
Figure 6: Effect of electron initialization: Initializing the electron positions using an exponential distribution instead of Gaussian, better fits the actual density (a), and thus leads to faster equilibration of observables during evaluation (b).
B Orbital localization
Our model uses orbital embeddings  $\pmb{x}^{\mathrm{orb}}$  as inputs to parameterize the backflows  $f^{\mathrm{orb}}$ , and exponents  $g^{\mathrm{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  $\pmb{x}^{\mathrm{orb}}$  are not uniquely defined, but only up to a linear transformation  $U$  with determinant  $\pm 1$
$$
\boldsymbol {x} _ {I k} ^ {\mathrm {o r b}} = \sum_ {n = 1} ^ {N _ {\mathrm {o r b}}} U _ {k n} \hat {\boldsymbol {x}} _ {I k} ^ {\mathrm {o r b}}, \quad U \in \mathbb {R} ^ {N _ {\mathrm {o r b}} \times N _ {\mathrm {o r b}}}, \quad \det  U = \pm 1. \tag {15}
$$
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  $\pmb{x}^{\mathrm{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  $\mathcal{L}$ :
$$
\phi_ {k} (\boldsymbol {r}, U) = \sum_ {I = 1} ^ {N _ {\mathrm {n u c}}} \sum_ {\mu = 1} ^ {N _ {\mathrm {b a s i s}}} b _ {I \mu} (\boldsymbol {r}) U _ {k n} \hat {x} _ {I n \mu} ^ {\mathrm {o r b}} \tag {16}
$$
$$
\mathcal {L} (U) = \sum_ {k} \int \phi_ {k} ^ {2} (\boldsymbol {r}, U) \boldsymbol {r} ^ {2} d \boldsymbol {r} - \left(\int \phi_ {k} ^ {2} (\boldsymbol {r}, U) \boldsymbol {r} d \boldsymbol {r}\right) ^ {2} \tag {17}
$$
Here  $b_{I\mu}(\boldsymbol{r})$  denotes  $\mu$ -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  $\mathcal{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 ThisNet
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 ThisNet 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 ThisNet which obtained sufficient accuracy for our purposes. We used 2 layers (instead of 5) and  $L_{\mathrm{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  $\nabla 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.
We train the PhisNet-model on a dataset of 47k molecules from QM7X [26], using the Adam optimizer [37] on the following loss