Variance extrapolation method for neural-network variational Monte Carlo

VMC is a stochastic method for solving quantum many-body problems. For a complex quantum system under Born-Oppenheimer approximation [34], the nonrelativistic Hamiltonian operator can be expressed as

$$\begin{align} \begin{split} \hat{H} = - \frac{1}{2}\sum_{i}\nabla^{2}_{i} - \sum_{I, i}\frac{Z_{I}}{\left|\mathrm{\mathbf{r}}_{i}-\mathrm{\mathbf{R}}_{I}\right|}+ \sum_{i\lt j}\frac{1}{\left|\mathrm{\mathbf{r}}_{i}-\mathrm{\mathbf{r}}_{j}\right|} + \sum_{I\lt J}\frac{Z_I Z_J}{\left|\mathrm{\mathbf{R}}_{I}-\mathrm{\mathbf{R}}_{J}\right|}, \end{split} \end{align} \tag{ 1 }$$

where $i, j$ are subscripts for electrons, and $I, J$ for nuclei. Z_I denotes the charges, and $r_i, R_I$ denote the positions. Regardless of time evolution, one of the main goals is to obtain the ground state wavefunction $| \psi_0 \rangle$ and corresponding energy E₀, which are determined by the stationary Schrödinger equation,

$$\begin{align} \hat{H}| \psi_i \rangle = E_i| \psi_i \rangle, i = 0,1,2\ldots \end{align} \tag{ 2 }$$

To accomplish this, VMC employs Monte Carlo methods to handle the high dimensional energy integral and variational principles to optimize the wavefunction towards the ground state. Given a wavefunction ansatz $\psi_\theta(\mathbf{r})$, its energy expectation is defined as

$$\begin{align} \begin{split} E\left[\psi_\theta\right] & = \frac{\langle{\psi_\theta}|\hat{H}|{\psi_\theta}\rangle}{\langle \psi_\theta | \psi_\theta\rangle} = \frac{\int\psi^*_\theta\left(\mathbf{r}\right)\hat{H}\psi_\theta\left(\mathbf{r}\right)\mathrm{d}\mathbf{r}}{\int\psi^*_\theta\left(\mathbf{r}\right)\psi_\theta\left(\mathbf{r}\right)\mathrm{d}\mathbf{r}} \\ & = \int p\left(\mathbf{r}\right) E_L\left(\mathbf{r}\right) \mathrm{d}\mathbf{r} = \mathbb{E}_{p\left(\mathbf{r}\right)}\left[E_L\right], \end{split} \end{align} \tag{ 3 }$$

where $p(\mathbf{r}) = \frac{\left|\psi_\theta(\mathbf{r})\right|^2}{\int \left|\psi_\theta(\mathbf{r})\right|^2 \mathrm{d}\mathbf{r}}$ is the normalized wavefunction squared, i.e. the probability density function, and $E_L(\mathbf{r}) = \frac{\hat{H}\psi_\theta(\mathbf{r})}{\psi_\theta(\mathbf{r})}$ is the so-called local energy function. The outcome of this energy integral must be greater than or equal to the ground state energy, which is the target of VMC, so we just need to optimize the wavefunction towards lower energy,

$$\begin{align} -\nabla_\theta E\left[\psi_\theta\right]\propto \mathbb{E}_{p\left(\mathbf{r}\right)}\left[\left(\mathbb{E}_{p\left(\mathbf{r}\right)}\left[E_L\right] -E_L\right)\nabla_\theta \log \left|\psi_\theta \right|\right]. \end{align} \tag{ 4 }$$

Keep optimizing until convergence, and we will get the approximation of ground state wavefunction and corresponding energy, whose accuracy limit depends on the expressiveness of the wavefunction ansatz. In addition, the quality of the optimization and the statistical error would also affect the accuracy.

NN-VMC reaps the benefit of a vast number of parameters to obtain high accuracy. However, this comes at the cost of a lengthy optimization process, which unfortunately results in less satisfactory error cancellation capability when calculating the relative energy. Since there is no established criterion for measuring the degree of convergence during the training, the common approach is to use the same optimization iteration in all the systems for calculating relative energies, as shown in figure 1(b) and labeled as 'Training Matched' (TM). Nevertheless, this can lead to an unbalanced treatment if the optimization is inadequate for convergence because of the different speed of convergence in different systems. Even if achieving complete convergence, this simple fairly-treating approach is still unbalanced as the more complicated system requires a larger network to obtain the same accuracy.

Zoom In Zoom Out Reset image size

To achieve balanced treatments, matching a quantity related to the error in a wavefunction is a potential strategy, which has been utilized in various studies to obtain excitation energy [35–39]. The energy variance is selected as the matching quantity and is defined as:

$$\begin{align} \sigma_\theta^2 = \frac{\langle{\psi_\theta }| \left(\hat{H} - E\right)^2 |{\psi_\theta}\rangle}{\langle \psi_\theta| \psi_\theta\rangle}, \end{align} \tag{ 5 }$$

where $E = \langle\hat{H}\rangle_{\psi_\theta}$ denotes the energy expectation. As shown in figure 1(c), the 'Variance Matched' (VM) method is thus comparing the energies of two systems where they have the same variance. To clarify further, we choose the data point with the most training steps from 'system 1' and compare its energy result with 'system 2', whose energy result is determined via linear regression at the corresponding variance.

Another strategy is the zero-variance extrapolation method, where the energy outcomes are extrapolated until the variance reaches zero as the real ground state should possess zero energy variance. In VMC, when the wavefunction is optimized to be near the ground state, the energy and variance exhibit a linear relationship, as demonstrated in figure 1(c). Although the fundamental reason for the linear relationship is not fully clear [40, 41], it leads to empirical extrapolation schemes that have been successfully used in various systems, including Fermi liquids and the Hubbard model [42–45]. The most commonly used schemes for extrapolation in previous studies are: (1) based on the converged results of different ansatz, such as different forms of backflow or different widths of the network [42, 43], and (2) based on the results of different converged degrees with the same ansatz, such as different Lanczos steps [44, 45]. In other words, there are two different routes to obtain a range of wavefunctions with different 'distance' to the ground truth. In this work, we employ extrapolation based on the training data at different converged levels of NN-VMC, as shown in figure 1(c), labeled as 'Variance Extrapolation' (VE).

To obtain effective information from the noisy training data, a non-overlapping rolling window, specifically characterized by intervals such as [1, n], [n + 1, 2n], and so forth, is employed. This method serves to calculate the robust mean of the energy and variance, which means in each window the outlier data are masked. We assume that the wavefunction in a window does not change much, and the energy data follow the normal distribution. For a window containing n data points, the data points whose absolute deviation of the energy from the median is larger than 3σ are considered to be outliers, where σ is the standard deviation. In all the calculations we mentioned in this paper, we set the window size n to be 2000. After masking, the energy and variance outcomes of this window are calculated following:

$$\begin{align} \left\{ \begin{aligned} E_\textrm{out} & = \overline{e} \\ V_\textrm{out} & = \overline{v} + \overline{e^2} - \overline{e}^2 \end{aligned} \right. \end{align} \tag{ 6 }$$

where e and v denote the energy and variance of the remaining data after masking, and the overline denotes taking the average. It is worth pointing out that in the training process of NN-VMC, there may be extreme outliers that can significantly impact the extrapolation outcome, thus making the masking step indispensable.

For more technical details about the selection of linear segments for extrapolation and the appropriateness of using training data rather than inference data, refer to supplementary Note 1.

In the VE method, we assume that there exists linear relationship between the energy and the energy variance when the wavefunction is nearing the ground state. As mentioned earlier, the fundamental reason for the linear relationship remains unclear. However, we can provide a possible explanation by offering a sufficient but unnecessary condition for such a relationship, following the idea of Kashima Imada [46].

In the late training period, firstly, we decompose the neural-network state into a sum of the ground state $| \psi_0 \rangle$ and the remaining orthogonal components $| \psi^{^{\prime}} \rangle$. That is,

$$\begin{align} | \psi_\theta \rangle = \alpha| \psi_0 \rangle+\beta| \psi^{^{\prime}} \rangle. \end{align} \tag{ 7 }$$

Here $|\alpha| \gg |\beta|$, corresponding to the ground state's dominance as the training process approaches convergence. Without loss of generality, here we assume the normalization condition $\alpha^2+\beta^2 = 1$. The energy expectation result is

$$\begin{align} E = \langle\hat{H}\rangle_{\psi_\theta} = E_0+\beta^2\left(E^{^{\prime}}-E_0\right), \end{align} \tag{ 8 }$$

where E₀ is the ground energy, and $E^{^{\prime}} = \langle\hat{H}\rangle_{\psi^{^{\prime}}}$. The energy variance result is

$$\begin{align} V = \langle\left(\hat{H}-E\right)^2\rangle_{\psi_\theta} = \beta^2 V^{^{\prime}} + \alpha^2\beta^2\left(E^{^{\prime}}-E_0\right)^2, \end{align} \tag{ 9 }$$

where $V^{^{\prime}} = \langle(\hat{H}-E^{^{\prime}})^2\rangle_{\psi^{^{\prime}}}$. Combining these two equations, we can obtain

$$\begin{align} E = \frac{1}{\alpha^2 \left(E^{^{\prime}}-E_0\right)}\left(V-\beta^2V^{^{\prime}}\right)+E_0. \end{align} \tag{ 10 }$$

Then, with two assumptions:

• (a)  
  $V \gg \beta^2 V^{^{\prime}}$, i.e. $(E^{^{\prime}}-E_0)^2 \gg V^{^{\prime}}$. This assumption requires $| \psi^{^{\prime}} \rangle$ to be far away from the ground state and its components to be concentrated, approaching an eigenstate;
• (b)  
  Eʹ remains constant,

we can derive the linear relationship

$$\begin{align} E = kV+E_0, \end{align} \tag{ 11 }$$

where the slope factor $k\approx\frac{1}{E^{^{\prime}}-E_0}$.

It is worth mentioning that all the calculations presented in this work are performed with KFAC optimizer [49, 52]. Given that the variance extrapolation relies on wavefunctions obtained at various stages of training, it raises the question of whether the choice of optimizer may impact the performance of the variance extrapolation approach.

In order to address this concern, we also performed calculations on N₂ and Li₂ molecules using FermiNet with the ADAM optimizer, whose results are detailed in supplementary note 3. The linear relationship between energy and variance persists, and the variance extrapolation results remain highly satisfactory. As a result, we can conclude that this linear relationship appears to be largely insensitive to the choice of optimizer.

We now discuss the tests of the VM and VE methods on N₂ molecule in more detail. N₂ molecule is strongly correlated at the breaking-bond region due to the complexity of triple-bond breaking, providing a challenge case for accurate wavefunction calculation. Energies at three bond lengths are calculated using FermiNet as the wavefunction ansatz, including $r_0 = 2.0743$ Bohr at the equilibrium geometry, $r_1 = 4.0$ Bohr in the strongly correlated region and $r_2 = 10.0$ Bohr in the completely dissociated region. As shown in figure 2(a), the energy and variance show linear relationship at all the three bond lengths, which makes it possible for VM and VE methods to be applied.

Zoom In Zoom Out Reset image size

The energy differences relative to the equilibrium geometry $\Delta E_1 = E(r_1) - E(r_0)$ and $\Delta E_2 = E(r_2) - E(r_0)$ are investigated respectively using TM, VM and VE methods, as illustrated in figure 2(b). Because of the high accuracy of FermiNet on the N₂ molecule at bond lengths r₀ and r₂, the TM method can already provide a good estimate of $\Delta E_2$, leaving only a little room for improvements. Still and all, it is obvious that the VE method performs better than the VM method. Regarding $\Delta E_1$, when the accuracy of FermiNet at r₁ decreases due to the strong correlation, the TM method becomes inadequate in providing a reliable result. In this case, the VM method can lead to only a minor improvement, whereas the VE method significantly reduces the error by 67.8%, from 2.84 mHa to 0.92 mHa.

In figure 2(c), we further investigate the convergence and fluctuation patterns of $\Delta E_1$ and $\Delta E_2$ along the training process when using respectively TM, VM and VE method. The VE method's curves begin at the training step of $3\times10^4$ in order to ensure a sufficient number of data points for the linear regressions. Compared with the results obtained through the TM and VM methods, the VE method exhibits faster convergence and a much quicker reduction in fluctuation. By the 10⁴th training step, the VE method has achieved convergence for both $\Delta E_1$ and $\Delta E_2$, with smooth curves and no fluctuations. This suggests that using the VE method, VMC training does not require complete convergence, resulting in significant computational cost savings.

To further verify the reliability of the VE method, we tested it on the N₂ dissociation curve, which contained calculations of N₂ molecule at 25 different bond lengths, involving intact and breaking triple bond (figure 2(d)). For comparison, the state-of-the-art r12-MR-ACPF results [47] are also displayed. Considering total energy, all the results show remarkable progress after variance extrapolation, and most results are refined to be within chemical accuracy. As for the non-parallelity error, i.e. the difference between the maximum and the minimum errors along the curve, the VE method enhances the original result of 4.53 mHa to 2.89 mHa, which is comparable to the r12-MR-ACPF result of 2.14 mHa. The corresponding variance extrapolation details are presented in supplementary figure 2. Since it is a bit arbitrary to choose which of the 25 different bond lengths to match the energy variance, the VM method is not tested here.

Finally, it is worth noting that by improving the accuracy of total energy estimations, it becomes easier to accurately derive certain types of relative energy, such as dissociation energy. Normally, we have two strategies for calculating the dissociation energy of a system AB containing two parts A and B,

$$\begin{align} \Delta E & = E\left(A\right)+E\left(B\right)-E\left(AB\right), \end{align} \tag{ 12 }$$

$$\begin{align} \Delta E & = E\left(A+B\right) - E\left(AB\right), \end{align} \tag{ 13 }$$

where E(AB) denotes the energy of the AB system, E(A) and E(B) denote the energies of the individual monomers A and B, and $E(A+B)$ denotes the energy of the system consisting of separate monomers A and B at a sufficient distance from each other. If a difference appears in $E(A) + E(B)$ and $E(A+B)$, then it is often regarded as the appearance of size inconsistency of the calculation, which would cause severe problems when computing useful quantities such as the dissociation energy of molecules. In traditional wavefunction methods the size consistency problem is often induced by an limited ansatz, and in NN-VMC the unbalanced optimization of two different systems is also an important contribution of the problem. Curing the relative energy error is effectively improving the size consistency property of NN-VMC calculations.

In figure 2(b), $\Delta E^* = 2E(\mathrm{N}) - E(r_0)$ and $\Delta E_2 = E(r_2) - E(r_0)$ are the dissociation energies of the N₂ molecule respectively following equations (12) and (13), where $E(\mathrm{N})$ denotes the energy of atom N reported by Chakravorty et al [57]. Apparently, both $E(r_0)$ and $E(r_2)$ would have upward biases due to the limit of ansatz expressiveness. Before the extrapolation, $\Delta E_2$ performs better than $\Delta E^*$ as the two biases cancel each other out to a large extent. Through extrapolation, we obtain an improved estimate of $E(r_0)$ and thus greatly improve the accuracy of $\Delta E^*$ without the requirement of $E(r_2)$, which saves half the computational cost. Generally speaking, E(A) and E(B) are much simpler to obtain precise results than $E(A+B)$ because of the computational complexity of at least $\mathcal{O}(N^3)$ in terms of the system size for accurate QMC calculations. Therefore, the VE method providing better total energy results may help a lot on calculating dissociation energy. Note that here $\Delta E^*$ has no VM result because it only involves one VMC calculation to obtain $E(r_0)$, since $E(\mathrm{N})$ is already established.

To summarize, we show that the VE approach is an effective method for enhancing the accuracy of total energy and relative energy outcomes without incurring additional computational expenses. In contrast, it appears that the VM approach may not provide significant benefits for calculations at this level of accuracy.

Periodic systems are different from molecules, where the system is in principle infinitely large while molecules contain a finite number of atoms. Nevertheless, a similar problem would occur because of an improper wavefunction ansatz or unbalanced optimization. To compute the useful quantities for periodic systems, such as the cohesive energy of solids, it is required to maintain the size extensiveness of the calculation, otherwise the errors are too significant. To verify whether the VE method also works for periodic systems and improves the size extensiveness, we further analyze two other calculations on periodic systems reported in DeepSolid [26], namely the one-dimensional hydrogen chain and the two-dimensional graphene.

Hydrogen chain is a simple but challenging and interesting system. Figure 3(a) shows the interaction energy results of different hydrogen chains containing different atom numbers in a supercell and the corresponding extrapolation result to the thermodynamic limit (TDL). The extrapolation simply follows the quadratic curve fitting with the symmetry axis fixed at 0. For comparison, the TDL auxiliary field QMC (AFQMC) energy result at complete basis set limit is also plotted [50]. Upon observing the magnified results depicted in the inset diagram, it is clear that the TDL energy result utilizing the VE method is more consistent with the AFQMC result. Returning to the four finite size results, the VE method has led to a downward correction on each of them. The larger the system, the bigger the VE method's correction, which is both expected and reasonable given that the accuracy of DeepSolid may progressively worsen on larger systems. The corresponding variance extrapolation details are shown in supplementary figure 3.

Zoom In Zoom Out Reset image size

Graphene is a famous two-dimensional material attracting broad attention over last two decades [58]. The real space and k-space structure of graphene are illustrated in figure 3(b). The positions of three special k points sampled following Monkhorst-Pack mesh [59, 60] are also marked, at which the cohesive energy is calculated. The calculations are carried out on a 2×2 supercell of graphene containing 8 C atoms [26]. The corrected result is regarded as the weighted average of the results at the three k points under twist-averaged boundary condition (TABC) plus an additional structure factor $\mathrm{S}(\mathbf{k})$ [26, 51, 53]. Compared with the experimental result, it is obvious that the VE method enhances the result. The corresponding variance extrapolation details are displayed in supplementary figure 4.

The homogeneous electron gas systems at different densities presented in DeepSolid [26] are also tested, and the corresponding results are shown in figure 4. For these systems, the calculations are performed on a simple cubic cell containing 54 electrons in a closed-shell configuration [26]. In each panel, the correlation energies (i.e. enhancements from the Hatree–Fock method) calculated by various methods are plotted, and an upper bound of the exact result is marked using red dashed lines, which utilizes the variational property of the BF-DMC results [55]. While the initial results prior to extrapolation fall above the upper bound, the results subsequent to extrapolation almost consistently surpass the upper bound. For the densities $\mathrm{\textit{r}_s} = 0.5, 1$ Bohr, where TC-FCIQMC results are available [56], the red full lines indicate high-accuracy benchmarks. The extrapolation results at both densities are closer to the red full lines. The corresponding variance extrapolation details are displayed in supplementary figure 5.

Zoom In Zoom Out Reset image size

The success of the VE method on the periodic systems further illustrates its universality and reliability.