Solving Many-Electron Schrödinger Equation Using Deep Neural Networks

Introduction

An accurate quantum mechanical treatment of the interaction between many electrons and ions is the foundation for modeling physical, chemical, and biological systems. Theoretically, these systems are described by the many-electron Schrödinger equation, which consists of the kinetic and Coulomb interaction terms dirac1929quantum. Solving the Schrödinger equation accurately for real physical systems has been prohibitively difficult due to the high dimensionality of the Hilbert space involved, the high degree of entanglement produced by the electron-electron and electron-ion interactions, and the notorious Pauli exclusion principle imposed on the wave-function, i.e., the wave-function has to change its sign when two identical electrons exchange places pauli1925principle.

Developing efficient algorithms for this problem is among the most heroic endeavors in computational science and has achieved remarkable successes. An incomplete list of the major methodologies developed so far includes the Hartree-Fock (HF) based methods roothaan1960RHF, pople1954UHF, configuration interaction (CI) based methods pople1987QCI, werner1988MRCI, knowles1988MRCI, shiozaki2011MRCI-F12, coupled cluster (CC) based schemes purvis1982CCSD, paldus1999CC, bartlett2007CC, shavitt2009CC, Monte Carlo-based approaches mcmillan1965VMC, ceperley1977VMC, bressanini1999vmc, blankenbecler1981AFQMC, zhang1997AFQMC, zhang2003QMC, van2010diaMC, foulkes2001QMC, and the more recently developed density matrix renormalization group (DMRG) theory white1992DMRG, white1999DMRG, chan2002DMRG, chan2016DMRG, stoudenmire2017slicedDMRG and density matrix embedding theory (DMET) knizia2012DMET, knizia2013DMET, wouters2016practicalDMET. We refer to Refs. grotendorst2000modernQC, szabo2012modernQC, motta2017towards for a more detailed review of these advances.

Of particular interest to this work is the variational Monte Carlo (VMC) scheme which uses variational principle and Monte Carlo sampling to obtain the best parametrized trial wave-function mcmillan1965VMC, ceperley1977VMC, bressanini1999vmc, foulkes2001QMC. Naturally the key component is the representation of the trial wave-functions. The most commonly used trial wave-functions typically consist of an anti-symmetric Slater determinant slater1930note multiplied by a symmetric Jastrow correlation factor jastrow1955many. There have been tremendous efforts on improving the nodal surface (a subspace on which the function value equals zero and across which it changes the sign) of the anti-symmetric part of the trial wave-function and the representability of the symmetric part umrigar1988optimized, umrigar2007alleviation, casula2003geminal, changlani2009CPS.

With remarkable advances in many fields such as computer vision and speech recognition, the deep neural network (DNN) has shown great capacity in approximating high-dimensional functions (see, e.g., review lecun2015deep and the references therein). Furthermore, DNN has been successfully used in solving general high-dimensional partial differential equations han2018solving, e2017deep, berg2017unified, khoo2017solving and certain quantum many-body problems for Bosonic and lattice systems saito2018method, carleo2017solving, gao2017QMB, saito2017QMB, cai2018QMB. However, there have been few attempts to solve the many-electron Schrödinger equations based on DNN, and this constitutes the main objective of this work.

To achieve this, we develop a general and efficient DNN representation for the many-electron wave-function satisfying the Pauli exclusion principle. The resulted trial wave-function can naturally fit into the framework of VMC to optimize the parameters in our model. As preliminary tests, we show that this DNN-based trial wave-function is able to produce reasonably well ground-state energies for some small systems, such as Be, B, LiH, and a chain of 10 hydrogen atoms (H10). In addition, learning from scratch without any prior knowledge and without resorting to a reference of atomic bases, the DNN-based wave-function is able to reproduce the electronic structures of the tested systems. We call the methodology introduced here the Deep WaveFunction method, abbreviated DeepWF. This paper only reports our initial results. There is still a huge room for improvement.

Method

We consider a system of $N$ electrons and $M$ ions, under the Born-Oppenheimer approximation Born1927BO. This system is described by the Hamiltonian

where $\bm{r}=(\bm{r}_{1},\dots,\bm{r}_{N})$ and $\bm{R}=(\bm{R}_{1},\dots,\bm{R}_{M})$ are the coordinates of the electrons and the ions, respectively, and $Z_{I}$ denotes the nuclear charge. Since $\hat{H}$ is spin-independent, it is valid to assume that the first $N_{\uparrow}$ electrons are of spin-up and the remaining $N_{\downarrow}=N-N_{\uparrow}$ electrons are of spin-down foulkes2001QMC. Accordingly, we can write the wave-function in a spin-independent form $\Psi(\bm{r};\bm{R})$. Let $N_{\text{tp}}$ be the number of ion types in the system, then there are $N_{\text{tp}}+2$ types of particles in total, taking into account the electrons of spin-up and spin-down separately.

In this work we restrict our attention to the ground state of the system. The variational principle states that the wave-function associated with the ground state minimizes within the required symmetry the following energy functional

We now discuss how to represent the wave-function $\Psi(\bm{r};\bm{R})$ with DNN. The trial wave-functions are assumed to be real-valued. The primary goal is to ensure that the represented wave-function satisfies the anti-symmetry property. To this end, we decompose the wave-function as

where $\bm{r}^{\uparrow},\bm{r}^{\downarrow}$ denote the positions of spin-up and spin-down electrons, respectively. We require $S(\bm{r},\bm{R})$ to be symmetric and $A^{\uparrow}(\bm{r}^{\uparrow}),A^{\downarrow}(\bm{r}^{\downarrow})$ to be both anti-symmetric. As an analogy, one can view $S(\bm{r};\bm{R})$ as a Jastrow factor and $A^{\uparrow}(\bm{r}^{\uparrow}),A^{\downarrow}(\bm{r}^{\downarrow})$ as Slater determinant-like functions. In practice, $\log(\Psi^{2}(\bm{r};\bm{R}))$ is more convenient for implementing VMC. From this perspective, the introduced decomposition of wave-function becomes

The basic building block for the anti-symmetric function $A^{\uparrow}(\bm{r}^{\uparrow})$ is an ansatz

where $\bm{a}^{\uparrow}_{0}$ is a general two-body anti-symmetric function with $M_{1}$-dimensional output, i.e., $\bm{a}^{\uparrow}_{0}(\bm{r}_{i},\bm{r}_{j})=-\bm{a}^{\uparrow}_{0}(\bm{r}_{j},\bm{r}_{i})$, and the products are component-wise. In (3) we have assumed $N_{\uparrow}\geq 2$. Such representation can be viewed as a generalization of the Laughlin wave-function laughlin1983anomalous, which describes well the anomalous quantum Hall effect. A related model is the so-called correlator product state (CPS) changlani2009CPS, which works well for lattice systems and has natural connections with some recently proposed neural-network quantum states glasser2018neural, clark2018unifying. In practice, we let $\bm{a}^{\uparrow}_{0}(\bm{r}_{i},\bm{r}_{j})=\text{Net}^{\uparrow}_{\text{anti}}(\bm{r}_{i},\bm{r}_{j},r_{ji})-\text{Net}^{\uparrow}_{\text{anti}}(\bm{r}_{j},\bm{r}_{i},r_{ji})$, where $\text{Net}^{\uparrow}_{\text{anti}}(\cdot)$ is represented by a DNN. By convention $r_{ji}=|\bm{r}_{ji}|=|\bm{r}_{j}-\bm{r}_{i}|$ denotes the Euclidean distance between particles $i$ and $j$. Next $\bm{a}^{\uparrow}$ is fed into another DNN $\text{Net}^{\uparrow}_{\text{odd}}$, which outputs a scalar and satisfies $\text{Net}^{\uparrow}_{\text{odd}}(\bm{x})=-\text{Net}^{\uparrow}_{\text{odd}}(-\bm{x})$. We can further adjust its scale and define $A^{\uparrow}(\bm{r}^{\uparrow})$ through $\log|A^{\uparrow}(\bm{r}^{\uparrow})|=f_{\text{odd}}\log|\text{Net}^{\uparrow}_{\text{odd}}(\bm{a}^{\uparrow}(\bm{r}))|$ where $f_{\text{odd}}$ is a positive scalar factor. The sign of $A^{\uparrow}(\bm{r}^{\uparrow})$ is taken as the same as $\text{Net}^{\uparrow}_{\text{odd}}(\bm{a}^{\uparrow}(\bm{r}^{\uparrow}))$ such that $A^{\uparrow}(\bm{r}^{\uparrow})$ inherits the anti-symmetry. The same construction introduced above are applied to spin-down electrons and yield $A^{\downarrow}(\bm{r}^{\downarrow})$.

While DNN-based anti-symmetric function has seldom been investigated in the literature, DNN-based symmetric function has gained wide attentions recently in modeling many-body potential energy surface behler2007generalized, han2017deep, zhang2018deep, zhang2018end, schutt2017schnet, free energy surface zhang2018reinforced, zhang2018deepcg, schneider2017stochastic, etc. For instance, the Deep Potential-Smooth Edition (DeepPot-SE) model is able to efficiently describe the interatomic potential energy of both finite and extended systems zhang2018end. A crucial step there is to faithfully map the input atomic coordinates onto a symmetry-preserving feature space. Therefore, it shares the symmetric property of the Jastrow factor, and might become a more general Jastrow factor that is able to better capture the many-body correlations between the electrons and ions.

We denote all the parameters in the DNNs ($\text{Net}^{\uparrow}_{\text{anti}}$, $\text{Net}^{\downarrow}_{\text{anti}}$, $\text{Net}^{\uparrow}_{\text{odd}}$, $\text{Net}^{\downarrow}_{\text{odd}}$, $\text{Net}_{\text{ebd}}^{\alpha_{i}\alpha_{j}}$, and $\text{Net}_{\text{fit}}^{\alpha_{i}}$), the two scalar factors ($f_{\text{odd}}$ and $f_{\text{dec}}$) together as $\bm{\theta}$ and the associated trial wave-function as $\Psi_{\bm{\theta}}(\bm{r})$ (below for convenience we ignore the dependence on $\bm{R}$, the clamped ion positions). $\bm{\theta}$ are initialized randomly from a Gaussian distribution without any pre-training on pre-calculated wave-functions. We use VMC to optimize the trial wave-function. Specifically, we keep track of $N_{\text{wk}}$ walkers to approximate the squared wave-function through an empirical distribution. The learning process consists of two phases, the sampling phase and the optimization phase, which we proceed with alternatively. In the sampling phase of step $t$, we run several steps of the Metropolis-Hasting algorithm to update the positions of the walkers, according to a target probability proportional to $\Psi^{2}_{\bm{\theta}_{t}}(\bm{r})$. In the optimization phase of step $t$, we aim to minimize the second moment of the local energy $E_{\text{loc}}(\bm{r})$ with respect to a fixed reference energy $E_{\text{ref}}$. Here the local energy is defined as $E_{\text{loc}}(\bm{r})\coloneqq\hat{H}\Psi_{\bm{\theta}}(\bm{r})/\Psi_{\bm{\theta}}(\bm{r})$. Accordingly, the objective function has the explicit form

In numerical computation, in order to reduce the variance when evaluating the objective function, we use a technique called correlated sampling. Assuming the walkers’ current positions $\{\bm{r}^{(1)},\bm{r}^{(2)},\dots,\bm{r}^{(N_{\text{wk}})}\}$ approximate the distribution $\Psi_{\bm{\theta}_{t}}^{2}(\bm{r})$ well, we define a reweighting factor $w_{i}\coloneqq\Psi_{\bm{\theta}}(\bm{r}^{(i)})/\Psi_{\bm{\theta}_{t}}(\bm{r}^{(i)})$ and rewrite the objective function as

Note here we take $\bm{\theta}_{t}$ as constants and view both $w_{i}$ and $E_{\text{loc}}(\bm{r}^{(i)})$ as functions of the parameters $\bm{\theta}$ given the walker’s position $\bm{r}^{(i)}$. The gradients $\nabla_{\bm{\theta}}\Omega_{E_{\text{ref}}}(\bm{\theta})$ are computed by the backpropagation algorithm and used to update the parameters through $\bm{\theta}_{t+1}=\bm{\theta}_{t}-\eta\nabla_{\bm{\theta}}\Omega_{E_{\text{ref}}}(\bm{\theta}_{t})$, with $\eta$ being the learning rate. Motivated by the idea of stochastic gradient descent (SGD), we actually evaluate $\nabla_{\bm{\theta}}\Omega_{E_{\text{ref}}}(\bm{\theta})$ with a random batch of data $i\in\mathcal{B}\subseteq\{1,\dots,N_{\text{wk}}\}$ and update $\bm{\theta}_{t}$ with a few steps.

Results and Discussion

Using the algorithm described above, for a series of benchmark systems, we obtain the optimized parameters $\bm{\theta}^{*}$ and run further Monte Carlo simulations for several statistical properties. The network structure and training scheme are detailed in the appendix. In Table 1, we present the ground state energies of H2, He, LiH, Be, B, and a chain of 10 hydrogen atoms under open boundary conditions. Note that for H2 and He we don’t have the anti-symmetric part in the wave-function. Overall the ground-state energies obtained by DeepWF show great consistency with the benchmarks. However, as the number of electrons increases, the accuracy deteriorates.

In Fig. 2, we plot the potential energy curves of H2 and H10 as a function of bond length. The result for H2 is very close to the benchmark. For H10, the DeepWF shows a good relative energy, with the correct prediction of local minimum, but is still a bit above the benchmark result.

Finally, Fig. 3 plots the spatial distributions of electrons in different systems. In the case of radial distribution function of electrons in a Be atom (Fig. 3 (a)), it is remarkable that the DeepWF learns from scratch the shell structure of the electrons and shows great consistency with the result obtained using the ccpv5z basis woon1995gaussian. As a comparison, the sto6g basis hehre1969sto is not enough to describe the electronic dispersion. In the case of axial distribution function of electrons in a LiH molecule (Fig. 3 (b)), DeepWF and the other two methods show excellent agreement.

What is appealing to us is the simplicity of the proposed approach. Obviously there is a huge room for further improvement. In terms of the representation of DeepWF, the symmetric part is relatively general, considering the success of a similar version in representing the inter-atomic potential energy surface. However, the anti-symmetric ansatz, although appealing due to its quadratic scaling, might not be sufficient in representing the electronic correlations caused by the Pauli exclusion rule. Second, in terms of optimization, techniques for accelerating VMC through more efficient sampling (see e.g. lee2011strategies, dewing2000improved) can be directly adapted into our DeepWF method. Optimization method other than the SGD-like correlated sampling method can also be employed. In any case, we hope that ideas presented here will add some ammunition to the heroic endeavor of attempting to solve the many-body Schrödinger equation.

Acknowledgments

The authors acknowledge M. Motta for helpful discussions. This work is supported in part by Major Program of NNSFC under grant 91130005, ONR grant N00014-13-1-0338 and NSFC grant U1430237. We are grateful for the computing time provided by the High-performance Computing Platform of Peking University and the TIGRESS High Performance Computing Center at Princeton University.

References

Appendix A Details of the Training Procedure for Each System