SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Variational Monte Carlo on a Budget – Fine-tuning pre-trained Neural Wavefunctions
Michael Scherbela*
University of Vienna  
michael.scherbela@univie.ac.at
Leon Gerard*  
University of Vienna  
leon.gerard@univie.ac.at
Philipp Grohs
University of Vienna philipp.grohs@univie.ac.at
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.
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_{\mathrm{nuc}}$  nuclei and  $n_{\mathrm{el}}$  electrons is an eigenvalue problem with Hamiltonian  $H$ :
$$
H \psi = E \psi , \quad H = - \frac {1}{2} \sum_ {i} \nabla_ {\boldsymbol {r} _ {i}} ^ {2} + \sum_ {i > j} \frac {1}{r _ {i j}} + \sum_ {I > J} \frac {Z _ {I} Z _ {J}}{R _ {I J}} - \sum_ {i, I} \frac {Z _ {I}}{r _ {i I}}. \tag {1}
$$
By  $\pmb{R} = (R_{1},\dots,R_{N_{\mathrm{nuc}}})\in \mathbb{R}^{N_{\mathrm{nuc}}\times 3}$  and  $\pmb {Z} = (Z_{1},\dots,Z_{N_{\mathrm{nuc}}})\in \mathbb{N}^{N_{\mathrm{nuc}}}$  we denote the nuclear positions and charges. The electron positions are denoted by  $\pmb {r} = (\pmb {r}_1,\dots,\pmb {r}_{n_\uparrow},\dots,\pmb {r}_{n_\mathrm{el}})\in \mathbb{R}^{n_\mathrm{el}\times 3}$  with  $n_\uparrow$  spin-up electrons and  $n_\downarrow$  spin-down electrons. The inter-particle difference and distance vectors are written as  $\pmb{R}_{IJ} = \pmb{R}_I - \pmb{R}_J$ ,  $R_{IJ} = |\pmb{R}_{IJ}|$ ,  $\pmb{r}_{iI} = \pmb{r}_i - \pmb{R}_I$ ,  $r_{iI} = |\pmb{r}_{iI}|$ ,  $\pmb{r}_{ij} = \pmb{r}_i - \pmb{r}_j$  and  $r_{ij} = |\pmb{r}_{ij}|$  with  $I,J = 1,\dots,N_{\mathrm{nuc}}$  and  $i,j = 1,\dots,n_{\mathrm{el}}$ . The eigenvalues  $E$  of Eq. 1 represent
*Equal contribution, author order random
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  $\psi : \mathbb{R}^{n_{\mathrm{el}} \times 3} \to \mathbb{R}$  can be found via the Rayleigh-Ritz principle, by minimizing
$$
\mathcal {L} (\psi_ {\theta}) = \mathbb {E} _ {\boldsymbol {r} \sim \psi_ {\theta} ^ {2} (\boldsymbol {r})} \left[ \frac {H \psi_ {\theta} (\boldsymbol {r})}{\psi_ {\theta} (\boldsymbol {r})} \right] \geq E _ {0}. \tag {2}
$$
Due to electrons being fermions, the solution must fulfill the anti-symmetry property, stating that the sign of the wavefunction must change for any permutation  $\mathcal{P}$  of two electrons of the same spin:  $\psi (\boldsymbol {r}) = -\psi (\mathcal{P}\boldsymbol {r})$ . Having access to the solution  $\psi$ , 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_{\mathrm{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  $\mathcal{O}(n_{\mathrm{el}}^7)$  [2]. Deep-learning-based Variational Monte Carlo (DL-VMC) has emerged as a promising alternative solution. A single step scales as  $\mathcal{O}(n_{\mathrm{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  $\theta$  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 GPUs to find  $\psi$  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 scalable, 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.
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
Figure 1: Overview of our approach: Wavefunction ansatz (top) and geometry sampling (bottom)
A single forward-pass for our wavefunction model
$$
\psi = \sum_ {d = 1} ^ {N _ {\mathrm {d e t}}} \det  \left[ \phi_ {k} ^ {d} \left(\boldsymbol {x} _ {i} ^ {\mathrm {e l}}\right) \right], \quad \phi_ {k} ^ {d} \left(\boldsymbol {x} _ {i} ^ {\mathrm {e l}}\right) = \sum_ {I = 1} ^ {N _ {\mathrm {n u c}}} \langle \boldsymbol {x} _ {i} ^ {\mathrm {e l}}, \boldsymbol {f} _ {I k d} ^ {\mathrm {o r b}} \rangle e ^ {- r _ {I} g _ {I k d} ^ {\mathrm {o r b}}}, \quad i, k = 1, \dots , n _ {\mathrm {e l}} \tag {3}
$$
can again be divided into three blocks: An electron model acting on nuclear and electron coordinates, generating electron embeddings  $\boldsymbol{x}^{\mathrm{el}}$ ; an SE(3)-equivariant orbital model acting only on nuclear coordinates, generating orbital embeddings  $g^{\mathrm{orb}}$  and  $f^{\mathrm{orb}}$ ; a Slater determinant combining electron- and 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),
$$
e _ {i j} ^ {\text {e l - e l}} = \operatorname {M L P} \left(\left[ \boldsymbol {r} _ {i j}, r _ {i j} \right]\right) \quad e _ {i I} ^ {\text {e l - n u c}} = \operatorname {M L P} \left(\left[ \boldsymbol {r} _ {i I}, r _ {i I} \right]\right) \quad e _ {I J} ^ {\text {n u c - n u c}} = \operatorname {M L P} \left(\left[ \boldsymbol {R} _ {I J}, R _ {I J} \right]\right), \tag {4}
$$
by using a concatenation  $([\cdot])$  of distance and difference vectors, and separate weights for each MLP. A single message passing step is decomposed into the following operations
$$
\begin{array}{l} \tilde {\boldsymbol {a}} _ {i} ^ {\text {r e c}} = \operatorname {M e s s a g e P a s s i n g} \left(\boldsymbol {a} _ {i} ^ {\text {r e c}}, \left\{\boldsymbol {a} _ {j} ^ {\text {s e n d}} \right\}, \left\{\boldsymbol {e} _ {i j} \right\}\right) (5) \\ = \sigma \left(\operatorname {L i n e a r} \left(\boldsymbol {a} _ {i} ^ {\text {r e c}}\right) + \operatorname {L i n e a r} \left(\sum_ {j} \operatorname {L i n e a r} \left(\boldsymbol {a} _ {j} ^ {\text {s e n d}}\right) \odot \operatorname {L i n e a r} \left(\boldsymbol {e} _ {i j}\right)\right)\right) (6) \\ \end{array}
$$
for a receiving particle  $\tilde{\pmb{a}}_i^{\mathrm{rec}}$  and the set of sending particles  $\{\pmb{a}_j^{\mathrm{send}}\}$ , connected via their edges  $\{e_{ij}\}$ . By  $\sigma$  we denote the non-linear activation and with  $\odot$  the element-wise multiplication along the feature dimension. An MPNN is obtained by stacking MessagePassing layers
$$
\operatorname {M P N N} \left(\boldsymbol {a} _ {i}, \left\{\boldsymbol {a} _ {j} \right\}, \left\{\boldsymbol {e} _ {i j} \right\}\right) = \text {M e s s a g e P a s s i n g} (\dots \text {M e s s a g e P a s s i n g} \left(\boldsymbol {a} _ {i}, \left\{\boldsymbol {a} _ {j} \right\}, \left\{\boldsymbol {e} _ {i j} \right\}\right)) \tag {7}
$$
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^{\mathrm{nuc}}$ :
$$
\boldsymbol {x} ^ {\mathrm {n u c}} = \operatorname {p h i s n e t} _ {\theta} (\boldsymbol {R}, \boldsymbol {Z}) \quad \boldsymbol {S} _ {I J} = s _ {\theta} \left(\boldsymbol {x} _ {I} ^ {\mathrm {n u c}}, \boldsymbol {x} _ {J} ^ {\mathrm {n u c}}\right) \quad \boldsymbol {F} _ {I J} = f _ {\theta} \left(\boldsymbol {x} _ {I} ^ {\mathrm {n u c}}, \boldsymbol {x} _ {J} ^ {\mathrm {n u c}}\right). \tag {8}
$$
Here  $\pmb{x}^{\mathrm{nuc}}\in \mathbb{R}^{N_{\mathrm{nuc}}\times (L + 1)^2\times N_{\mathrm{channels}}}$ , and  $S_{IJ}$  and  $F_{IJ}$  are each in  $\mathbb{R}^{N_{\mathrm{basis}}\times N_{\mathrm{basis}}}$ . The basis-set size of the predicted orbitals is denoted by  $N_{\mathrm{basis}}$  and the feature dimension of the nuclear embeddings by  $N_{\mathrm{channels}}$ . The full overlap- and Fock-matrices are assembled from the corresponding blocks  $S_{IJ}$  and  $F_{IJ}$ , leading to matrices of shape  $[N_{\mathrm{nuc}}N_{\mathrm{basis}}\times N_{\mathrm{nuc}}N_{\mathrm{basis}}]$ . Each layer of PhisNet is SE(3)-equivariant, ensuring that any 3D-rotation or inversion of the input coordinates  $\pmb{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,\dots ,L$ , each with components  $m = -l,\ldots ,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:
$$
\boldsymbol {F} \boldsymbol {C} _ {k} = \boldsymbol {S} \boldsymbol {C} _ {k} \epsilon_ {k} \quad \hat {\boldsymbol {x}} _ {I k} ^ {\text {o r b}} = \operatorname {r e s h a p e} \left(\boldsymbol {C} _ {k}, \left[ N _ {\mathrm {n u c}}, N _ {\mathrm {b a s i s}} \right]\right) _ {I} \tag {9}
$$
It has been shown empirically that it is beneficial to obtain the orbital coefficients  $\hat{x}_{Ik}^{\mathrm{orb}}$  as solutions to this generalized eigenvalue problem, rather than predicting them directly [27] as functions of  $\pmb{R}$  and  $\pmb{Z}$ . This is because the orbital coefficients are neither unique, nor do they share the molecule's symmetry. The matrices  $\pmb{F}$  and  $\pmb{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  $\epsilon_{k}$  and obtain the backflow factors  $\pmb{f}$  and exponents  $\pmb{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.
$$
\widetilde {\boldsymbol {x}} _ {I k} ^ {\mathrm {o r b}} = \sum_ {n = 1} ^ {N _ {\mathrm {o r b}}} U _ {k n} ^ {\mathrm {l o c}} \cdot \hat {\boldsymbol {x}} _ {I n} ^ {\mathrm {o r b}}, \quad \boldsymbol {x} _ {I k} ^ {\mathrm {o r b}} = \operatorname {M P N N} \left(\widetilde {\boldsymbol {x}} _ {I k} ^ {\mathrm {o r b}}, \left\{\widetilde {\boldsymbol {x}} _ {J k} ^ {\mathrm {o r b}} \right\}, \left\{\boldsymbol {e} _ {I J} ^ {\mathrm {n u c - n u c}} \right\}\right) \tag {10}
$$
$$
\boldsymbol {f} _ {I k} ^ {\text {o r b}} = \operatorname {M L P} \left(\boldsymbol {x} _ {I k} ^ {\text {o r b}}\right), \quad \boldsymbol {f} _ {I k} ^ {\text {o r b}} \in \mathbb {R} ^ {N _ {\mathrm {n u c}} \times N _ {\mathrm {o r b}} \times N _ {\mathrm {d e t}} \times N _ {\mathrm {e m b}}} \tag {11}
$$
$$
g _ {I k} ^ {\mathrm {o r b}} = \operatorname {M L P} \left(\boldsymbol {x} _ {I k} ^ {\mathrm {o r b}}\right), \quad g ^ {\mathrm {o r b}} \in \mathbb {R} ^ {N _ {\mathrm {n u c}} \times N _ {\mathrm {o r b}} \times N _ {\mathrm {d e t}}} \tag {12}
$$
**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  $\pmb{x}^{\mathrm{nuc}}$ , by first performing a message passing step between the nuclear embeddings
$$
\hat {\boldsymbol {x}} _ {I} ^ {\mathrm {n u c}} = \operatorname {M L P} \left(\boldsymbol {x} _ {I} ^ {\mathrm {n u c}}\right) \quad \tilde {\boldsymbol {x}} _ {I} ^ {\mathrm {n u c}} = \text {M e s s a g e P a s s i n g} \left(\hat {\boldsymbol {x}} _ {I} ^ {\mathrm {n u c}}, \left\{\hat {\boldsymbol {x}} _ {J} ^ {\mathrm {n u c}} \right\}, \left\{\boldsymbol {e} _ {I J} ^ {\mathrm {n u c - n u c}} \right\}\right)
$$
and then using these features to initialize the electron embeddings
$$
\boldsymbol {x} _ {i} ^ {\text {e l}, 0} = \operatorname {M e s s a g e P a s s i n g} (\mathbf {0}, \{\tilde {\boldsymbol {x}} _ {J} ^ {\text {n u c}} \}, \{e _ {i J} ^ {\text {e l - n u c}} \}),
$$
by using a zero vector  $\mathbf{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
$$
\boldsymbol {x} _ {i} ^ {\mathrm {e l}} = \operatorname {M P N N} \left(\boldsymbol {x} _ {i} ^ {\mathrm {e l}, 0}, \left\{\boldsymbol {x} _ {j} ^ {\mathrm {e l}, 0} \right\}, \left\{\boldsymbol {e} _ {i j} ^ {\mathrm {e l} - \mathrm {e l}} \right\}\right),
$$
resulting in a  $N_{\mathrm{emb}}$  -dimensional representation for each electron  $\pmb {x}_i^{\mathrm{el}}\in \mathbb{R}^{N_{\mathrm{emb}}}$ $i = 1,\dots ,n_{\mathrm{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 therefore 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  $\psi(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  $\psi^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  $\psi^2$  and thus requires  $\approx 10^5$  MCMC steps 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  $\pmb{R}$  is required. Starting with an initial set of geometries  $\pmb{R}^0$ , we update  $\pmb{R}$  on the fly, by perturbing each geometry every 20 optimization steps by adding random noise  $\Delta R$  to the nuclear coordinates. Using uncorrelated, isotropic random noise for  $\Delta R$  would yield many non-physical geometries  $\pmb{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  $\Delta R$  from a correlated normal distribution
$$
\Delta \boldsymbol {R} \sim \mathcal {N} \left(\alpha \left(\boldsymbol {R} ^ {0} - \boldsymbol {R}\right), \beta \boldsymbol {H} _ {\text {p h i s}} ^ {- 1}\right) \quad \boldsymbol {R} ^ {\prime} = \boldsymbol {R} + \Delta \boldsymbol {R}. \tag {13}
$$
The bias term  $\alpha (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^{\mathrm{phis}}$  which is predicted from the scalar component of the nuclear embeddings using a pre-trained MLP. By using  $H_{\mathrm{phis}}^{-1}$  as covariance matrix, we take large steps along soft directions and small steps along stiff directions, thus avoiding unphysical geometries with very high energies.
$$
E ^ {\text {p h i s}} = \sum_ {I = 1} ^ {N _ {\mathrm {n u c}}} \operatorname {M L P} \left(\boldsymbol {x} _ {I} ^ {\mathrm {n u c}}\right), \quad H _ {I \zeta , J \xi} ^ {\text {p h i s}} = \frac {\partial^ {2} E ^ {\text {p h i s}}}{\partial \boldsymbol {R} _ {I \zeta} \partial \boldsymbol {R} _ {J \xi}}, \quad I, J = 1 \dots N _ {\mathrm {n u c}} \quad \zeta , \xi = 1 \dots 3 \tag {14}
$$
After distorting the nuclear coordinates  $\mathbf{R}$ , we also adjust the electron positions  $\mathbf{r}$ , using the space-warp 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.