Deep learning representations for quantum many-body systems on heterogeneous hardware

To meet the strictly demanding requirement, we develop two deep CNN structures for the spin models and fermion models, which are denoted as CNN1 and CNN2 respectively. The computational complexity of CNN is much lower compared to the fully connected structure. By taking advantage of translational invariance, the CNN can scale to very large lattices.

The structure of CNN1 is depicted in figures 2(a) and (b), and the network is built by stacking the building block denoted in figure 2(a) six times [36]. A building block consists of two-dimensional convolution and one-dimensional maxpooling and transposed-convolution. To maintain the dimensions, paddings are employed on the input spin lattice and between two building blocks. The padding scheme is based on the PBCs. The final outputs are the products of the neurons, which are the wavefunction coefficients for the input spin configurations.

Zoom In Zoom Out Reset image size

The structure of CNN2 is depicted in figures 2(c) and (d). CNN2 originates from CNN1, with two major modifications: (1) the input spin lattice is processed by an embedding layer; (2) the one-dimensional maxpooling is performed with stride one. To maintain the dimensions, padding is performed by copying the first several neurons to the last. The deep structure is built by stacking the building block in figure 2(c) six times. Finally, with a convolution operation, the channel number is decreased to one, the final output is the product of the neurons.

The deep CNN structures used in this work have several important differences compared to other CNN structures for quantum many-body problems [21, 24, 28]. First, the nonlinearity in our deep CNN is induced by the maxpooling, instead of traditional activation functions. The maxpooling picks up the most important degree of freedom in a convolution filter, which is similar to the coarse-grained process in a renormalization group theory [37]. Second, the wavefunction coefficients are generated by the products of neurons, thus the CNN can give the ±1 signs within real network parameters, which differs from the exponential function in RBM based structures.

The CNN1 structure in this work is based on the deep CNN previously reported in [26, 36], and the wavefunction is preconditioned by the Marshall-sign-rule (MSR)[38, 39]:

$$\begin{equation} \tilde{W}(S)=(-1)^{\mathrm{MA}}W_{\mathrm{CNN1}}(S), \end{equation} \tag{ 5 }$$

where $\mathrm{MA}$ is the magnetization on the equivalent sublattice A. Although the MSR is only exact for the case of $J2 = 0$, however, because the CNN1 is able to give the ±1 signs, it is able to correct the MSR for large $J2$.

Enforcing symmetries can significantly reduce the optimization difficulty [40]. In the case of $J2 = 0.5$, the final wavefunction coefficient has an enforced rotational symmetry on the square lattice [24]:

$$\begin{equation} W(S)=\sum_i\tilde{W}(\hat{T}^{\,i} S), \end{equation} \tag{ 6 }$$

where $\hat{T}$ is the spin rotation operator and $i = 0,1,2,3$.

Based on the structure of CNN1, we design CNN2 to solve the t-J model. To encode the Fermions into spins, Jordan–Wigner transformation is used. Figure 3 denotes the spin ordering for the JW transformation on $4\times 4$ square lattice. We have compared several spin orderings on the $8\times 8$ square lattice, and the spin ordering denoted in figure 3 gives the lowest variational energy.

Zoom In Zoom Out Reset image size

There are three degrees of freedom in each site for the t-J model, which correspond to no-occupation, one-occupation with spin up and spin down. The Hilbert space of the t-J model is much larger compared to spin models, thus a higher representation ability of the CNN is required. To enhance the representation ability, an embedding layer is employed to transfer the input of each degree of freedom to a trainable two-element vector. Even though the original Hamiltonian of the t-J model has rotational symmetry, the JW transformation breaks this symmetry, and we do not enforce any additional symmetry of the CNN wave-function, and the final wave-function coefficient is directly generated from the CNN2: $W(S) = W_{\mathrm{CNN2}}(S)$.

For the $J1$–$J2$ model, direct training of the model in large lattice starting from random parameters is very difficult. The convolution operation is intrinsically scalable for different lattice sizes. Furthermore, the energy convergence on a small lattice is fast, as the lattice size increases, the ground energy will converge with respect to the lattice size. Initially the network is trained on $6\times 6$ lattice from randomly initialized parameters. The network parameters are then checkpointed and used for optimizing the larger $10\times 10$ system as the pretrained model. With multiple stages of transfer learning, we finally increased the lattice size to $36\times 36$ and obtained a good initial energy. Thanks to transfer learning, the local optimizations are avoided and the training steps for large lattices are significantly reduced.

For the t-J model, a proper initial model parameter and the initial spin configuration are beneficial for avoiding local optimizations, and therefore crucial for fast energy convergence. Here we calculate the energy of a randomly initialized CNN by performing MCMC sampling, assuming that the CNN with lower initial energy will have a better convergence after SR optimizations. With a preselection stage, the CNN parameters together with the initial spin configuration that has the lowest energy result are selected. After the initial state selections, the optimization steps are significantly reduced especially on the lattice as large as $12\times 12$.

Different from the distributed data parallel training in conventional deep learning tasks, the proposed AI-HPC framework only holds the same model state, while it replaces training data with on-the-fly MC sampling. In principle, importance sampling in the Hilbert space is achieved by multiple independent Markov chains. Considering that the MC sampling tasks are intrinsically parallel, they can be distributed across all participating processes, where the parallel execution among the 4 processes is pictured in figure 4(a). Inside each process, there is an independent chain of MC sampling, which generates $E_{\mathrm {loc}}$ and $O_{\mathrm {loc}}$ with model forward and backward computations. In the figure, each process produces B samples, where R equals the number of network parameters. The input data is naturally distributed along with Markov chains. Implementing the process-level parallel sampling only requires broadcasting the initial parameter number and the parameter update.

The MCMC sampling is time sequential, and numerous state transitions exist between two collected samples. The state transition along the Markov chain requires little computation, which leads to significant resource under-utilization, especially for today’s powerful heterogeneous devices. In this case, an alternative by managing multiple independent chains in one process is proposed, as depicted in figure 4(b). In each step, one sampling step from many chains increases the batch size during the model computation. The CNN computation is further divided into multiple computing layers and then offloaded onto heterogeneous many-core processor with modest thread-level parallel optimization. The repeated CNN execution with varying lengths can be organized with the fixed-length batch to improve performance. Additionally, such method theoretically increases the sample independence, which may help faster model convergence comparing to sampling from a single chain.

The flexible parallel MCMC scheme (i.e. MPI-parallel among different processes and batch-parallel among independent chains) provides an alternative to navigate the aforementioned initial state selection in section 2.2.4. To adapt to the two-level MCMC sampling, the selection paradigm is described in figure 4(d). The parallel processes is designed to find suitable initial model parameters, while the chain batch inside each process aims to find the best initial spin lattice. Taking the advantage of HPC infrastructure to simultaneously examine millions of sampling results, it dramatically offsets the difficulty in finding a reasonable initial configuration in the exponentially large Hilbert space.

With sufficient samples collected from MCMC, the values for updating network parameters δ are obtained through SR. The procedure for computing δ includes three steps: (1) adjusting the data format; (2) constructing the covariance matrix; and (3) solving a system of linear equations.

After MCMC sampling, the input for the SR procedure is the evenly distributed among all participating processes. These samples are then organized in a 2D-mesh style. For P processes, the collected samples are divided into $\sqrt{P}$ blocks. Naturally, the processes are split into $\sqrt{P}$ groups, where each group contains $\sqrt{P}$ processes. The data block is scattered into corresponding processes within the group.

Figure 4(c) denotes an example of 4 processes (Proc-0 and Proc-1 in group-1, Proc-2 and Proc-3 in group-2). The block exchange between 2 processes can be illustrated by swapping blocks of ② and ③ in group 1 (⑥ and ⑥ in group 2). As depicted in the figure, for a CNN with the parameter number R, it generates the 2D-mesh distributed matrix [4 B, R]. The following covariance matrix [R, R] is constructed in parallel. The final δ is obtained by combining the matrix with the first-order gradients of the energy. Similar to data parallelism training, all processes shares the same δ by MPI broadcast.

The scalability of SR is limited on constructing and solving the covariance matrix. In previous investigations, small network parameter number achieves moderate energy precision [36]. In this work, the network parameter number is further increased to the magnitude of 10⁵. With the help of modern parallel programming paradigm and high performance heterogeneous devices, accurately constructing and storing the full covariance matrix is still feasible. However, there still exists scalability issue as the parameter number continues to increase. One solution is to adopt approximate solutions, such as the sparse distributed solver as introduced in NetKet [41] and hybrid low-rank natural gradient method in [42]. Investigating the performance of these highly scalable sparse algorithms is one of our future directions.

Figure 4(e) depicts the proposed AI-HPC software stack in solving CNQS. In this work, CNQS is built by a series of common deep learning operators, such as Conv, Deconv, Maxpool and Embedding.

The gradients of CNQS are obtainable through the MCMC. As introduced in section 2.1, the forward and backward calculations of the CNN give the local energy and the local derivative during the MCMC. With the gradients, it is straightforward to use optimizers, such as SGD, Adam, etc. The forward and backward computation of CNN, as well as the optimizers are supported by mainstream deep learning libraries, such as TensorFlow and PyTorch, etc. The SR method is handled by the MPI and ScaLAPACK, which further supports heterogeneous devices through the MAGMA library.

The computation of CNN can be significantly accelerated by heterogeneous hardware, such as Intel Xeon Phi, NVIDIA (AMD) GPUs, sw26010pro CPUs, etc. The sw26010pro many-core processor is the upgraded version of sw26010 released with Sunway TaihuLight [43–45]. Taking sw26010pro as an example, there are six core groups on one chip. In each core group, there is one MPE (management process element) that acts as control elements and 64 CPEs (computing process elements) that act as computing elements. By offloading the computation into the CPEs, a high acceleration ratio is achieved for deep learning operators, e.g. Conv and Deconv [26].

Thanks to the modern programming paradigm and flexible libraries (i.e. sw/cu/roc-dnn or BLAS), the optimization of CNQS can be easily adapted onto both the traditional CPU and diverse heterogeneous accelerators.

For the sw26010pro CPU delivered with the new generation Sunway supercomputer, most computing power for each core group is provided by the 64 CPEs. Therefore it is important to utilize the CPEs to achieve high performance. Taking the compute-intensive operators used in CNN1 and CNN2 as examples, the many-core acceleration ratios for Conv2D, Conv2DBackpropInput and Conv2DBackpropFilter are 538×, 211× and 505×, respectively. Comparing to the compute-intensive operators, the many-core acceleration ratios of memory-intensive operators are limited due to the inability to achieve the same high CPE utilizations. The many-core acceleration ratio reaches 33.21–35.93 for Tile and 34.2–44.8 for Slice.

With the performance optimized operators, the overall acceleration ratio achieves considerable performance improvement over original MPE version. The many-core acceleration ratio for the CNN1 with 106 529 and 421 953 parameters are 90× and 130×, respectively.

For the MCMC+SR implemented on Sunway HPC, the elapsed time for one optimization step with respect to the core number is depicted in figure 5(a). In this case, $J2 = 0.5$ and L = 10, the total sample number is about $5\times 10^6$ and the network parameter number is 106 529. Under 40 000 processes, the elapsed time of one MCMC+SR step requires 1053 s with 64 CPEs acceleration, while the elapsed time for the MPE version is over 39 258 s. Therefore an acceleration ratio of approximately 37 times is achieved by utilizing the CPE cluster, compared to the MPE only. The MCMC+SR can be easily scaled to the whole system, with a maximum of 608 400 processes. It reports nearly 10× throughput improvement and achieves 63% parallel efficiency when scaling to nearly 40 million cores.

Zoom In Zoom Out Reset image size

For the Lanczos optimization, it is necessary to calculate the high order expectation of Hamiltonian like $\langle\hat{H}^2\rangle$ and $\langle\hat{H}^3\rangle$. Figure 5(b) depicts the elapsed time for calculating $\langle\hat{H}^2\rangle$ with respect to the NVIDIA A100 CUDA core number. In this case $J2 = 0.5$ and L = 8, the total sample number is 30 240 and the total number of spin configurations is on the order of magnitude of $30\,240\times L^4$. Different from the SR stage for handling the distributed covariance matrix, the Lanczos optimization only requires substantial CNN computation, which incurs little communication. Consequently, the forward calculations of all spin configurations are perfectly scalable to all the GPUs. As depicted by the figure, the parallel efficiency on 995 328 CUDA cores (144 GPUs) is 94%. In addition to the high scalability, for the double precision calculation of CNN wavefunctions, the speedup ratio of 55 296 CUDA cores (8 GPUs) to 640 CPU cores is 6.23.

Therefore, the optimization of CNN based wavefunctions has a high scalability and can achieve high acceleration ratio on heterogeneous systems.

Based on the linear relationship between energy and variance introduced in section 2.1: $E\approx E_{\mathrm{exact}}+\mathrm{const}.\times\sigma^2/N$, figure 6(a) reveals the energy extrapolations achieved by CNN1 within 106 529 parameters for the $J1$–$J2$ model, and figure 6(b) reveals the energy extrapolations achieved by CNN2 within 113 815 parameters for t-J model. Note that the extrapolated values are used only for estimation purposes, as only two data points are used.

Zoom In Zoom Out Reset image size

As depicted by figure 6(a), for the $J1$–$J2$ model in the maximal frustration region $J2 = 0.5$, the variational energies on the $10\times 10$ lattice for p = 0 and p = 1 are −0.497 168 and −0.497 468 respectively, and the extrapolated energy is −0.497 752. Furthermore, on the $12\times 12$ lattice, the variational energies for p = 0 and p = 1 are −0.496 636 and −0.496 989 respectively, and the extrapolated energy is −0.497 341.

As depicted by figure 6(b), for the t-J model with hole doping $n_h = 0.125$, the variational energies on the $12\times 8$ lattice for p = 0 and p = 1 are −0.608 793 and −0.620 9887 respectively, and the extrapolated energy is −0.642 031. On the $12\times 12$ lattice, the variational energies for p = 0 and p = 1 are −0.622 963 and −0.632 605 respectively, and the extrapolated energy is −0.648 179. From above results, we see that the Lanczos step can further improve the ground state energy.

The CNN1 used in this work has high quantum state representation ability, and therefore is more efficient comparing to other kinds of neural networks [46]. For example, for the $10\times 10$ lattice with $J2 = 0.5$, the ground state energy obtained by the CNN1 with 16 443 parameters without Lanczos step (i.e. p = 0) is −0.496 687, which is lower than −0.4960 obtained by the GNN-2 and the −0.4957 obtained by the lattice convolutional network (LCN) reported in [46]. Remarkably, the number of parameters of both GNN-2 and LCN is on the order of a million, which is significantly higher than that of CNN1.

Here we present the competitive energy results of the $J1$–$J2$ model and t-J model achieved in this work, and compare them to other state-of-the-art results. The energies are obtained by the CNN wavefunctions with a Lanczos step.

For the $J1$–$J2$ model, the energy values achieved by CNN1 with 106 529 parameters for various $J2$ and lattice sizes are listed in table 1. The energies are obtained by CNN1 optimized by MCMC+SR, and a Lanczos step. Specifically, in the maximal frustration region $J_2 = 0.5$, the convergence of energy with respect to the square lattice size L is depicted in figure 7 by red dots. In the figure, the solid line is fitted with the function $E = a/L^3+b$. Other state-of-the-art results are also shown in the figure for comparison. When L = 10, the energy in this work is −0.497 468, which has a small distance comparing to the −0.497 629 achieved by the PP+RBM method [25]. Increasing CNN parameter number and more than once independent optimization can further improve the energy precision. However, because of transfer learning, the energy obtained in this work for L = 18 is −0.496 513, which is $4.5\times 10^{-4}$ lower than the −0.496 275 reported in [25].

Zoom In Zoom Out Reset image size

Table 1. Energy values for various $J2$ values and lattice sizes, achieved by CNN1 with 106 529 parameters. The values are obtained by the Lanczos step with p = 1.

Description	0.46	0.48	0.50	0.55	0.60
$10\times 10$	−0.507 44	−0.502 31	−0.497 47	−0.486 70	−0.478 39
$12\times 12$	−0.506 97	−0.501 83	−0.496 99	−0.486 03	−0.477 15
$14\times 14$	−0.506 75	−0.501 59	−0.496 74	−0.485 81	−0.476 78
$16\times 16$	−0.506 62	−0.501 47	−0.496 59	−0.485 65	−0.476 56
$18\times 18$	−0.506 51	−0.501 39	−0.496 51	−0.485 57	−0.476 52

To date, there are no reference ground state energies for the t-J with PBC, therefore we compare the results to those of OBC calculated via PEPS in [47]. The parameter number of CNN2 is 113 815. On $8\times 8$ lattice, the energy achieved by CNN2 is −0.60 958, which is $3\times 10^{-3}$ higher than that reported in [47]. On $12\times 12$ lattice, the energy achieved by CNN2 is −0.63 261, which is $3.8\times 10^{-3}$ lower than that reported in [47]. We would like to note that the energies of different boundary conditions cannot be compared directly, which may have small differences. The results show here are only a demonstration of the effectiveness of CNN2 for fermionic lattice systems.