A Real Neural Network State for Quantum Chemistry

Abstract

The restricted Boltzmann machine (RBM) has recently been demonstrated as a useful tool to solve the quantum many-body problems. In this work we propose tanh-FCN, which is a single-layer fully connected neural network adapted from RBM, to study ab initio quantum chemistry problems. Our contribution is two-fold: (1) our neural network only uses real numbers to represent the real electronic wave function, while we obtain comparable precision to RBM for various prototypical molecules; (2) we show that the knowledge of the Hartree-Fock reference state can be used to systematically accelerate the convergence of the variational Monte Carlo algorithm as well as to increase the precision of the final energy.

1. Introduction

Ab initio electronic structure calculations based on quantum-chemical approaches (Hartree–Fock theory and post-Hartree–Fock methods) have been successfully applied in molecular systems [1]. For strongly correlated many-electron systems, the exponentially growing Hilbert space size limits the application scale of most numerical algorithms. For example, the full configuration interaction (FCI), which takes the whole Hilbert space into account, is currently limited within around 24 orbitals and 24 electrons [2]. The density matrix renormalization group (DMRG) algorithm [3,4] has been used to solve larger chemical systems of several tens of electrons [5,6]; however, it is essentially limited by the expressive power of its underlying variational ansatz: the matrix product state (MPS), which is a special instance of the one-dimensional tensor network state [7]. Therefore, DMRG could also be extremely difficult to to use to approach even larger systems. The coupled cluster (CC) [8,9] method expresses the exact wave function in terms of an exponential form of a variational wave function ansatz, and a higher level of accuracy can be obtained by considering electronic excitations up to doublets in CCSD or triplets in CCSD(T). In practice, it is often accurate with a durable computational cost and is thus considered as the “gold standard” in electronic structure calculations. However, the accuracy of the CC method is only restricted in studying weakly correlated systems [10]. The multi-configuration self-consistent field (MCSCF) [11,12,13] method is crucial for describing molecular systems containing nearly degenerate orbitals. It introduces a small number of (active) orbitals; then, the configuration interaction coefficients and the orbital coefficients are optimized to minimize the total energy of the MCSCF state. It has been applied to systems with around 50 active orbitals [14], but they are still limited by the exponential complexity that grows with the system size.

In recent years, the variational Monte Carlo (VMC) method in combination with a neural network ansatz for the underlying quantum state (wave function) [15], referred to as the neural network quantum states (NNQS), has been demonstrated to be a scalable and accurate tool for many-spin systems [16,17,18] and many-fermion systems [19]. NNQS allow very flexible choices of the neural network ansatz, and with an appropriate variational ansatz, it could often achieve comparable or higher accuracy compared to existing methods. NNQS has also been applied to solve ab-initio quantum chemistry systems in real space with up to 30 electrons [20,21,22], as well as in discrete basis after second quantization [23,24,25]. Up to now, various neural networks have been used, such as the restricted Boltzmann machine (RBM) [15], the convolutional neural network [16], recurrent neural networks [26] and the variational auto-encoder [25]. In all of those neural networks, the RBM is a very special instance in that: (1) it has a very simple structure that contains only a fully connected dense layer plus a nonlinear activation; (2) with such a simple structure, RBM can be more expressive than MPS [27]; in fact, it is equivalent to certain two-dimensional tensor network states [28] and can even represent certain quantum states with volume-law entanglement [29]. In practice, RBM achieves comparable accuracy to other more sophisticated neural networks for complicated applications such as frustrated many-spin systems [30,31].

For the ground state of molecular systems, the wave function is real. However, if one uses a real RBM as the variational ansatz for the wave function, then all of the amplitudes of the wave function will be positive, which means that it may be good for ferromagnetic states but will be completely wrong for anti-ferromagnetic states. Therefore, even for real wave functions one would have to use complex RBMs or two RBMs [32] in general. In this work, we propose a neural network with real numbers that is slightly modified from the RBM, such that its output can be both positive and negative, and use it as the neural network ansatz to solve quantum chemistry problems. To accelerate convergence of the VMC iterations, we explicitly use the Hartree–Fock reference state as the starting point for the Monte Carlo sampling after a number of VMC iterations such that the wave function ansatz has become sufficiently close to the ground state. We show that this technique can generally improve the convergence and the precision of the final result, even when using other neural networks. Our paper is organized as follows. In Section 2, we present our neural network ansatz. In Section 3, we present our numerical results demonstrating the effectiveness of our neural network ansatz and the technique of initializing the Monte Carlo sampling with the Hartree–Fock reference state. We conclude in Section 4.

2. Methods

2.1. Real Neural Network Ansatz

Before we introduce our model, we first briefly review the RBM used in NNQS. For a classical many-spin system, one could embed the system into a larger one consisting of visible spins (corresponding to the system) and hidden spins with the total (classical) Hamiltonian

$$\begin{array}{c}\hfill \mathcal{H}=\sum _{j=1}^{N_v}{a}_j{x}_j+\sum _{i=1}^{N_h}{b}_i{h}_i+\sum _{i,j}{W}_{ij}{h}_i{x}_j,\end{array}$$

(1)

where $x_j$ represents the visible spin and $h_i$ the hidden spin. $N_v$ and $N_h$ are the number of visible and hidden spins, respectively. The coefficients $\theta =\{a,b,W\}$ are variational parameters of the Hamiltonian. Since there is no coupling between the hidden spins, one could explicitly integrate them out and obtain the partition function of the system $\mathcal{Z}$ as

$$\begin{array}{c}\hfill \mathcal{Z}=\sum _{\mathbf{x}}p\left(\mathbf{x}\right),\end{array}$$

(2)

with $\mathbf{x}=\{x_1,x_2,\dots ,x_{N_v}\}$ a particular configuration and $p\left(\mathbf{x}\right)$ the unnormalized probability (in case of real coefficients) of x, which can be explicitly written as

$$\begin{array}{cc}\hfill p\left(\mathbf{x}\right)& =\sum _{\mathbf{h}}{e}^{\mathcal{H}}\hfill \\ \hfill & =e^{{\sum }_{j=1}^{N_v}{a}_j{x}_j}\times \prod _{i=1}^{N_h}2cosh(b_i+\sum _{j=1}^{N_v}{W}_{ij}{x}_j).\hfill \end{array}$$

(3)

When using RBM as a variational ansatz for the wave function of a quantum many-spin system, $p\left(\mathbf{x}\right)$ is interpreted as the amplitude (instead of the probability) of the configuration x. Equation (3) can be seen as a single-layer fully connected neural work that accepts a configuration (a vector of integers) as input and outputs a scalar. For real coefficients, the output will always be positive by definition; therefore, one generally has to use complex coefficients even for real wave functions. In this work, we slightly change Equation (3) as follows so as to be able to output any real numbers with a real neural network:

$$\begin{array}{c}\hfill p\left(\mathbf{x}\right)=tanh\left(\sum _{j=1}^{N_v}{a}_j{x}_j\right)\times \prod _{i=1}^{N_h}2cosh(b_i+\sum _{j=1}^{N_v}{W}_{ij}{x}_j).\end{array}$$

(4)

In the following, we will write $p\left(\mathbf{x}\right)$ as ${\Psi }_{\theta }\left(\mathbf{x}\right)$ to stress its dependence on the variational parameters and that it is interpreted as a wave function instead of a probability distribution. We will also refer to our neural network in Equation (4) as tanh-FCN since it contains a fully connected layer followed by hyperbolic tangent as the activation function. The difference between RBM and tanh-FCN is demonstrated in Figure 1.

2.2. Variational Monte Carlo

The electronic Hamiltonian ${\hat{H}}^e$ of a chemical system can be written in a second-quantized formulation:

$$\begin{array}{c}\hfill {\hat{H}}^e=\sum _{p,q}{h}_q^p{a}_p^{†}{a}_q+\frac{1}{2}\sum _{\begin{array}{c}p,q\\ r,s\end{array}}{g}_{rs}^{pq}{a}_p^{†}{a}_q^{†}{a}_r{a}_s\end{array}$$

(5)

where $h_q^p$ and $g_{rs}^{pq}$ are one- and two-electron integrals in molecular orbital basis, and $a_p^{†}$ and $a_q$ in the Hamiltonian are the creation and annihilation operators. To treat the fermionic systems, we first use the Jordan–Wigner transformation to map the electronic Hamiltonian to a sum of Pauli operators, following Ref. [23], and then use our tanh-FCN in Equation (4) as the ansatz for the resulting many-spin system. The resulting spin Hamiltonian $\hat{H}$ can generally be written in the following form:

$$\begin{array}{c}\hfill \hat{H}=\sum _i{c}_i\prod _{j=1}^N{\sigma }_j^{v_{i,j}},\end{array}$$

(6)

where $N=N_v$ is the number of spins, $c_i$ is a real coefficient, and ${\sigma }_j^{v_{i,j}}$ is a single spin Pauli operator acting on the j-th spin ($v_{i,j}\in \{0,1,2,3\}$ and ${\sigma }^0=I$, ${\sigma }^1={\sigma }^x$, ${\sigma }^2={\sigma }^y$, ${\sigma }^3={\sigma }^z$).

Given the wave function ansatz ${\Psi }_{\theta }\left(\mathbf{x}\right)$, the corresponding energy can be computed as

$$\begin{array}{c}\hfill E\left(\theta \right)=\frac{\langle {\Psi }_{\theta }\left|\hat{H}\right|{\Psi }_{\theta }\rangle }{\langle {\Psi }_{\theta }|{\Psi }_{\theta }\rangle }=\frac{{\sum }_{\mathbf{x}}{E}_{\mathrm{loc}}\left(\mathbf{x}\right){\left|{\Psi }_{\theta }\left(\mathbf{x}\right)\right|}^2}{{\sum }_{\mathbf{y}}{\left|{\Psi }_{\theta }\left(\mathbf{y}\right)\right|}^2},\end{array}$$

(7)

where the “local energy” $E_{\mathrm{loc}}\left(\mathbf{x}\right)$ for a configuration x is defined as

$$\begin{array}{c}\hfill E_{\mathrm{loc}}\left(\mathbf{x}\right)=\sum _{{\mathbf{x}}^{\prime }}\frac{{\Psi }_{\theta }\left({\mathbf{x}}^{\prime }\right)}{{\Psi }_{\theta }\left(\mathbf{x}\right)}{H}_{{\mathbf{x}}^{\prime }\mathbf{x}},\end{array}$$

(8)

with $H_{{\mathbf{x}}^{\prime }\mathbf{x}}=\langle {\mathbf{x}}^{\prime }|\hat{H}|\mathbf{x}\rangle$. The VMC algorithm evaluates Equation (7) approximately using Monte Carlo sampling, namely,

$$\begin{array}{c}\hfill \tilde{E}\left(\theta \right)=\langle E_{\mathrm{loc}}\rangle ,\end{array}$$

(9)

where the average is over a set of samples $\{{\mathbf{x}}^1,{\mathbf{x}}^2,\dots ,{\mathbf{x}}^{N_s}\}$ ($N_s$ is the total number of samples), generated from the probability distribution $|{\Psi }_{\theta }{\left(\mathbf{x}\right)|}^2$. $\tilde{E}\left(\theta \right)$ will converge to $E\left(\theta \right)$ if $N_s$ is large enough. In this work, we use the Metropolis–Hastings sampling algorithm to generate samples [33]. A configuration x is updated using the SWAP operation between nearest-neighbor pairs of spins to preserve the electron-number conservation. We also use the natural gradient of Equation (9) for the stochastic gradient descent algorithm in VMC, namely, the parameters are updated as

$$\begin{array}{c}\hfill {\theta }^{k+1}={\theta }^k-\alpha S^{-1}F,\end{array}$$

(10)

where k is the number of iterations, $\alpha$ is the learning rate ($\alpha$ is dependent on k in general), S is the stochastic reconfiguration matrix [34,35], and F is the gradient of Equation (9). Concretely, S and F are computed by

$$\begin{array}{c}\hfill S_{ij}\left(k\right)=\langle O_i^{*}{O}_j\rangle -\langle O_i^{*}\rangle \langle O_j\rangle ,\end{array}$$

(11)

and

$$\begin{array}{c}\hfill F_i\left(k\right)=\langle E_{\mathrm{loc}}{O}_i^{*}\rangle -\langle E_{\mathrm{loc}}\rangle \langle O_i^{*}\rangle \end{array}$$

(12)

respectively, with $O_i\left(\mathbf{x}\right)$ defined as

$$\begin{array}{c}\hfill O_i\left(\mathbf{x}\right)=\frac{1}{{\Psi }_{\theta }\left(\mathbf{x}\right)}\frac{\partial {\Psi }_{\theta }\left(\mathbf{x}\right)}{\partial {\theta }_i}.\end{array}$$

(13)

In general, S can be non-invertible, and a simple regularization is to add a small shift to the diagonals of S, namely, using $S^{reg}=S+ϵI$ instead of S in Equation (10), with $ϵ$ a small number. The calculation of S can become the bottleneck in case the number of parameters is too large. This issue could be leveraged by representing S as a matrix function instead of building it explicitly [36], or by freezing a large portion of S during each iteration similar to DMRG [37]. Here, this is not a significant concern because we use at most about 1000 parameters to specify the network. To further enhance the stability of the algorithm, we add the contribution of an L2 regularization term when evaluating the gradient in Equation (10), that is, instead of directly choosing F as the gradient of $\tilde{E}\left(\theta \right)$, F is chosen as the gradient of the function $\tilde{E}\left(\theta \right)+{\lambda \left|\right|\theta \left|\right|}^2$ instead where $\left|\right|·{\left|\right|}^2$ means the square of the Euclidean norm. In this work, we choose $ϵ=0.02$ and $\lambda ={10}^{-3}$ for our numerical simulations if not particularly specified.

3. Results

3.1. Training Details

In this work, we use the Adam optimizer [38] for the VMC iterations, with an initial learning rate of $\alpha =0.001$, and the decay rates for the first- and second-moment are ${\beta }_1=0.9$, ${\beta }_2=0.99$, respectively. For the Metropolis–Hastings sampling, we will use a fixed $N_s=4\times {10}^4$ for our numerical simulations if not particularly specified (in principle, one should use a larger $N_s$ for larger systems; however, in this work we focus on molecular systems with at most 30 qubits). We will also use a thermalization step of $N_{th}=2\times {10}^4$ (namely, throwing away $N_{th}$ samples starting from the initial state). To avoid auto-correlation between successive samples we will only pick one out of every $10N_v$ samples. In addition, for each simulation we run 8 Markov chains, and the energy is chosen to be the lowest of them. Since the energy will always contain some small fluctuations when $N_s$ is not large enough, the final energy is evaluated by averaging over the energies of the last 20 VMC iterations.

3.2. Effect of Hidden Size

We first study the effect of $N_h$, which essentially determines the number of parameters and thus the expressivity of our tanh-FCN (analogously to RBM). The result is shown in Figure 2, where we have taken the N2 molecule as an example. We can see that by enlarging $N_h$, the precision of tanh-FCN can be systematically improved. With $N_h=4N_v=80$, we can already obtain a final energy that is lower than the CCSD results.

3.3. Potential Energy Surfaces

Now we demonstrate the accuracy of our tanh-FCN by studying the potential energy surfaces of the two molecules ${\mathrm{H}}_2$ and LiH in the STO-3G basis, as shown in Figure 3(a1,b1). We can see that for both molecules under different bond lengths, our simulation can reach lower or very close to the chemical precision, namely error within $1.6\times {10}^{-3}$ Hatree (Ha) or 1 kcal/mol (CCSD results are extremely accurate for these two molecules). To demonstrate of the effectiveness of our method for weakly correlated systems, we have also studied the potential energy surfaces of the two inert gas dimers ${\mathrm{He}}_2$ and ${\mathrm{Ne}}_2$ for completeness, which are shown in Figure 3(c1,d1). We can see that in the later cases our tanh-FCN can converge extremely well with the FCI results. We note that for ${\mathrm{Ne}}_2$ one may need to use a very large basis set to faithfully reproduce the actual potential energy surface, while here we have used the minimal STO-3G basis set due to the limitation of our current implementation.

3.4. Final Energies for Several Molecular Systems

We further compare the precision of tanh-FCN with RBM and CCSD for several small-scale molecules, which are shown in

Table 1. List of molecules and the ground state energies computed using RBM, tanh-FCN, and CCSD. The FCI energy is also shown as a reference. The column $N_v$ shows the number of qubits. We have used $N_h/N_v=2$ for all of the molecules studied.

Molecule	${\mathit{N}}_{\mathit{v}}$	RBM [23]	tanh-FCN	CCSD	FCI
${\mathbf{H}}_2$	4	$-1.1373$	$-1.1373$	$-1.1373$	$-1.1373$
Be	10	-	$-14.4033$	$-14.4036$	$-14.4036$
C	10	-	$-37.2184$	$-37.1412$	$-37.2187$
${\mathbf{Li}}_2$	20	-	$-14.6641$	$-14.6665$	$-14.6666$
LiH	12	$-7.8826$	$-7.8816$	$-7.8828$	$-7.8828$
${\mathbf{NH}}_3$	16	$-55.5277$	$-55.5101$	$-55.5279$	$-55.5282$
${\mathbf{H}}_2\mathbf{O}$	14	$-75.0232$	$-75.0021$	$-75.0231$	$-75.0233$
${\mathbf{C}}_2$	20	$-74.6892$	$-74.6134$	$-74.6744$	$-74.6908$
${\mathbf{N}}_2$	20	$-107.6767$	$-107.622$	$-107.6716$	$-107.6774$
${\mathbf{CO}}_2$	30	-	$-185.1247$	$-184.8927$	$-185.2761$

. For these simulations we have used $N_h/N_v=2$, while the RBM results are taken from Ref. [23]. As a proof of principle demonstration, we have mostly used the STO-3G basis set. However, we have also considered the LiH molecule in a larger basis set (6-31G) as well as in the localized molecular basis set (we have used the canonical molecular basis set for the rest ones) to show the effectiveness of our method in more general cases. Unlike DMRG which uses a one-dimensional matrix product state as the wave function ansatz, our neural network ansatz has an all to all structure which could represent certain volume-law quantum states [29], therefore it does not significantly rely on localized orbitals and it seems that using a localized basis set could not improve the precision or significantly reduce the computational cost for us. From the runtime performance point of view, properly selected orbital localization scheme could reduce the number of Pauli terms in the Hamiltonian thus effectively accelerate the algorithm. However, this improvement is not universally achieved. For example, the number of Pauli terms for an equi-spaced H₁₂ molecule with R(H-H) = 2.5 Angstrom can be reduce from 14,905 to 4377 if natural atomic orbitals (NAO) [39] are used, while this number is increased to 23,109 if the bond length is changed to 0.7 Angstrom. An optimal choice of orbital localization methods is usually system-specific [40] and requires benchmark for the neural network ansatz.

These results show that even with a relatively small number of parameters and a real neural network, we can still obtain the ground state energies of a wide variety of molecules to very high precision (close to or lower than the CCSD energies). In the meantime, we note that the energies obtained using tanh-FCN is not as accurate as those obtained using RBM, however the computational cost of tanh-FCN is at least two times lower than RBM under with the same $N_h$ and we could relatively easily study larger systems such as ${\mathrm{CO}}_2$ with 30 qubits. It should be noted that the total energy depends on the basis set size and the basis type, in principle, we should use a large basis set to obtain more reliable results.

3.5. Effect of Hartree–Fock Re-Initialization

There are generally two ingredients which would affect the effectiveness of the NNQS algorithm: (1) the expressivity of the underlying neural network ansatz and (2) the ability to quickly approach the desired parameter regime during the VMC iterations. The former is dependent on an intelligent choice of the neural network ansatz. The effect of the latter is more significant for larger systems, and one generally needs to use a knowledged starting point such as transfer learning [41,42] for the VMC algorithm to guarantee success. For molecular systems, it is difficult to explore transfer learning since the knowledge for different molecules can hardly be shared. However, for molecular systems the Hartree–Fock reference state may have a large overlap with the exact ground state and is often used as a first approximation of the ground state. Here, we show that for quantum chemistry problems the ability to reach faster the ground state can be improved by using the knowledge of the Hartree–Fock reference state. Concretely, during the VMC iterations, after the energies have become sufficiently close to the ground state energy, we stop using random initialization for our Metropolis–Hastings sampling and use the Hartree–Fock reference state instead (Hartree–Fock re-initialization). The effect of the Hartree–Fock re-initialization is demonstrated in Figure 4, where we have taken the H₂O molecule as our example. To show the versatility of the Hartree–Fock re-initialization, we demonstrate its effect for RBM as well. We can see that for both tanh-FCN and RBM, using Hartree–Fock re-initialization after a number of VMC iterations can greatly accelerate the convergence and reach a lower ground state energy than using random initialization throughout the VMC optimization. We can also see that for the H₂O molecule, tanh-FCN is less accurate than RBM using the same $N_h$, which is probably due to the fact that under the same $N_h$, tanh-FCN has a different expressive power to RBM for H₂O.

4. Conclusions

We propose a fully connected neural network inspired from the restricted Boltzmann machine to solve quantum chemistry problems. Compared to RBM, our tanh-FCN is able to output both positive and negative numbers even if the parameters of the network are purely real. As a result, we can directly study quantum chemistry problems using tanh-FCN with real numbers. In our numerical simulation, we demonstrate that tanh-FCN can be used to compute the ground states with high accuracy for a wide range of molecular systems with up to 30 qubits. In addition, we propose to explicitly use the Hartree–Fock reference state as the initial state for the Markov chain sampling used during the VMC algorithm and demonstrate that this technique can significantly accelerate the convergence and improve the accuracy of the final energy for both tanh-FCN and RBM. Our method could be used in combination with existing high performance computing devices that are well optimized for real numbers, such as to provide a scalable solution for large-scale quantum chemistry problems.