Implementing advanced trial wave functions in fermion quantum Monte Carlo via stochastic sampling

TL;DR

The paper presents a scalable framework to incorporate advanced correlated trial wave functions into auxiliary-field quantum Monte Carlo by stochastically sampling the trial state. It couples branching random walks with Metropolis sampling of multi-dimensional auxiliary variables, enabling overlaps and constraints to be evaluated as stable ratios while maintaining polynomial scaling. Demonstrations on N$_2$ bond dissociation and transition-metal-oxide diatomics show energies within chemical accuracy and improved efficiency compared with conventional multi-determinant trials. The approach generalizes to other trial forms, including neural-network states, and opens avenues for broader quantum-classical hybrid QMC applications and finite-temperature extensions.

Abstract

We introduce an efficient approach to implement correlated many-body trial wave functions in auxiliary-field quantum Monte Carlo (AFQMC). To control the sign/phase problem in AFQMC, a constraint is derived from an exact gauge condition but is typically imposed approximately through a trial wave function or trial density matrix, whose quality can affect the accuracy of the method. Furthermore, the trial wave function can also affect the efficiency through importance sampling. The most natural form of the trial wave function has been single Slater determinants or their linear combinations. More sophisticated forms, for example, with the inclusion of a Jastrow factor or other explicit correlations, have been challenging to use and their implementation is often assumed to require a quantum computer. In this work, we demonstrate that a large class of correlated wave functions, written in the general form of multi-dimensional integrals over hidden or auxiliary variables times Slater determinants, can be implemented as trial wave function by coupling the random walkers to a generalized Metropolis sampling. We discuss the fidelity of AFQMC with stochastically sampled trial wave functions, which are relevant to both quantum and classical algorithms. We illustrate the method and show that an efficient implementation can be achieved which preserves the low-polynomial computational scaling of AFQMC. We test our method in molecules under bond stretching and in transition metal diatomics. Significant improvements are seen in both accuracy and efficiency over typical trial wave functions, and the method yields total ground-state energies systematically within chemical accuracy. The method can be useful for incorporating other advanced wave functions, for example, neural quantum state wave functions optimized from machine learning techniques, or for other forms of fermion quantum Monte Carlo.

Figures (8)

• Figure 1: Illustration of (a) the MCMC and (b) the BRW algorithms to evaluate Eq. \ref{['eq:overview2']}. The figures depict $m=2$ and $n=2$. Circle and rectangular denote one-body operators $\hat{B}(\textbf{x})$ and $\hat{\mathcal{B}}_T^j(\textbf{y})$, respectively. Blue/red indicates current/updated state in the MC. In (a), paths of auxiliary-fields, $\{\textbf{y}^j,\textbf{x}^i\}$, and sampled by the Metropolis algorithm. In the figure, a schematic is shown for one path. In (b), a population of walkers $\{|\phi^i_k\rangle\}$ (with weight) are propagated from right to left. In the figure, the trajectory of one walker (labeled by $k$) is shown. The rectangular box shaded orange grows from right to left, indicating that often in ground-state calculations only the current walker state is kept and the AF path history can be discarded.
• Figure 2: Illustration of the structure of our Monte Carlo sampling algorithm. One random walker and its tethered MCMC paths are shown, during one leapfrog update step. The colors blue/red/orange have the same meaning as in Fig. \ref{['Fig.Metro_BRW']}, and circle and rectangular again denote $\hat{B}$ and $\hat{\mathcal{B}}_T$, respectively. (0) current state of the walker $|\phi^{i}_k \rangle$ and the $|\Psi_T \rangle$ paths attached to it. (Taken literally, this picture depicts $i=2$, with $m=3$ and $P$ given by the number of rows on the left.) (1). the walker is propagated one step, $|\phi^i_k \rangle \rightarrow |\phi^{i+1}_k \rangle$ and its weight updated $W^i_k \rightarrow W^{i+1}_k$, using the current samples of $\langle \Psi_T|$ as importance function and constraint. (2). MCMC sweeps are performed to update each path (each $p$, represented by one row of squares), $Y_k^i[p] \rightarrow Y_k^{i+1}[p]$, according to the walker's new position, $|\phi^{i+1}_k \rangle$.
• Figure 3: Ground state energy systematic errors for AFQMC with different trial wave functions on the F atom. The trial wave function is $|\Psi_T\rangle=e^{-m\,\tau\,\hat{H}}|\phi_{\rm RHF}\rangle$ with $\tau = 0.01$. The cc-pVDZ basis is used. The error is given by $E_{\rm AFQMC}-E_{\rm FCI}$ and $\beta_T = m\tau$ is also given in atomic units. The AFQMC systematic error with the RHF itself as trial wave function is presented as the "shaded blue background" whose width indicates the statistical errorbar. "Orange star" denotes AFQMC results, and related trial energy is specified with "black square".
• Figure 4: The energy trajectory of the AFQMC propagation with the VAFQMC trial wave function of Eq. (\ref{['eq:VAFQMC']}), labeled AFQMC/VAFQMC, in the N$_2$ molecule at a stretched bondlength of $R=3.6\,$Bohr. The ground-state energy $E(\beta)$ is evaluated along the AFQMC imaginary-time propagation with increased $\beta$, and plotted as $E(\beta)-E_{\rm DMRG}$, relative the the near-exact value from DMRG N2_bondbreaking_DMRG. The converged and final AFQMC ground state energy is shown as the shaded orange line whose width indicates the estimated statistical uncertainty. The variational energy of the VAFQMC trial wave function itself is given as the shaded blue line, again with the width indicating statistical error.
• Figure 5: Fidelity and convergence of the stochastic sampling of $|\Psi_T\rangle$, and computational efficiency. (a) Energy computed from AFQMC (relative to the near-exact DMRG results N2_bondbreaking_DMRGAFQMC_bondBreaking) as a function of $P$, the number of paths tethered to each walker in the Metropolis sampling of the VAFQMC $|\Psi_T\rangle$. The thin dashed black line indicates chemical accuracy. The inset shows, for comparison, results from AFQMC using a multi-determinant trial wave function taken from CASSCF, truncated to $P$ determinants according to the absolute value of the CI coefficient. The converged results of AFQMC/MD and AFQMC/VAFQMC (in their respective limit of large $P$) are shown by a thick line shaded in the corresponding color, with the line width indicating statistical errors (b) The computational accuracy and efficiency of AFQMC/VAFQMC as a function of $P$. The statistical accuracy is shown as stars (left scale), with the dashed orange curve to guide the eye. The computational cost (normalized to that of $P=5$) is shown as squares (right scale).