Solving the Electronic Schrödinger Equation by Pairing Tensor-Network State with Neural Network Quantum State

Abstract

Neural network methods have shown promise for solving complex quantum many-body systems. In this study, we develop a novel approach through incorporating the density-matrix renormalization group (DMRG) method with the neural network quantum state method. The results demonstrate that, when tensor-network pre-training is introduced into the neural network, a high efficiency can be achieved for quantum many-body systems with strong correlations.

1. Introduction

In 1929, Dirac wrote his famous remark: “The fundamental laws necessary for the mathematical treatment of a large part of physics and the whole of chemistry are thus completely known, and the difficulty lies only in the fact that application of these laws leads to equations that are too complex to be solved”. Today Dirac’s vision has turned into one of the greatest challenges in modern electronic structure calculation. To address this, a variety of methods have been developed to tackle the many-body Schrödinger equation for realistic chemical systems. Methods like full configuration interaction (FCI) have been considered, which encompass the entirety of Hilbert space, growing exponentially and thus limiting simulation scales. More scalable alternatives exist for approximating the precise ground state, including perturbation theory and truncated configuration interaction methods (such as CISD, MCSCF) [1], as well as the coupled-cluster (CC) method [2]. However, these techniques primarily offer accuracy in analyzing weakly correlated chemical systems [3] or face challenges in dealing with very large systems.

The density matrix renormalization group (DMRG) method was initially formulated in 1992 as an innovative approach for evaluating the one-dimensional strongly correlated quantum lattices [4]. Later, this method was introduced to ab initio quantum chemistry for solving systems with strong electron correlation effect [5]. The computational complexity of the DMRG method is characterized by a polynomial order of $O(k^4{m}^2+k^3{m}^3)$, which makes it capable of manipulating a relatively large active space. Here, k is the number of orbitals, which is decided by the active space, and m is called the bond dimension, which is manually assigned to make the computational cost acceptable. When m is relatively small, DMRG can generate a cursory approximation in a very short time. Through certain transformation methods, this cursory approximation can form a good starting point for further evaluations. There are several transformation methods used to transfer the DMRG wave function into a complete active space-configuration interaction (CASCI) type wave function, including the sampling-reconstructed complete active space algorithm [6] and the entanglement-driving genetic algorithm (EDGA) [7].

In the year 2017, Carleo and Troyer introduced the neural network quantum state (NNQS) algorithm, specifically for systems with multiple spins. This approach models the wave function through a neural network and employs stochastic optimization of parameters via the variational Monte Carlo (VMC) method [8]. Specifically, their work shows that a straightforward neural network model, the restricted Boltzmann machine (RBM)—essentially a dense layer with a nonlinear activation function—can match the accuracy of current tensor network methods. So far, the effectiveness of the NNQS approach has been established across a broad range of problems involving multiple spins [8,9,10,11,12,13] and molecular systems [14,15,16,17,18,19,20].

However, in the NNQS method, learning a good representation of the wave function usually requires many iterations, since the neural network models have a large number of parameters, which are hard learn quickly. To deal with this problem, the neural network states can be further enhanced with a pre-training method. Here, in this work, we integrate the DMRG method with our NNQS model to incorporate the physics of valid wave functions into the neural networks, thereby ensuring both high expressiveness and reliable convergence, while maintaining computational efficiency, and our method achieved good performance for molecular systems with up to 38 qubits.

2. Related Work

Neural network quantum states (NNQS) represent a groundbreaking intersection of machine learning and quantum physics, offering novel approaches to modeling and understanding complex quantum systems. This field has evolved rapidly, marked by several pivotal studies.

Carleo and Troyer introduced this concept in their groundbreaking work in 2017 [8]. They utilized a neural network known as a restricted Boltzmann machine (RBM) to represent the wave function of quantum many-body systems and achieved significant success in solving problems related to many-body localization and quantum spin systems. This work demonstrated that NNQS can effectively capture the complexity of quantum systems, while also providing a new perspective for understanding quantum many-body systems. Nomura, Y. et al. [21] highlighted advancements in using RBMs to solve strongly correlated quantum systems, a significant challenge in quantum physics. The research of Huang, L. and Wang, L. utilized RBMs to accelerate quantum Monte Carlo simulations, a critical aspect of quantum computations [22]. Another study conducted in 2017 [23] by Deng and colleagues explored how to enhance the accuracy of wave function representations in quantum many-body systems using deep neural networks. They discovered that deep learning can effectively capture quantum entanglement and interactions, which are crucial for understanding phenomena such as quantum phase transitions and quantum entanglement. In 2018, Cai, Z. and Liu, J. provided insights into the potential of NNQS for approximating wave functions for quantum many-body systems [24]. Choo, K. et al. expanded NNQS applications to include many-body excitations and symmetry considerations, integral for understanding quantum phenomena [25].

The largest CI calculation has been carried out for a molecule with about 24 electrons (24 orbitals) [26], since ${\mathrm{C}}_2{\mathrm{H}}_4\mathrm{O}$ contains 24 electrons (qubits = 38), it is not possible to calculate this with the full CI method. For scalability, the scaling of our method is $\mathrm{O}\left({\mathrm{N}}_{\mathrm{u}}{\mathrm{N}}_{\mathrm{o}}^4\right)$, where ${\mathrm{N}}_{\mathrm{u}}$ refers to the number of unique samples and ${\mathrm{N}}_{\mathrm{o}}^4$ accounts for the computation involved in assessing the local energy. For huge molecules such as proteins, a multiscale method [27] can be adapted, in which the protein systems are first divided into a molecule and a molecular environment, whereby the environment can be treated using a force field method or even DFT method, the central region is divided into many fragments, and the energy of each fragment is described using the method in this work. We employ DMRG for pre-training NNQS, which can accelerate the convergence rate of NNQS. For systems with more than 30 qubits, NNQS can rapidly provide reference energies using this method. Compared to previous works, our approach offers significant advantages in terms of operational speed and system size.

3. Methods

3.1. Quantum Chemistry Hamiltonians

The main task in ab initio quantum chemistry calculation is to solve the static Schrödinger equation $H|\Psi \rangle =E|\Psi \rangle$ to obtain the ground state $|\Psi \rangle$ and the ground-state energy E of the many-body interacting Hamiltonian:

$$\widehat{H}=-\sum _{i=1}^N\frac{1}{2}{\nabla }_i^2-\sum _{i=1}^N\sum _{A=1}^M\frac{Z_A}{|{\mathbf{r}}_i-{\mathbf{R}}_A|}+\sum _{i=1}^N\sum _{j>i}^N\frac{1}{|{\mathbf{r}}_i-{\mathbf{r}}_j|}$$

(1)

In this context, N represents the aggregate count of electrons, M stands for the total number of nuclei present. The term ${\nabla }_i$ refers to the kinetic operator for a single particle, specifically for the i-th electron. The notation ${\mathbf{r}}_i$ is used to denote the coordinates of electrons. ${\mathbf{R}}_A$ and $Z_A$ are used to signify the respective coordinates and charges of the A-th nucleus within the molecular structure. As shown in Ref. [20], we transform the Hamiltonian in Equation (1) into a many-spin Hamiltonian $H={\sum }_{i=1}^{N_h}{c}_i{P}_i$, where each $P_i$ represents a Pauli string, which is essentially the tensor product of the Pauli spin operators $\{I,X,Y,Z\}$ over N elements, paired with a real coefficient $c_i$. The total count of such Pauli strings is denoted by $N_h$. In quantum chemistry Hamiltonians, the scaling of $N_h$ is typically proportional to $O\left(N^4\right)$. This implies that for a given input bitstring x, there could be a polynomial number of ${\mathbf{x}}^{\prime }$ bitstrings, specifically on the order of $O\left(N^4\right)$, that have a non-zero interaction $H_{{\mathrm{xx}}^{\prime }}$ with it. Consequently, for larger values of N, the computation of local energy becomes increasingly resource-intensive, and storing every ${\mathbf{x}}^{\prime }$ that significantly interacts with x demands substantial memory. To address these challenges, we have developed an effective method for computing local energy. This approach uses a compactly structured Hamiltonian and a combined process for both assessing non-zero Hamiltonian entries and performing local energy computation [19].

3.2. NNQS-RNN: The NNQS Method Based on RNN Architecture

Using an RNN neural network as a representation for the ground state, expressed as $|{\psi }_{\overrightarrow{\theta }}\rangle$, where $\overrightarrow{\theta }$ signifies the parameters undergoing optimization, the system’s energy can be articulated as a function dependent on $\overrightarrow{\theta }$. In the second quantized formalism, with a basis set (single-electron quantum states or spin-orbitals) introduced, the many-electron wave function can be written as a linear combination of configurations

$$|{\psi }_{\overrightarrow{\theta }}\rangle =\sum _{\mathbf{x}}\langle \mathbf{x}|{\psi }_{\overrightarrow{\theta }}\rangle |\mathbf{x}\rangle =\sum _{\mathbf{x}}{\psi }_{\overrightarrow{\theta }}\left(\mathbf{x}\right)|\mathbf{x}\rangle$$

(2)

where each configuration is represented by an occupation number vector (‘configuration string’) $|\mathbf{x}\rangle =\{x_1,x_2,\dots x_N\}$ with $x_i\in \{0,1\}$ denoting whether the i-th spin orbital is occupied or not, then we have [8]

$$\begin{array}{cc}\hfill E(\overrightarrow{\theta })& =\frac{\langle {\psi }_{\overrightarrow{\theta }}|H|{\psi }_{\overrightarrow{\theta }}\rangle }{\langle {\psi }_{\overrightarrow{\theta }}|{\psi }_{\overrightarrow{\theta }}\rangle }=\frac{{\sum }_{\mathbf{x},{\mathbf{x}}^{\prime }}\langle {\psi }_{\overrightarrow{\theta }}|\mathbf{x}\rangle \langle \mathbf{x}|H|{\mathbf{x}}^{\prime }\rangle \langle {\mathbf{x}}^{\prime }|{\psi }_{\overrightarrow{\theta }}\rangle }{{\sum }_{\mathbf{y}}\langle {\psi }_{\overrightarrow{\theta }}|\mathbf{y}\rangle \langle \mathbf{y}|{\psi }_{\overrightarrow{\theta }}\rangle }\hfill \\ \hfill & =\frac{{\sum }_{\mathbf{x}}{E}_{loc}\left(\mathbf{x}\right)p_{\overrightarrow{\theta }}\left(\mathbf{x}\right)}{{\sum }_{\mathbf{y}}{p}_{\overrightarrow{\theta }}\left(\mathbf{y}\right)}={\mathbb{E}}_p\left[E_{loc}\left(\mathbf{x}\right)\right]\hfill \end{array}$$

(3)

Here, x, ${\mathbf{x}}^{\prime }$, and y represent distinct bitstrings. In the second line of Equation (3), the concept of local energy, denoted as $E_{loc}\left(\mathbf{x}\right)$, is established as

$$\begin{array}{c}\hfill E_{loc}\left(\mathbf{x}\right)=\sum _{{\mathbf{x}}^{\prime }}{H}_{\mathbf{x}{\mathbf{x}}^{\prime }}{\psi }_{\overrightarrow{\theta }}\left({\mathbf{x}}^{\prime }\right)/{\psi }_{\overrightarrow{\theta }}\left(\mathbf{x}\right)\end{array}$$

(4)

In this context, $H_{\mathbf{x}{\mathbf{x}}^{\prime }}=\langle \mathbf{x}|H|{\mathbf{x}}^{\prime }\rangle$ signifies the matrix element, and ${\psi }_{\overrightarrow{\theta }}\left(\mathbf{x}\right)=\langle \mathbf{x}|{\psi }_{\overrightarrow{\theta }}\rangle$ represents the probability amplitude for the wave function hypothesis $|{\psi }_{\overrightarrow{\theta }}\rangle$ in the $|\mathbf{x}\rangle$ basis. Additionally, the expression $p_{\overrightarrow{\theta }}\left(\mathbf{x}\right)={|{\psi }_{\overrightarrow{\theta }}\left(\mathbf{x}\right)|}^2$ is used to denote the probability.

Accurately calculating Equation (3) is typically unfeasible, due to the vast, exponentially large set of varying bitstrings. Nevertheless, an approximate assessment of Equation (3) can be achieved by drawing samples from the probability distribution $p_{\overrightarrow{\theta }}\left(\mathbf{x}\right)$, thereby acquiring a collection of $N_s$ samples, labeled as $\{{\mathbf{x}}^1,{\mathbf{x}}^2,\cdots ,{\mathbf{x}}^{N_s}\}$, followed by their subsequent averaging:

$$\begin{array}{c}\hfill \tilde{E}\left(\overrightarrow{\theta }\right)=\frac{1}{N_s}\sum _{i=1}^{N_s}{E}_{loc}\left({\mathbf{x}}^i\right)\end{array}$$

(5)

In this instance, $\tilde{E}\left(\overrightarrow{\theta }\right)$ is employed rather than $E\left(\overrightarrow{\theta }\right)$, to underscore that it is merely an estimation of the actual value. Consequently, if one can effectively draw samples from $p_{\overrightarrow{\theta }}\left(\mathbf{x}\right)$ (feasible when ${\psi }_{\overrightarrow{\theta }}\left(\mathbf{x}\right)$ can be efficiently calculated for each x) and identify those significant $H_{\mathbf{x}{\mathbf{x}}^{\prime }}$ values along with corresponding ${\mathbf{x}}^{\prime }$, then a more efficient evaluation of Equation (3) is achievable. Employing a gradient-based optimizer can expedite these calculations compared to methods that do not use gradients. Using the gathered samples, it is possible to estimate the gradient of Equation (3) through automatic differentiation, as suggested in [17]:

$$\begin{array}{c}\hfill {\nabla }_{\overrightarrow{\theta }}\tilde{E}=2\left({\mathbb{E}}_p\left[\left(E_{loc}\left(\mathbf{x}\right)-{\mathbb{E}}_p\left[E_{loc}\left(\mathbf{x}\right)\right]\right){\nabla }_{\overrightarrow{\theta }}ln\left({\Psi }_{\overrightarrow{\theta }}^{*}\left(\mathbf{x}\right)\right)\right]\right)\end{array}$$

(6)

In a comparable manner, ${\nabla }_{\overrightarrow{\theta }}\tilde{E}$ serves as an approximation for the precise gradient ${\nabla }_{\overrightarrow{\theta }}E$. Subsequently to this, the parameters $\overrightarrow{\theta }$ are updated using ${\nabla }_{\overrightarrow{\theta }}\tilde{E}$ in conjunction with the optimizer, culminating in the completion of a single cycle of the variational Monte Carlo algorithm(VMC).

3.3. DMRG Method

Training the NNQS model with completely random initial parameters can suffer from optimization problems, similarly to a variational quantum algorithm. Refs. [28,29] Specifically, a proper initial state will significantly influence the convergence and accuracy of NNQS. In the context of chemical applications, it can be assumed that a chemically-motivated pre-training would be beneficial. In conventional quantum chemistry approaches, the initial step involves simplifying the electronic wave function to a single Slater determinant, for ease of computation. Subsequently, to capture electron correlation, additional determinants are incorporated into the analysis. Commonly used methods such as configuration interaction (CI) and coupled-clusters (CC) theory can generate a approximated ground-state wave function; however, its high-order complexity is excessively expensive for serving as an initial guess.

Recently, the DMRG algorithm, which implements a matrix product state (MPS) wave function, has been widely applied in quantum chemistry calculations. Refs. [30,31,32,33,34] expressed the electronic wave function in a CI-expansion form:

$$\left|\Psi \right.〉=\sum _{i_1{i}_2\dots i_N}{c}_{i_1{i}_2\dots i_N}|i_1{i}_2\dots i_N\rangle ,$$

(7)

where N is the number of spin orbitals, $|i_1{i}_2\dots i_N\rangle$ is the computational basis, and the MPS ansatz factorizes this rank-N coefficient tensor $\left\{c_{i_1{i}_2\dots i_N}\right\}$ into lower-rank tensors ${\{}^{k}T\}$, which can be written as

$$c_{i_1{i}_2\dots i_N}=\sum _{u_0\dots u_N}{}^{1}T_{u_0{u}_1}^{i_1}{}^{2}T_{u_1{u}_2}^{i_2}\dots {}^{N}T_{u_{N-1}{u}_N}^{i_N}$$

(8)

where ${\{}^k{T}_{u_{k-1}{u}_k}^{i_k}\}$ represents elements of the rank-3 tensor ${}^{k}T$ at the k-th site (orbital), with $i_k$ called the physical index and $u_k$ called the auxiliary index. The singular value decomposition (SVD) procedure could be a good choice to generate these MPSs.In addition, the MPSs can be truncated according to the singular value, so that the maximum size of the auxiliary indices, which is defined as the bond dimension of the MPS, denoted as $M={max}_{0\le k\le N}\left\{u_k\right\}$, is constrained. This tensor decomposition procedure can also be applied to the operators to generate the matrix product operator (MPO) and develop a tensor network formula, as presented in Figure 1. The ground-state energy $E_0$ is the lowest energy of the system, defined as

$$E_0=\underset{|\Psi \rangle }{min}\frac{\langle \Psi |\widehat{\mathbf{H}}|\Psi \rangle }{\langle \Psi |\Psi \rangle }$$

(9)

The Lagrangian multiplier method is applied to evaluate this ground-state energy, which derives a eigenvalue problem:

$$\langle \Psi |\widehat{\mathbf{H}}|\Psi \rangle =E\langle \Psi |\Psi \rangle$$

(10)

It is difficult to solve this problem directly, so the DMRG adopts a divide-and-conquer strategy, which is to adjust the MPS tensor one at a time, forward and backward, until the energy converges. This is referred to as the one-site DMRG method. In quantum chemistry, to avoid the local minimum, the two-site DMRG method is used, which entails adjusting two adjacent MPSs simultaneously. On each pair of sites, the corresponding MPSs are updated according to the effective operator ${\widehat{\mathbf{H}}}_{\mathrm{eff}}^i$, which is developed by contracting the current MPOs and the left and right part of the tensor network. Thus, the new MPSs are derived through this eigenvalue problem:

$${\widehat{\mathbf{H}}}_{\mathrm{eff}}^i|{\psi }_i\rangle =\lambda |{\psi }_i\rangle$$

(11)

This formula is resolved using a method of iterative diagonalization, employing techniques like the Lanczos, Davidson, or Jacobi–Davidson methods.

The DMRG algorithm performs variational optimization for the MPSs to search for an approximation of the ground state. An important feature of DMRG is its linear scaling with respect to the system size (number of sites, or orbitals) if the bond dimension is fixed, which makes it particularly appealing as a lightweight pre-training algorithm for chemical significance for NNQS. Although, the MPSs generated from the DMRG procedure can be used directly as the initial guess to train the NNQS, a more efficient method is to generate a CASCI-type wavefunction that contains the most important configurations. There are several methods sued to accomplish this task through Monte-Carlo-based sampling. In this work, we adopt the entanglement-driving genetic algorithm (EDGA) to perform the wave function transformation.

Here, the EDGA method can develop the significant terms of Equation (7) from the pretrained MPS as a initial guess for the subsequent NNQS optimization. The EDGA method treats the Slater determinate (SD) as a pair of DNA sequences consisting of the alpha spin (spin-up) configurations and beta spin (spin-down) configurations, respectively. Based on the genetic algorithm, these configurations will evolve based on two steps, “crossover” and “mutation”, as illustrated in Figure 2.

The EDGA method introduces quantum entanglement into these two steps, to make the evolution more reasonable and efficient. In the crossover step, two SDs are randomly selected, and the alpha spin configuration of the first SD and the beta spin configuration of the second SD are combined, to generate a new SD. The possibility of accepting this new SD is generated using

$$\rho =\frac{{\tilde{c}}_i}{{\sum }_i^N{\tilde{c}}_i}$$

(12)

where ${\tilde{c}}_i$ is determined by a segmentation function of the CI coefficient $c_i$:

$${\tilde{c}}_i=\left\{\begin{array}{cc}sin(|c_i|),\hfill & \hfill \mathrm{for}\mathrm{the}\mathrm{first}15\%\mathrm{evolution}\mathrm{steps};\\ |c_i|,\hfill & \hfill \mathrm{for}\mathrm{the}\mathrm{middle}40\%\mathrm{evolution}\mathrm{steps};\\ c_i^2,\hfill & \hfill \mathrm{for}\mathrm{the}\mathrm{rest}\mathrm{evolution}\mathrm{steps};\end{array}\right.$$

(13)

In the mutation step, the determinants obtained from the “crossover” step are randomly changed. The occupation status of orbitals i and j are exchanged according to the probability defined as

$${\rho }_{ij}=\frac{I_{ij}}{{\sum }_i^N{I}_{ij}}$$

(14)

where $I_{ij}$ is the mutual information of the corresponding orbital pair. This EDGA procedure can be easily conducted using the Kylin quantum chemistry package [35], to develop a satisfying initial guess.

3.4. Pre-Training NNQS with DMRG Method

As shown in Figure 3, first, we transfer the data generated by DMRG into many {state, psi} pairs and load these data into our NNQS model. The model extracts the state part from the DMRG data, serving as a supervised learning dataset for training. Initially, the model starts with random parameters and computes the psi for each state. Then, it uses the absolute error between this psi value and the corresponding DMRG state’s psi value as the loss, undergoing backward and parameter updates to reduce the loss. Through multiple iterations, the model’s parameters gradually stabilize, and the psi values it trains for each state closely approximate those of the corresponding DMRG states. At this point, we consider the model’s wave function fit as having a high similarity with the DMRG-fitted wave function, achieving comparable computational results to DMRG. Then, our NNQS generates numerous checkpoint files, each associated with specific steps and loss values. The step value records at which step of our NNQS pre-training the checkpoint was created, while the loss value records the loss at the time of checkpoint generation in comparison to the DMRG pre-training data.

During the gradient update process, we use weight decay to prevent overfitting of the model. This is essentially a regularization method that applies additional constraints to the weight parameters in the network. Weight decay is usually implemented by adding a term proportional to the weight magnitude to the loss function. When using weight decay, the loss function includes not only the original loss (like cross-entropy loss or mean squared error loss) but also an additional term proportional to the square of the weights, usually a coefficient multiplied by the sum of the squares of all weights. This coefficient is often referred to as the weight decay coefficient or regularization parameter. In the construction of the loss function, weight decay serves as a factor applied to the regularization term. This term typically reflects the complexity of the model; thus, the weight decay functions to modulate how much the model’s complexity influences the overall loss function. We use the gradient descent optimizer AdamW [36], which is an optimization algorithm used in deep learning and that is a variant of the Adam [37] optimization algorithm. AdamW introduces direct control over weight decay, rather than indirectly implementing it through L2 regularization, as in the traditional Adam algorithm. This modification makes AdamW more effective than the standard Adam in dealing with certain types of problems, especially those requiring a substantial weight decay. During model training, the learning rate schedule we use is $lr=d_{model}^{-0.5}\times min(i^{-0.5},i\times S_{warmup}^{-1.5})$, where $lr$ is the learning rate of the i-th training epoch, and we set the warm up steps as $S_{warmup}=4000$ [38].

To ascertain whether pre-training with DMRG can expedite the convergence speed of the energy inference system of our NNQS, we initiated a test without pre-training, operating within the same system. We identified a set of NNQS model parameters that showcased superior convergence results. These parameters included n_hidden, nip, and layer based on the RNN model. Leveraging this combination of parameters, we conducted DMRG pre-training experiments. Our tests on small molecular systems (qubits < 30) such as ${\mathrm{H}}_2\mathrm{O}$ (14 qubits) and strongly correlated systems like ${\mathrm{H}}_{10}$ STO-6G (20 qubits) indicated that DMRG pre-training accelerated the convergence speed of the our NNQS model. Upon evaluating larger molecular systems (qubits > 30), we found that the convergence speed of the NNQS fluctuated based on the different loss values of the pre-training checkpoints. The experimental outcomes revealed that utilizing checkpoints with reduced loss values for continued computations led to a swifter convergence. Moreover, after DMRG pre-training, our NNQS was capable of achieving a lower energy for certain larger molecular systems (qubits > 30), like ${\mathrm{C}}_2{\mathrm{H}}_4\mathrm{O}$ (38 quibts), a feat unattainable when relying solely on direct NNQS inference without any pre-training.

4. Results

4.1. The Influence of Different Hyper-Parameters on NNQS

After a neural network model is built, it is typically fine-tuned through hyper-parameter tuning to enhance its expressive capability. In our study, we introduce the NNQS model, which possesses the following hyper-parameters: n_hidden, indicating the number of neurons in each hidden layer of the NNQS network; nip, representing the number of features in the input data; and in the context of NNQS, this denotes the representation of each qubit using a nip-dimensional vector; layer, signifying the number of layers in the NNQS network. The network-type parameter is employed to select between the gated recurrent unit (GRU) [39] and long short-term memory (LSTM) [40] is used for training. LSTM, a specialized form of recurrent neural network(RNN), is designed to address the long-term dependency issues encountered by a standard RNN in processing long sequences. Conversely, GRU, a variant of LSTM, combines the forget and input gates of LSTM into a single “update gate”, and it also merges the cell state with the hidden state, alongside other modifications. Due to its relatively simpler structure, GRU offers a higher computational efficiency. When testing the NNQS, we found that for molecular systems with more than 30 orbitals (qubits), the NNQS converged slowly and had difficulty converging to chemical precision (CP = 0.0016). Our goal was to choose a set of parameters that enabled better convergence for most systems. On an Nvidia A100 GPU with 40G of GPU memory, we meticulously tuned hyper-parameters such as n_hidden, nip, layer, and network-type. In addition, we used Pytorch 1.13.1 + cu116, Julia 1.8.0 to build our NNQS model.

Initially, we chose a small molecule system, LiH (qubit = 12), with less than 20 qubits for the hyper-parameter testing of our NNQS model. We selected LiH as the subject for our hyper-parameter testing due to its rapid computational efficiency, limited number of configurations, and suitability for streamlined parameter optimization and debugging processes. In Figure 4, the x-axis represents the iteration count of the NNQS model, while the y-axis shows the absolute error in energy at each step for LiH compared to FCI. When the absolute error diverged from the chemical precision (CP = 0.0016, denoted by the red dashed line in Figure 4), we considered the model to have converged.

Figure 4a demonstrates that the convergence performance for n_hidden = 64 was similar to that for n_hidden = 32, but the computational time per step was longer than n_hidden = 32. When n_hidden = 16, the model exhibited inferior convergence performance (indicated by the blue solid line in Figure 4a). In the second set of tests, it is shown that for nip = 16, 32, 64, the convergence curves are similar, but nip = 16 required a shorter computation time per step. Figure 4c reveals that at layer = 6, the model achieved chemical precision earlier (below the red dashed line), indicating a better convergence performance. Figure 4d illustrates that for the NNQS model, employing GRU resulted in a superior fit compared to LSTM. In summary, considering both the convergence performance and operational speed, we used n_hidden = 32, nip = 16, layer = 6, and GRU for subsequent testing

We compare our test results with the NAQS [17], RBM, and FCI in

Table 1. Comparison of the energy (in Hartree) between this work, NAQS, RBM, and FCI.

Molecular Systems	This Work	NAQS	RBM	FCI
$\mathrm{LiH}$	−7.7845	−7.7845	−7.7777	−7.7845
${\mathrm{H}}_2\mathrm{O}$	−75.0155	−75.0155	−74.9493	−75.0155
${\mathrm{N}}_2$	−107.6599	−107.6595	−107.5440	−107.6602
${\mathrm{CH}}_4$	−39.8062	−39.8062	−39.7571	−39.8063
${\mathrm{C}}_2$	−74.6904	−74.6899	−74.5147	−74.6908
$\mathrm{LiF}$	−105.1661	−105.1662	−105.1414	−105.1662
${\mathrm{PH}}_3$	−338.6979	−338.6984	−338.6472	−338.6984
${\mathrm{Li}}_2\mathrm{O}$	−87.8916	−87.8909	−87.3660	−87.8927

. The results show that our tested approach achieved a better convergence in certain systems and the data with a dark underline indicate the best values.

4.2. The Influence of DMRG Pre-Training on NNQS

In this section, we demonstrate the use of DMRG configuration coefficients as pre-training data for the NNQS. This approach enables supervised learning for the NNQS model, aiming to closely approximate the wave function predicted by DMRG before proceeding to inference. The checkpoints generated at different steps during the pre-training process have varying impacts on the model. A lower number of pre-training steps might result in suboptimal model fitting, leading to slower convergence and a higher initial inference energy during inference. Conversely, an excessive number of pre-training steps may lead to overfitting, where the model accurately fits the DMRG pre-training data’s wave function but may struggle to generate new samples or may exhibit stagnant energy levels during the subsequent inference. In such cases, we describe the NNQS as converging to a local minimum. We use the configuration coefficients generated using DMRG data to pre-train the NNQS. At first, we verify the convergence effect of this method with two smaller systems ${\mathrm{H}}_2\mathrm{O}$ (14 qubits) and ${\mathrm{H}}_{10}$ STO-6G(20 qubits), as shown in Figure 5.

We selected the ${\mathrm{H}}_2\mathrm{O}$ and ${\mathrm{H}}_{10}$ STO-6G systems to test the effectiveness of pre-training. Figure 5a,b represent the convergence of the ${\mathrm{H}}_{10}$ STO-6G system, while (c,d) depict the convergence of ${\mathrm{H}}_2\mathrm{O}$. The y-axis in Figure 5a,c shows the absolute error between the energy values inferred by the NNQS and those of the FCI, with the red dashed line representing the chemical precision. When the curve converged below this red dashed line, we considered the model to have inferred the energy of the molecular system with an accuracy acceptable by FCI standards. The y-axis in Figure 5b,d displays the actual energy values inferred by our NNQS model. We compare the energy values at the 1000-th (purple solid line), 2000-th (yellow solid line), and 3000-th (green solid line) pre-training steps checkpoints with those obtained from inference without pre-training (blue solid line). The experiment showed that for ${\mathrm{H}}_{10}$ STO-6G, the energy in the first step of pre-training was lower than that of starting from scratch, as shown in

Table 2. Comparison of the first step energy between ${\mathrm{H}}_{10}$ STO-6G and ${\mathrm{H}}_2\mathrm{O}$ with different pre-training steps.

Molecular Systems	${\mathbf{H}}_{10}$ STO-6G	${\mathbf{H}}_2\mathbf{O}$
no ckpt	5.104793	−69.437209
1000-th step ckpt	−2.485850	−73.123730
2000-th step ckpt	−2.384864	−73.343999
3000-th step ckpt	−2.202542	−73.386116

. However, for the strongly correlated ${\mathrm{H}}_{10}$ STO-6G system, the opposite was observed; the first step energy from scratch was erroneous, and the energy became accurate only after pre-training. For the 1000-th step checkpoints, the ${\mathrm{H}}_{10}$ STO-6G system exhibited an abnormal convergence curve (purple solid line in Figure 5a,b), where the energy after pre-training was higher than that without pre-training (blue solid line), indicating overfitting of the model. Both ${\mathrm{H}}_{10}$ STO-6G and ${\mathrm{H}}_2\mathrm{O}$ converged more quickly to chemical precision (below the red dashed line) after pre-training, demonstrating that using DMRG as a pre-training method for NNQS can accelerate its convergence.

As shown in

Table 3. Comparison of the elapse time between ${\mathrm{H}}_{10}$ STO-6G and ${\mathrm{H}}_2\mathrm{O}$ with different pre-training steps.

Molecular Systems	${\mathbf{H}}_{10}$ STO-6G	${\mathbf{H}}_2\mathbf{O}$
no ckpt	103.7935	1524.3876
1000-th step ckpt	65.6547	1580.0200
2000-th step ckpt	76.3441	944.5145
3000-th step ckpt	113.5518	1032.1916

and

Table 4. Comparison of the convergence step between ${\mathrm{H}}_{10}$ STO-6G and ${\mathrm{H}}_2\mathrm{O}$ with different pre-training Steps.

Molecular Systems	${\mathbf{H}}_{10}$ STO-6G	${\mathbf{H}}_2\mathbf{O}$
no ckpt	2732	7376
1000-th step ckpt	890	6925
2000-th step ckpt	347	2871
3000-th step ckpt	507	2358

, we compiled the number of iterations required for ${\mathrm{H}}_{10}$ STO-6G and ${\mathrm{H}}_2\mathrm{O}$ to converge to chemical precision. For the unpretrained tests, we calculated the total runtime from scratch inference to convergence to chemical precision as the convergence time. For the pretrained tests, we recorded the time from the beginning of pretraining to the generation of the corresponding step’s checkpoint, and the time from loading the checkpoint to convergence to chemical precision as the convergence time. The results indicate that the convergence time and the number of iterations for the NNQS after pretraining were generally less than those for the unpretrained versions.

As depicted in Figure 6, DMRG data were employed for the pre-training of the NNQS. This pre-training notably enhanced the convergence trajectory for ${\mathrm{C}}_2{\mathrm{H}}_4\mathrm{O}$, as evidenced by a marked decline. The convergence curve without pre-training is illustrated by the blue solid line in Figure 6. Evaluations were conducted utilizing DMRG checkpoints at the 20,000-th, 40,000-th, and 80,000-th steps. As delineated in Figure 6b , the unpretrained curve exhibits a discrepancy exceeding 0.02 relative to DMRG at the 29,000-th step (as indicated by the blue solid line). Furthermore, the energy extrapolated from the 20,000-th step checkpoint surpasses that derived from the 40,000-th and 80,000-th iteration checkpoints. The convergence curves of energy for the 40,000-th and 80,000-th iterations demonstrate similarity, suggesting a negligible variation in model parameters between these iterations in the pre-training phase. Figure 6a reveals that the initial energy inferred without pre-training (−134.756129) significantly exceeded the initial energies with pre-training (20,000-th checkpoint: −149.3210498, 40,000-th checkpoint: −149.337339, 80,000-th checkpoint: −149.344904).

4.3. The Influence of the Optimizer on NNQS

Numerous experiments have been conducted to investigate the effects of different optimizers on the calculation of wave functions using neural networks. Traditional first-order optimization algorithms, commonly used in these calculations, such as Adam [37], and AdamW [36], depend on the loss function’s first-order derivatives to minimize the loss. In contrast, second-order optimization algorithms, including Newton’s method and the LBFGS [41] optimizer, incorporate a Hessian matrix of the loss function. However, the computational challenges associated with computing and storing the Hessian matrix have impeded their widespread use.

At present, the Kronecker-factored approximate curvature (KFAC) optimizer has proven its efficiency in accelerating neural network convergence by leveraging approximations of the Hessian matrix, without imposing substantial computational overheads [42]. This method has been well tested and has shown great results in quantum chemistry computations [43,44]. However, when interfacing with certain complex network layers in PyTorch, the KFAC implementation requires specific layer formatting, which can pose challenges in its applicability. This has driven us to explore simpler second-order optimization algorithms that better fit with these network designs and can improve training speed.

We chose AdamW as the optimizer for NNQS for the following reasons: (1) it employs a more intuitive weight decay strategy compared to the standard Adam optimizer, making it easier to fine-tune the weight decay hyper-parameters. This aids in mitigating the risk of overfitting during model training and enhances generalization performance; (2) it also typically exhibits a more stable convergence performance during training, particularly in deep neural networks. This contributes to reducing the training instability and overall training time; (3) the Adam optimizer may introduce biases in parameter updates during the initial stages of training, leading to training instability. AdamW, through its improved weight decay method, mitigates such biases, thereby accelerating model convergence.

Weight decay is a regularization technique employed in deep learning models to mitigate the risk of overfitting. Its function entails the addition of an extra penalty term to the loss function, penalizing the model’s weight parameters, thereby encouraging them to tend towards smaller values. This aids in preventing the model from overfitting the training data during the training process, ultimately enhancing its generalization performance. In our study, we selected four quantum systems with qubits less than or equal to 30 and investigated the impact of different weight decay values on the performance of the AdamW optimizer within the NNQS-model. As shown in Figure 7, when the weight decay was set to 0.003, all four molecular systems exhibited favorable convergence curves.

4.4. Application: Ferrocene

For the application, we studied the transition barrier energy between two structures of ferrocene (Fe(C₅H₅)₂), as shown in Figure 8. Ferrocene is an important organometallic compound, and it has stimulated an explosion of interest in compounds of d-block metals with hydrocarbons. The 1973 Chemistry Nobel Prize was bestowed upon Ernst Otto Fischer and Geoffrey Wilkinson “for their pioneering work, performed independently, on the chemistry of the organometallic, so called sandwich compounds”. The “sandwich” structure of ferrocene consists of two cyclopentadienyl (Cp) rings bound to a central iron atom. There are two kinds of structures with a small energy difference, called eclipsed ferrocene with D5h symmetry and staggered conformer with D5d symmetry, and the staggered structure has a lower energy than the eclipsed structure [45]. The ground state of ferrocene is dominated by a single configurations, including the coupling of Fe and C orbitals [46]. We used cc-pVTZ-DK basis set and used the automated valence active space (AVAS) [47] method to generate the desired active orbitals by projecting the canonical molecule into the Fe d and C p_z atomic orbitals (30 spin-orbitals). Using our method, we obtained the energy barrier between the two structures as $9.16\times {10}^{-3}$ Hartree, while the best estimates from the CASCI method predicted $9.30\times {10}^{-3}$ Hartree. We can see that chemical accuracy (1 kcal/mol = 0.0016 Hatree) was achieved with our method. These results gave us confidence in our method’s ability to deal with complex system with transition metals, and even strongly correlated materials.

We compiled the NNQS in hybrid single-precision and double-precision on a new generation Sunway supercomputer, Julia 1.6.5 and Pytorch 1.13 was used, which had been integrated with the underlying runtime and computing libraries; for example, the MPI communication library (3.2b2) customized for the interconnection network topology of the new-generation Sunway supercomputer, deeply optimized BLAS and LAPACK libraries (3.8), etc. We also used an AMD EPYC 7742 CPU and NVIDIA A100 PCIe 80GB for the CPU and GPU computation environments, respectively.

5. Conclusions

This study represents a significant advancement in the field of quantum many-body physics, addressing the critical challenge of the complexity of equations in physics and chemistry. By integrating the density-matrix renormalization group (DMRG) method with transformer-based neural networks, we introduced a novel approach that efficiently tackles the computational difficulties in solving complex quantum systems, especially those with strong correlations. Future work will focus on further refining these methods, extending their applicability to even more complex systems, and exploring the integration of emerging machine learning and quantum computing techniques, to continue advancing our understanding of the quantum world.