NNQS-AFQMC: Neural network quantum states enhanced fermionic quantum Monte Carlo

TL;DR

This work addresses the challenge of accurately simulating strongly correlated electrons by combining the expressive power of neural network quantum states (NNQS) with the robustness of AFQMC through a stochastic trial-wavefunction framework. The authors implement NNQS-AFQMC by representing the NNQS trial as a stochastic sum over configurations and integrating it into AFQMC via MCMC sampling of a precomputed dataset, achieving near-exact energies for $N_2$ across bond lengths and basis sets. The method shows robustness to the quality of the underlying NNQS and extends naturally to CI-based trial states, offering a scalable route to high-accuracy electronic structure calculations while reducing NNQS training overhead. Overall, NNQS-AFQMC provides a versatile, efficient approach that could broaden the applicability of NNQS in strongly correlated quantum chemistry and materials problems.

Abstract

We introduce an efficient approach to implement neural network quantum states (NNQS) as trial wavefunctions in auxiliary-field quantum Monte Carlo (AFQMC). NNQS are a recently developed class of variational ansätze capable of flexibly representing many-body wavefunctions, though they often incur a high computational cost during optimization. AFQMC, on the other hand, is a powerful stochastic projector approach for ground-state calculations, but it normally requires an approximate constraint via a trial wavefunction or trial density matrix, whose quality affects the accuracy. Recently it has been shown (Xiao et al, arXiv2505.18519) that a broad class of highly correlated wave-functions can be integrated into AFQMC through stochastic sampling techniques. In this work, we apply this approach and present a direct integration of NNQS with AFQMC, allowing NNQS to serve as high-quality trial wavefunctions for AFQMC with manageable computational cost. We test the NNQS-AFQMC method on the challenging nitrogen molecule (N$_2$) at stretched geometries. Our results demonstrate that AFQMC with an NNQS trial wavefunction can attain near-exact total energies, highlighting the potential of AFQMC with NNQS to overcome longstanding challenges in strongly correlated electronic structure calculations. We also outline future research directions for improving this promising methodology.

Figures (6)

• Figure 1: Illustration of the integration of NNQS with AFQMC. One random AFQMC walker and its tethered MCMC sampled configurations are shown, during one leapfrog update step. Circle and rectangular denote one-body operators $\hat{B}(\textbf{y})$ and configurations $\langle \mathbf{x}_p|$, respectively. Blue/red indicates current/updated state in the MC. 0). current state of the walker $|\phi^{i}_k \rangle$ and MCMC sampled configurations $\langle \bar{\Psi}_T|$ attached to it. 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 \bar{\Psi}_T|$ as importance function and constraint. 2). MCMC sweeps are performed to update each configurations (each $p$, represented by one row of rectangular), $\mathbf{x}_{k,p}^i \rightarrow \mathbf{x}_{k,p}^{i+1}$, according to the walker's new position, $|\phi^{i+1}_k \rangle$.
• Figure 2: Deviations in the computed ground-state energy from NNQS and NNQS-AFQMC, for the N$_2$ molecule at different bond lengths, in the cc-pVDZ basis. Energies are compared with near-exact DMRG results N2_bondbreaking_DMRGAFQMC_bondBreaking. Here, NNQS results do not stand for the optimized QiankunNet results, but the trial energy of the NNQS trial wave function for the related NNQS-AFQMC results. MD-AFQMC results, with trial wave functions from truncated MCSCF, was taken from Ref. AFQMC_bondBreaking. The dashed green line represents the chemical accuracy (1 kcal/mol).
• Figure 3: The convergence of NNQS-AFQMC along the training process of NNQS for $\mathrm{N}_2$ at 4.2 Bohr (cc-pVDZ basis). Blue curves indicate the NNQS energy during the training process. AFQMC is carried out with NNQS trial wave functions optimized at $4, 6, 8, 10, 12\times10^4$ steps correspondingly, denoted in the figure with blue dots. The green shaded region indicates the chemical accuracy (1 kcal/mol), and the red shaded region specifies the convergence of NNQS-AFQMC.
• Figure 4: Energy deviation of NNQS-AFQMC vs. the number of Metropolis samples (P) for $\mathrm{N}_2$ at 4.2 Bohr (cc-pVDZ basis). Calculations are performed with the same NNQS trial wave function and parameters, except for the number of Metropolis samples $P$. NNQS denotes the trial energy of the NNQS trial wave function. The green shaded region indicates the chemical accuracy (1 kcal/mol), and the red shaded region specifies the convergence of NNQS-AFQMC.
• Figure 5: Comparison of NNQS-AFQMC with MD-AFQMC for $\mathrm{N}_2$ at 4.2 Bohr (cc-pVDZ basis). Both methods aim to implement the same NNQS trial wave function. NNQS-AFQMC is implemented through stochastic sampling with $P$ sampled configurations. MD-AFQMC employing multi-determinant trial wavefunctions that consist of the $P$ configurations with the largest coefficients in NNQS. The green shaded region indicates the chemical accuracy (1 kcal/mol), and the red shaded region specifies the convergence of NNQS-AFQMC.