Enabling ab initio geometry optimization of strongly correlated systems with transferable deep quantum Monte Carlo
Authors: P. Bernát Szabó, Zeno Schätzle, Frank Noé
Organizations: FU Berlin, Department of Mathematics and Computer Science, Arnimallee 6, 14195 Berlin, Germany · Microsoft Research AI4Science, Karl-Liebknecht Str. 32, 10178 Berlin, Germany · FU Berlin, Department of Physics, Arnimallee 14, 14195 Berlin, Germany · Rice University, Department of Chemistry, Houston, TX 77005, USA
A faithful description of chemical processes requires exploring extended regions of the molecular potential energy surface (PES), which remains challenging for strongly correlated systems. Transferable deep-learning variational Monte Carlo (VMC) offers a promising route by efficiently solving the electronic Schrödinger equation jointly across molecular geometries at consistently high accuracy, yet its stochastic nature renders direct exploration of molecular configuration space nontrivial. Here, we present a framework for highly accurate ab initio exploration of PESs that combines transferable deep-learning VMC with a cost-effective estimation of energies, forces, and Hessians. By continuously sampling nuclear configurations during VMC optimization of electronic wave functions, we obtain transferable descriptions that achieve zero-shot chemical accuracy within chemically relevant distributions of molecular geometries. Throughout the subsequent characterization of molecular configuration space, the PES is evaluated only sparsely, with local approximations constructed by estimating VMC energies and forces at sampled geometries and aggregating the resulting noisy data using Gaussian process regression. Our method enables accurate and efficient exploration of complex PES landscapes, including structure relaxation, transition-state searches, and minimum-energy pathways, for both ground and excited states. This opens the door to studying bond breaking, formation, and large structural rearrangements in systems with pronounced multi-reference character.
Figures & tables
Figure 1: Sketch of the DeepQMC/GPR ab initio geometry optimization method. The flowchart depicts the process of optimizing the minimum energy pathway (MEP) of the formaldehyde isomerization. a) A transferable wave function is optimized on continuously sampled nuclear configurations of a two dimensional cut of the formaldehyde configuration space. Since the training loss consists of expectation values over small electron batches at random nuclear configurations, during transferable training the full PES is never materialized beyond noisy estimates. b) To extract energy and gradient information from the optimized transferable wave function, Monte Carlo estimators of the molecular energy and Hellmann-Feynman force are evaluated. The predictions are further improved by evaluating these expectation values at a range of molecular configurations in the vicinity of the geometry under investigation and fitting a local approximation of the PES with Gaussian process regression (GPR). From the Gaussian process energy, force and Hessian are computed analytically, which suffice to compute second order geometry updates. c) By a consecutive application of optimization steps, molecular geometries can be relaxed, transition states can be found and intermediates can be obtained when employing additional constrains. d) Combining multiple geometry optimizations, the full MEP of a chemical process can be characterized. Notably, for the MEP optimization the PES was only approximated sparsely via the samples depicted in gray.
Figure 2: Optimizing equilibrium bond lengths for diatomic molecules. DeepQMC/GPR results are compared to experimental data [ 34 ] and to MP2@CBS and CCSD(T)@cc-pV6Z calculations [ 35 ] .
Figure 3: Minimum energy path search on two dimensional cuts of the ammonia and formaldehyde potential energy surfaces. On panel a) the minimum energy pathway for the inversion of ammonia is shown, while panel b) depicts the isomerization of formaldehyde. Reference and baseline results for the ammonia inversion are taken from Ref. [ 23 ] , CCSD results for formaldehyde are computed in house. The heatmaps visualize PESs extracted from the transferable deep-learning VMC solution and serve as a guide to the eye.
Figure 4: Adiabatic excitation energies for the ethylene molecule. On panels a)-d) the internal coordinates of the relaxation trajectory of the triplet state after a vertical excitation from the singlet ground state equilibrium geometry are depicted. The VMC [ 22 ] reference equilibrium geometries of the singlet and triplet states are well reproduced. Corresponding estimates of the relative energy with respect to the triplet equilibrium geometry are displayed on panel e). The inserts show the optimized equilibrium geometry of the singlet and triplet state. Panel f) and panel g) compare the adiabatic and the vertical excitation energies to reference calculations [ 22 ] , where the shaded area indicates chemical accuracy.
Figure 5: Potential energy surfaces of the HX2 + NH \makebox[0.0pt]→H⋅ + ⋅NHX2 reaction. The reaction might proceed either on the a1A′′ or the b1A′ PES. Inset plots show the relevant reactant, transition state, and product structures highlighting the differences between the two electronic states. The FCI results taken from Wu et al. [ 38 ] were computed on geometries optimized at the CASPT2 level of theory, with the relatively small basis set of 6-311G*. The energies of these geometries were also evaluated with DeepQMC, shown with solid lines.
Figure 6: MEP for the HOX2⋅ + OH⋅\makebox[0.0pt]→HX2O + OX2 reaction. The results of the deep-learning VMC MEP optimization are compared with reference values obtained with a permutation invariant polynomial (PIP) neural network potential trained on 108k CCSD(T)-F12 single-point energies [ 40 ] .
reactant
CP1
TS
product
relative energy (kcal/mol)
CCSD(T)/PIP
0.0
−3.7
−1.4
−70.6
DeepQMC/GPR
0.0
−4.9±0.5
0.4±0.5
−69.8±0.4
DeepQMC on PIP geom
0.0
−4.3±0.6
−0.5±0.6
−69.9±1.0
DeepQMC v. PIP geometry MAE
bonds (a 0 )
0.02
0.03
0.01
<0.01
angles (deg)
0.7
7.7
4.3
<0.1
torsions (deg)
—
12
3
—
Table 1: Comparison of the deep-learning VMC and CCSD(T) characterizations of the HOX2⋅ + OH⋅\makebox[0.0pt]→HX2O + OX2 reaction. TS denotes the transition state structure, while CP1 corresponds to the local minimum found near the entrance of the reactive channel, first described in Liu et al. [ 40 ] . Two sets of geometries are considered: those obtained from the PIP neural-network potential fitted to CCSD(T)-F12a energies, and those optimized via the presented DeepQMC/GPR method.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Continous transferable training converges to single point results. a) Relative energy error from converged single point optimizations for 16 geometries of the N2 molecule. Continuous sampling and training on the test set converge to chemical accuracy wrt. to the reference on all 16 geometries. Training on a sparse subset of four geometries (1.5 bohr, 2.93 bohr, 4.36 bohr, 6.5 bohr) results in large interpolation errors. b) Mean absolute error in relative energies wrt. converged single point simulations along the entire dissociation curve throughout the optimization. Each data point is an evaluation of the wave function(s) on all geometries at the respective training iteration.
Figure 8: Benchmarking the cost-accuracy ratio of various VMC nuclear force estimators. The comparison is carried out on ammonia, formaldehyde, and O 2 + H 2 O molecular geometries. Finite difference forces obtained with the CCSD(T) based W1 protocol are used as reference. The mean absolute deviations are computed across all nuclei and all cartesian directions of each geometry. The wall clock runtimes were measured using four Nvidia A100 GPUs in parallel.
Description
Value
Unit
C–C distance
-0.5 – 1.0
bohr
C–H distance (4x)
0.5
bohr
H–C–H angle (2x)
0.5
rad
H–C–C–H dihedral
π
rad
Appendix
Table 2: The widths (or min – max if unsymmetrical) of the uniform distributions used to define the training geometry distribution for C2H4. The atoms are not numbered, since all occurrences of the internal coordinates that are considered as bonds in the equilibrium geometry are equal.
Description
Value
Unit
H1 – H2 distance
0.755
bohr
H2 – N distance
0.516
bohr
N – H3 distance
0.2
bohr
H2 – N – H3 angle
1.1
rad
Appendix
Table 3: The widths of the 25 uniform distributions used to define the training geometry distribution for the H 2 + NH reaction. The hydrogen atoms of the attacking hydrogen molecule are designed H1 (farther from NH) and H2 (closer to NH), while the hydrogen of NH is denoted H3.
Description
Value
Unit
O1 – O2 distance
1.0
bohr
O2 – H1 distance
0.12
bohr
O1 – O2 – H1 angle
1.4
rad
H1 – O3 distance
0.22
bohr
O2 – H1 – O3 angle
0.7
rad
O1 – O2 – H1 – O3 dihedral
1.516
rad
Appendix
Table 4: The widths of the 32 uniform distributions used to define the training geometry distribution for the O 2 H + OH reaction. The atoms of the O–O–H molecule are designated O1, O2, and H1, respectively, while the atoms of the OH molecule are denoted O3 and H2.
Hyperparameter
formaldehyde
ammonia
ethylene
H 2 + NH
O 2 H + OH
initial trust radius (bohr)
0.2
0.2
0.05
0.2
0.25
Newton-step regularizer ( λ )
10−2
0
0
0
10−2
data radius (bohr)
0.2
∞
∞
0.2
0.2
energy convergence criterion (Hartree)
10−3
5×10−4
5×10−4
10−3
10−3
Appendix
Table 5: Collection of hyperparameters employed during the geometry optimizations and minimum energy pathway searches. ‘data radius’ refers to the maximum Euclidian distance in between the target geometry and the sampled data points considered for fitting the local Gaussian approximate potential.
Variational Monte Carlo (VMC) is a central algorithm in electronic structure theory and has gained renewed importance through modern neural-network ansätze such as FermiNet. At its core, VMC seeks ground states by minimizing the Rayleigh quotient by stochastic optimization. In this work, we show that the resulting stochastic optimization problem is intrinsically governed by the nodal geometry of the underlying wave function. More precisely, we establish that properties of the nodal set determine the integrability of the local energy and gradient estimators that drive VMC. For broad and practically relevant ansatz classes, including Slater-Jastrow wave functions with variable-exponent Slater-type orbitals, we prove that these estimators are generically heavy-tailed and fail to admit higher moments. At the same time, for general analytic ansätze, we prove weak moment bounds for the relevant estimators and identify precise low-moment regimes, showing how generic and degenerate nodal structures lead to different integrability thresholds. Building on this analysis, we introduce a new robust variant of VMC \unicodex2013 coined PS-Clip-VMC \unicodex2013 which is based on clipping both the local energy and the gradient random variable. We prove that PS-Clip-VMC converges both in expectation and with high probability in the weak moment regime of VMC. Preliminary experiments for training FermiNet on Atoms with up to 18 electrons suggest that PS-Clip-VMC is significantly more robust than standard methods.
Philipp Grohs, Davide Nobile
Faculty of Mathematics, University of Vienna · RICAM, Austrian Academy of Sciences
Fundamental tasks in computational chemistry, from transition state search to vibrational analysis, rely on molecular Hessians, which are the second derivatives of the potential energy. Yet, Hessians are computationally expensive to calculate and scale poorly with system size, with both quantum mechanical methods and neural networks. In this work, we demonstrate that Hessians can be predicted directly from a deep learning model, without relying on automatic differentiation or finite differences. We observe that one can construct SE(3)-equivariant, symmetric Hessians from irreducible representations (irrep) features up to degree l=2 computed during message passing in graph neural networks. This makes HIP Hessians one to two orders of magnitude faster, more accurate, more memory efficient, easier to train, and enables more favorable scaling with system size. We validate our predictions across a wide range of downstream tasks, demonstrating consistently superior performance for transition state search, accelerated geometry optimization, zero-point energy corrections, and vibrational analysis benchmarks. We open-source the HIP codebase and model weights to enable further development of the direct prediction of Hessians at https://github.com/BurgerAndreas/hip
Andreas Burger, Luca Thiede, Nikolaj Rønne +5
University of Toronto · Vector Institute for Artificial Intelligence · Technical University of Denmark +4
Reliable description of bond breaking remains a major challenge for quantum chemistry due to the multireference character of the electronic structure in dissociating species. Multireference methods in particular suffer from large computational cost, which under the normal paradigm has to be paid anew for each system at a full price, ignoring commonalities in electronic structure across molecules. Quantum Monte Carlo with deep neural networks uniquely offers to exploit such commonalities by pretraining transferable wavefunction models, but all such attempts were so far limited in scope. Here, we bring this paradigm to fruition with Orbformer, a transferable wavefunction model pretrained on 22,000 equilibrium and dissociating structures that can be fine-tuned on unseen molecules reaching an accuracy-cost ratio rivalling classical multireference methods. On established benchmarks as well as more challenging bond dissociations and Diels-Alder reactions, Orbformer is the only method that consistently converges to chemical accuracy (1 kcal/mol). This work turns the idea of amortizing the cost of solving the Schr"odinger equation over many molecules into a practical approach in quantum chemistry.
Adam Foster, Zeno Schätzle, P. Bernát Szabó +7
Microsoft Research AI for Science · Freie Universität Berlin