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.