Transfer learning for scalability of neural-network quantum states

Neural-network quantum states have shown great potential for the study of many-body quantum systems. In statistical machine learning, transfer learning designates protocols reusing features of a machine learning model trained for a problem to solve a possibly related but different problem. We propose to evaluate the potential of transfer learning to improve the scalability of neural-network quantum states. We devise and present physics-inspired transfer learning protocols, reusing the features of neural-network quantum states learned for the computation of the ground state of a small system for systems of larger sizes. We implement different protocols for restricted Boltzmann machines on general-purpose graphics processing units. This implementation alone yields a speedup over existing implementations on multicore and distributed central processing units in comparable settings. We empirically and comparatively evaluate the efficiency (time) and effectiveness (accuracy) of different transfer learning protocols as we scale the system size in different models and different quantum phases. Namely, we consider both the transverse field Ising and Heisenberg $XXZ$ models in one dimension, as well as in two dimensions for the latter, with system sizes up to 128 and $8\times 8$ spins. We empirically demonstrate that some of the transfer learning protocols that we have devised can be far more effective and efficient than starting from neural-network quantum states with randomly initialized parameters.

Figure 1

Structure of a restricted Boltzmann machine with $L_v$ visible nodes and $L_h$ hidden nodes. The visible layer consists of visible nodes $x_1,\cdots ,x_{L_v}$. The hidden layer consists of hidden nodes $h_1,\cdots ,h_{L_h}$. The connections between the visible and hidden layer are given by the weight matrix $\mathit{W}$. The biases for the visible and hidden layers are represented by $a_1,\cdots ,a_{L_v}$ and $b_1,\cdots ,b_{L_h}$, respectively.

Figure 2

Schematic representation of the different transfer learning protocols used to scale up one-dimensional systems. Here $L_v$ and $L_h$, and $L_v^{\prime }$ and $L_h^{\prime }$ correspond to the number of visible and hidden nodes for the base and target networks, respectively. The different panels show how to construct the target weight matrix $\mathit{W}$ from the base weight matrix $\tilde{\mathit{W}}$ by replicating the colored rows and filling the gray cells with random entries: (a) (1,2) tiling, (b) (2,2) tiling, (c) $(L,2)$ tiling (the scaling factor is 2), and (d) $(L,4)$ tiling (scaling up the network by a factor 4; see the text for a detailed explanation).

Figure 3

Computation times $t$ (in seconds) obtained for different numerical strategies for the one-dimensional ferromagnetic Ising model ($J_I=1$ and $h=1$) as a function of the number of spins $L_v=8$, 16, 32, 64, and 128. The purple dash-dotted line with diamonds shows $t$ using four central processing units with netket. The red dashed line with triangles shows $t$ using eight central processing units with netket. The blue solid line with squares shows $t$ using 16 central processing units with netket. The black dotted line with circles shows $t$ using a graphics processing unit with tensorflow.

Figure 4

Efficiency of different transfer learning protocols: (a) Ising antiferromagnetic, (b) Heis antiferromagnetic, (c) Ising paramagnetic, (d) Heis $XY$, (e) Ising ferromagnetic, and (f) Heis ferromagnetic. Each panel shows the computation time $t$ (in seconds) obtained for different transfer learning protocols until the stopping criterion is reached as a function of the number of spins $L_v$. The blue solid line with squares shows the cold start, the red dashed line with circles $(L,2)$ tiling, the black dotted line with diamonds (2,2) tiling, and the green dash-dotted line with triangles (1,2) tiling. The error bars represent the $\pm 1$ standard deviation of the interval. The left column shows the Ising model with $h=1$ and (a) $J_I=-2$, (c) $J_I=0.5$, and (e) $J_I=2$, while the right column shows the Heisenberg model with $J_{XXZ}=1$ and (b) $\mathrm{\Delta }=-2$, (d) $\mathrm{\Delta }=-0.5$, and (f) $\mathrm{\Delta }=2$.

Figure 5

Weight matrices for a single configuration of the one-dimensional ferromagnetic Ising model ($J_I=2$ and $h=1$) with 16 spins and $\alpha =2$. The rows and columns of the matrix correspond to the visible ($L_v=16$) and the hidden ($L_h=\alpha L_v=32$) nodes of the restricted Boltzmann machine, respectively. For each row, the left panel shows the initial weight matrix $W^{\left(\mathrm{init}\right)}$ obtained for a given transfer learning protocol from an eight-spin chain and the right panel shows the final fine-tuned weight matrix $W^{\left(\mathrm{final}\right)}$ obtained after the stopping criterion is met: (a) and (b) (1,2)-tiling protocol for transfer distance $\mathcal{D}=0.01145$, (c) and (d) (2,2)-tiling protocol for $\mathcal{D}=0.01704$, and (e) and (f) $(L,2)$-tiling protocol for $\mathcal{D}=0.00162$.

Figure 6

Same as Fig. 5 for the one-dimensional antiferromagnetic Ising model ($J_I=-2$ and $h=1$). The transfer distances $\mathcal{D}$ for the (1,2)-tiling, (2,2)-tiling, and $(L,2)$-tiling protocols are $0.05423, 0.01858$, and $0.00224$, respectively.

Figure 7

Relative ground-state energy errors $\mathrm{\Delta }E/E_0$ obtained for different transfer learning protocols (including the cold start) and relative to accurate matrix product state simulations. Computations are run on a one-dimensional system with 128 spins initialized with a weight matrix generated from the fine-tuned weight matrix obtained with the most effective transfer learning protocol for the same system with 64 spins until the stopping criterion is reached; $\mathrm{\Delta }E/E_0$ is computed when the stopping criterion is met. (a) Ising model with $J_I=-2$, 0.5, and 2 ($h=1$) and (b) Heisenberg model with $\mathrm{\Delta }=-2, -0.5$, and 2 ($J_{XXZ}=1$). The blue open bar shows the cold start, the red closed bar $(L,2)$ tiling, the green closed bar (1,2) tiling, and the black open bar (2,2) tiling. The error bars represent the interval of the $9\text{th}$ percentile and the $91\text{st}$ percentile.

Figure 8

Ferromagnetic ($C_d^F$) and antiferromagnetic ($C_d^A$) ground-state correlators for one realization of a one-dimensional system of 128 spins as a function of the spin-spin separation $d$ ($d=1,...,60$). The different panels give the results obtained with the matrix product state algorithm, the cold-start, and the different transfer learning protocols when the stopping criterion is met. The initial weight matrix of the different tiling protocols is generated from the fine-tuned weights produced by the most effective transfer learning protocol operated on the same system with 64 spins. The ferromagnetic correlator $C_d^F$ is shown for the (a) Ising ($J_I=2$ and $h=1$) and (b) Heisenberg ($\mathrm{\Delta }=2$ and $J_{XXZ}=1$) models and the antiferromagnetic correlator $C_d^A$ is shown for the (c) Ising ($J_I=-2$ and $h=1$) and (d) Heisenberg ($\mathrm{\Delta }=-2$ and $J_{XXZ}=1$) models. The cyan solid line shows the matrix product state result, the blue solid line the cold-start protocol, the red dashed line the $(L,2)$-tiling protocol, the green dash-dotted line the (1,2)-tiling protocol, and the black dotted line the (2,2)-tiling protocol.

Figure 9

Zoomed-in version of Fig. 8 to better see the performance achieved by each transfer learning protocol. We refer to the caption of Fig. 8 for a complete description of this figure.

Figure 10

Efficiency and effectiveness of different size-increase scenarios to reach 128 spins starting from four spins for the one-dimensional Ising model at $h=1$ for different $J_I$. All quantities are computed until the stopping criterion is reached. (a) Efficiency as reflected by the computation time $t$. (b) Effectiveness as reflected by the ground-state energy error $\mathrm{\Delta }E/E_0$ relative to the matrix product state simulation. (c) Effectiveness as reflected by the spin-spin correlation errors $\mathrm{\Delta }{C}^F/C_0^F$ ($J_I=2$) and $\mathrm{\Delta }{C}^A/C_0^A$ ($J_I=-2$) computed at spin separation distance $d=60$ and relative to the matrix product state simulation. The different scenarios tested are (see the text for full details) SCE:JUMP (black open bar), SCE:DOUBLE (blue closed bar), SCE:4-8 (green open bar), SCE:8-4 (red closed bar), SCE:2-4-4 (cyan closed bar), SCE:4-2-4 (blue open bar), and SCE:4-4-2 (green closed bar). The error bars represent the interval of the $9\text{th}$ percentile and of the $91\text{st}$ percentile.

Figure 11

Efficiency and effectiveness of cold-start and different transfer learning protocols run for the two-dimensional isotropic antiferromagnetic Heisenberg model ($J_{XXZ}=1$ and $\mathrm{\Delta }=-1$) on a square lattice of linear size $D$ with open boundary conditions until the stopping criterion is reached. (a) and (b) Efficiency as reflected by the computation times $t$ (in seconds) achieved by the different transfer learning protocols for a lattice of linear size (a) $D=4$ and (b) $D=8$. (c) and (d) Effectiveness as reflected by the computed ground-state energy for (c) $D=4$ and (d) $D=8$, with the mean relative error $\mathrm{\Delta }E/E_0$ achieved by the cold-start and different transfer learning protocols as compared to quantum Monte Carlo simulations. The blue open bars are for the cold-start protocol, the red closed bars are for the $(L,2)$-tiling protocol, the green closed bars are for the (1,2)-tiling protocol, and the black open bars are for the (2,2)-tiling protocol. (e) and (f) Effectiveness as reflected by computing the antiferromagnetic correlator ${(-1)}^l\langle {\sigma }_{\mathbf{i}}^z{\sigma }_{\mathbf{i}+l\mathbf{e}}^z\rangle$ for $D=8$ as a function of $l$ from lattice node $\mathbf{i}=(4,5)$. Note that the correlator equals 1 for $l=0$. (e) Correlations along the diagonal direction $\mathbf{e}=(1,1)$. (f) Correlations along the row direction $\mathbf{e}=(1,0)$. The cyan solid line corresponds to quantum Monte Carlo simulations, the blue solid line corresponds to the cold-start protocol, the red dashed line corresponds to the $(L,2)$-tiling protocol, the green dash-dotted line corresponds to the (1,2)-tiling protocol, and the black dotted line corresponds to the (2,2)-tiling protocol.

Figure 12

Ground-state energy $E$ as a function of the number of iterations $N\ge 1$ for one realization of the one-dimensional Ising model with 128 spins: (a) the ferromagnetic ($J_I=2$ and $h=1$) Ising model and (b) the antiferromagnetic ($J_I=-2$ and $h=1$) Ising model. To smooth out the computed curves, the energy value plotted at iteration $N$ is the average of the energies obtained in the iteration window $[N_0,N]$, where $N_0=\max (N-100,0)$. The blue solid line shows the cold-start protocol, the red dashed line the $(L,2)$-tiling protocol, the black dotted line the (2,2)-tiling protocol, the green dash-dotted line the (1,2)-tiling protocol, and the cyan solid line the matrix product state simulation.