Ab initio solution of the many-electron Schrödinger equation with deep neural networks

Given access to accurate solutions of the many-electron Schrödinger equation, nearly all chemistry could be derived from first principles. Exact wave functions of interesting chemical systems are out of reach because they are NP-hard to compute in general, but approximations can be found using polynomially scaling algorithms. The key challenge for many of these algorithms is the choice of wave function approximation, or Ansatz, which must trade off between efficiency and accuracy. Neural networks have shown impressive power as accurate practical function approximators and promise as a compact wave-function Ansatz for spin systems, but problems in electronic structure require wave functions that obey Fermi-Dirac statistics. Here we introduce a novel deep learning architecture, the Fermionic neural network, as a powerful wave-function Ansatz for many-electron systems. The Fermionic neural network is able to achieve accuracy beyond other variational quantum Monte Carlo Ansatz on a variety of atoms and small molecules. Using no data other than atomic positions and charges, we predict the dissociation curves of the nitrogen molecule and hydrogen chain, two challenging strongly correlated systems, to significantly higher accuracy than the coupled cluster method, widely considered the most accurate scalable method for quantum chemistry at equilibrium geometry. This demonstrates that deep neural networks can improve the accuracy of variational quantum Monte Carlo to the point where it outperforms other ab initio quantum chemistry methods, opening the possibility of accurate direct optimization of wave functions for previously intractable many-electron systems.

Figure 1

The Fermionic neural network (FermiNet). Top: Global architecture. Features of one or two electron positions are inputs to different streams of the network. These features are transformed through several layers, a determinant is applied, and the wave function at that position is given as output. Bottom: Detail of a single layer. The network averages features of electrons with the same spin together, then concatenates these features to construct an equivariant function of electron position at each layer.

Figure 2

Comparison of the FermiNet against the Slater-Jastrow Ansatz, with and without backflow. (a) First-row atoms with a single determinant. Baseline numbers are from Chakravorty et al. [34]. The Slater-Jastrow neural network yields slightly lower energies than VMC with a conventional Slater-Jastrow Ansatz, while the FermiNet is substantially more accurate. (b) The CO and ${\mathrm{N}}_2$ molecules (bond lengths 2.17328 $a_0$ and 2.13534 $a_0$, respectively) with increasing numbers of determinants. All-electron CCSD(T)/CBS results are used as the baseline. No matter how many determinants are used, the FermiNet far exceeds the accuracy of the Slater-Jastrow Net.

Figure 3

Optimization progress for first-row atoms, ${\mathrm{H}}_2$, LiH and the hydrogen chain with KFAC (blue) vs. ADAM (orange). The qualitative advantage of KFAC is clear. For clarity, the median energy over the last 10% of iterations is shown. Note that the small overshoot with KFAC between ${10}^3$ and ${10}^4$ iterations is due to the slow equilibration of the MCMC chain and goes away with a larger Metropolis-Hastings proposal step size.

Figure 4

The ${\mathrm{H}}_4$ rectangle, $R=3.2843a_0$. Coupled cluster methods incorrectly predict a cusp and energy minimum at $\mathrm{\Theta }={90}^{\circ }$, while the FermiNet approach agrees with exact FCI results.

Figure 5

The dissociation curve for the nitrogen triple-bond. The difference from experimental data [53] is given in the main panel. In the region of largest UCCSD(T) error, the FermiNet prediction is comparable to highly accurate $r_{12}$-MR-ACPF results [54].

Figure 6

The ${\mathrm{H}}_{10}$ chain. All energies except the FermiNet are taken from Motta et al. (2017) [55]. The absolute energies (inset) cannot be distinguished by eye. The difference from highly accurate MRCI+Q+F12 results are shown in the main panel, where the shaded region indicates an estimate of the basis-set extrapolation error. The errors in the coupled cluster and conventional VMC energies are large at medium atomic separations but the FermiNet remains comparable to AFQMC at all separations. See also Appendix pp5 for data on larger separations.

Figure 7

Electron dimerization in the hydrogen chain. The gap between alternating minima of the density shrinks with increasing nuclear separation.

Figure 8

The pair-correlation function $1-g(\mathbf{r},{\mathbf{r}}^{\prime })$ for the neon atom, where $n(\mathbf{r},{\mathbf{r}}^{\prime })=n\left(\mathbf{r}\right)g(\mathbf{r},{\mathbf{r}}^{\prime })n\left({\mathbf{r}}^{\prime }\right)$. Different columns show the hole for electrons of the same spin (left), different spins (middle), or all electrons (right). Different rows show the hole when ${\mathbf{r}}^{\prime }$ is between 0 and 0.3 Bohr radii from the nucleus.

Figure 9

Comparison of the runtime for one optimization iteration on atoms up to zinc. Polynomial regressions up to fourth order are fit to the data. The small difference between the cubic and quartic fit suggests that the determinant computation is not the dominant factor at this scale.

Figure 10

Effects of network architecture on FermiNet performance on the hydrogen chain ${\mathrm{H}}_{10}$ with separation $R=2.0a_0$. Each point is one run of the same model. (a) Effect of network depth. The marginal improvement with 4 layers is small but not zero. (b) Effect of number of hidden units in the one-electron stream. There is a continuous improvement with wider layers, with the error decreasing roughly as $\mathcal{O}\left(N^{-0.395\pm 0.067}\right)$. (c) Effect of number of hidden units in the two-electron stream. The accuracy plateaus above 16 units.

Figure 11

Evaluation of the FermiNet wave function for the helium atom. The second electron is clamped at position $(0.5,0,0)a_0$ and the first electron is moved along the path $(x,0,0)a_0$, through both the nucleus and the second electron (top), and along the path $(0.5cos\theta ,0.5sin\theta ,0)a_0$, through the second electron (bottom).