Recurrent neural network wave functions

A core technology that has emerged from the artificial intelligence revolution is the recurrent neural network (RNN). Its unique sequence-based architecture provides a tractable likelihood estimate with stable training paradigms, a combination that has precipitated many spectacular advances in natural language processing and neural machine translation. This architecture also makes a good candidate for a variational wave function, where the RNN parameters are tuned to learn the approximate ground state of a quantum Hamiltonian. In this paper, we demonstrate the ability of RNNs to represent several many-body wave functions, optimizing the variational parameters using a stochastic approach. Among other attractive features of these variational wave functions, their autoregressive nature allows for the efficient calculation of physical estimators by providing independent samples. We demonstrate the effectiveness of RNN wave functions by calculating ground-state energies, correlation functions, and entanglement entropies for several quantum spin models of interest to condensed-matter physicists in one and two spatial dimensions.

Figure 1

(a) Left-hand side: An RNN cell (green box) takes a sequence of inputs $\left\{{\mathit{\sigma }}_n\right\}$, where at each step $n$ the input ${\mathit{\sigma }}_{n-1}$ and the vector ${\mathit{h}}_{n-1}$ are fed in the RNN cell which generates a vector ${\mathit{h}}_n$ called the hidden state of the RNN. ${\mathit{h}}_n$ is meant to encode the history of the previous inputs ${\mathit{\sigma }}_{n^{\prime }<n}$. Moreover, the hidden state ${\mathit{h}}_n$ is fed to a fully connected layer with Softmax activation $S$ (magenta circles) to compute conditional probabilities. Right-hand side: The unrolled version of the RNN layer on the left-hand side. (b) A deep RNN model with $N_l$ stacked single RNN cells (green blocks) followed by a fully connected layer with activation function $A$ (magenta circle). Each single RNN cell at the $\ell \mathrm{th}$ layer has its corresponding hidden state ${\mathit{h}}_n^{\ell }$, which serves also as an input for the RNN cell at the $(\ell +1)\mathrm{th}$ layer. (c) A graphical representation of autoregressive sampling of RNNs.

Figure 2

(a) pRNN wave function: A graphical representation of the computation of positive amplitudes using one RNN cell along with a Softmax layer (magenta circles) to compute the modulus ${\left|\psi \left(\mathit{\sigma }\right)\right|}^2=P\left(\mathit{\sigma }\right)$. (b) cRNN wave function: A graphical representation of the computation of complex amplitudes using one RNN cell along with a Softmax layer (magenta circles) and a Softsign (SS) layer (orange circles). The first computes the modulus ${\left|\psi \left(\mathit{\sigma }\right)\right|}^2=P\left(\mathit{\sigma }\right)$; the second computes the phase $\varphi \left(\mathit{\sigma }\right)$ of $\psi \left(\mathit{\sigma }\right)$.

Figure 3

Results for the pRNN wave function compared with DMRG when targeting the ground state of a 1D TFIM at the critical point. Our pRNN wave function has one layer with 50 units. (a) The relative error $\epsilon$ and the energy variance per spin ${\sigma }^2$ against the number of training steps (i.e., gradient descent steps) for $N=1000$ spins. We use only 200 samples per gradient step, which are enough to achieve convergence. (b) The two-point correlation function $\langle {\widehat{S}}_{40}{\widehat{S}}_n\rangle$ along the $x$ axis and $z$ axis of the optimized pRNN wave function for sites $n>40$ using ${10}^6$ samples. DMRG results are also shown for comparison. (c) The Rényi entropy $S_2$ against the relative size of subregion $A$ for system sizes $N=20$ and 80. In both (b) and (c), the error bars are smaller than the data points.

Figure 4

The relative error (compared to DMRG) of the cRNN wave function trained on the 1D $J_1\text{−}{J}_2$ model with $N=100$ spins for different values $J_2$, both without a prior sign (represented by “No Sign”) and with a prior Marshall sign as in Eq. (20) (represented by “Marshall Sign”). We observe that applying a Marshall sign improves the accuracy.

Figure 5

(a) Autoregressive sampling path of 2D spin configurations using 1D RNN wave functions. The 2D configurations are generated through raster scanning, such that in order to generate spin ${\sigma }_i$ one has to condition on the spins that are previously generated. (b) Autoregressive sampling path of 2D spin configurations using 2D RNN wave functions through a zigzag path, where each site receives two hidden states and two spins from the horizontal and the vertical neighbors that were previously generated. For both panels (a) and (b), the digits and the green dashed arrows indicate the sampling path, while the red arrows indicate how the hidden states are passed from one site to another. (c) A comparison of the variational energy per spin between a 2D pRNN wave function (labeled as 2DRNN), 1D pRNN wave function (labeled as 1DRNN), PixelCNN wave function [56], and DMRG with bond dimension $\chi$ for the 2D TFIM on a system with $L_x\times L_y=12\times 12$ spins. The shaded regions represent the error bars of each method. Note the broken $y$ axis on the plots for $h=3$ and 4, denoting a change in scale between the upper and lower portions of the plots. These results show that 2D pRNN wave functions can achieve a performance comparable to PixelCNN wave functions and DMRG with a large bond dimension, while requiring only a fraction of their variational parameters.

Figure 6

The energy variance per spin against the number of memory units of a 1D pRNN wave function trained at the critical point of (a) the 1D TFIM and (b) the 2D TFIM. Both scalings show that we can systematically reduce the bias in the estimation of the ground-state energy.

Figure 7

Graphical representation of the gated recurrent unit cell described in Eq. (A1). The magenta circles/ellipses represent pointwise operations such as vector addition or multiplication. The blue rectangles represent neural network layers labeled by the nonlinearity we use. Merging lines denote vector concatenation and forking lines denote a copy operation. The sigmoid activation function is represented by $\sigma$.

Figure 8

The Swap operator acting on the tensor product of two samples $\mathit{\sigma }$ and ${\mathit{\sigma }}^{\prime }$.

Figure 9

The energy variance per spin against the number of samples, which suggests that the energy variance saturates and does not improve further by using a larger number of samples for training.

Figure 10

Scaling study of the energy variance per spin vs the number of layers of a pRNN wave function such that all pRNN wave functions with different layers have approximately the same number of variational parameters. The results show that fixing the number of parameters while changing the number of layers does not affect the energy variance obtained by the pRNN wave function.