SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. Spin flip symmetry exists for quantum chemistry systems that do not couple spin up and spin down electron, so that the wavefunction is invariant to inputs when spin up and down are flipped. (Barrett et al., 2022) has shown that autoregressive neural network can preserve  $\mathbb{Z}_2$  spin flip symmetry as a function symmetry, which ANTN inherits.
Corollary B.5.2 (Discrete Abelian and Non-Abelian Symmetries). Autoregressive Neural TensorNet can realize discrete Abelian and Non-Abelian symmetries.
Proof. Gauge symmetry is a local symmetry such that the function is invariant under local transformation. (Luo et al., 2021b) has shown that autoregressive neural networks can preserve discrete Abelian and non-Abelian symmetries as either the mask symmetry or the function symmetry, which ANTN inherits.
C Additional Experiments
In this section, we perform additional experiments to further compare ANTN and ARNN, and show that
- It is challenging to surpass ANTN by ARNN even with increased number of parameters.
- The improvement from ANTN mainly comes from the effective inductive bias of MPS structure instead of the DMRG initialization.
- ANTN has a favorable scaling compared to ARNN and MPS for reaching accurate results.
Energy per site ↓ and additional information for 10 × 10 system
Algorithms	J2=0.2	J2=0.5	J2=0.8	Nparams,TN	Nparams,NN	Nparams	Ttraining
MPS (8)	-1.997537	-1.734326	-1.869871	1.2 × 104	0	1.2 × 104	< 1 hr
MPS (70)	-2.191048	-1.866159	-2.007845	8.8 × 105	0	8.8 × 105	< 1 hr
MPS (1024)	-2.255633	-1.938770	-2.062730	1.7 × 108	0	1.7 × 108	242 hrs
PixelCNN (7-48)	-2.22462(24)	-1.74098(29)	-1.98331(17)	0	5.0 × 105	5.0 × 105	43 hrs
PixelCNN (P) (7-48)	-2.1920(9)	-1.6458(10)	-1.9468(6)	0	5.0 × 105	5.0 × 105	19 hrs†
PixelCNN (8-48)	-2.23581(14)	-1.86922(16)	-2.01093(16)	0	5.7 × 105	5.7 × 105	59 hrs
EW (7-48-8)	-2.26034(6)	-1.94001(6)	-2.05707(43)	1.2 × 106	5.0 × 105	1.7 × 106	89 hrs
EW (RI) (7-48-8)	-2.25946(6)	-1.94030(9)	-2.05125(11)	1.2 × 106	5.0 × 105	1.7 × 106	46 hrs†
BW (7-48-70)	-2.25755(8)	-1.93842(42)	-2.05088(14)	8.9 × 105	5.0 × 105	1.4 × 106	136 hrs
Table 3: Energy per site  $\downarrow$ , tensor network parameters  $N_{\mathrm{params,TN}}$ , neural network parameters  $N_{\mathrm{params,NN}}$ , total parameters  $N_{\mathrm{params}}$  and runtime  $T_{\mathrm{training}}$  for  $10 \times 10$  system with various algorithms. Elementwise (EW) and blockwise (BW) are two constructions of ANTN with PixelCNN as the underlying ARNN. For MPS, the number in the parenthesis labels the bond dimension. For PixelCNN, the two numbers label the number of layers and hidden dimension respectively. For ANTN, the numbers label (in this order) the number of layers, hidden dimension, and bond dimension. For PixelCNN, we also test pertaining (P) using MPS result (with bond dimension the last number in parenthesis), and for EW, we in addition test random initialization (RI) of the MPS part. The algorithms in boldface are new experiments not previously shown in the main paper. Best energies for each  $J_2$  are also highlighted in boldface. The MPS with DMRG algorithm is run on CPUs, while the PixelCNN, elementwise (EW) and blockwise (BW) ANTN are trained using GPUs, with PixelCNN (P) (7-48) and EW (RI) (7-48-8) using A100 GPUs (labeled with  $\dagger$ ), and the rest using V100 GPUs.
Figure 4: Energy per site vs total number of parameters and runtime (in hours) for various algorithms (MPS with DMRG algorithm, PixelCNN, elementwise (EW), and blockwise (BW) ANTNs) and system sizes  $(10 \times 10$  and  $12 \times 12)$  for  $J_{2} = 0.5$ . The ANTN construction uses PixelCNN as the underlying ARNN. The total number of parameters includes parameters from both the TN part and the ARNN part.
C.1 Comparison between ANTN and ARNN with More Layers
In Table 3 and 4 we include additional results.
Energy per site  $\downarrow$  and additional information for  $10\times 10$  system at  $J_{2} = 0.5$
Algorithms	J2=0.5	Nparams,TN	Nparams,NN	Nparams	Ttraining
MPS (8)	-1.734326	1.2 × 104	0	1.2 × 104	< 1 hr
MPS (70)	-1.866159	8.8 × 105	0	8.8 × 105	< 1 hr
MPS (1024)	-1.938770	1.7 × 108	0	1.7 × 108	242 hrs
PixelCNN (7-48)	-1.74098(29)	0	5.0 × 105	5.0 × 105	19 hrs (A100) or 43 hrs (V100)
PixelCNN (8-48)	-1.86922(16)	0	5.7 × 105	5.7 × 105	59 hrs (V100)
PixelCNN (9-48)	-1.90021(15)	0	6.3 × 105	6.3 × 105	21 hrs (A100)
PixelCNN (10-48)	-1.90440(14)	0	7.0 × 105	7.0 × 105	23 hrs (A100)
PixelCNN (11-48)	-1.92826(11)	0	7.7 × 105	7.7 × 105	26 hrs (A100)
PixelCNN (12-48)	-1.92986(10)	0	8.4 × 105	8.4 × 105	30 hrs (A100)
PixelCNN (15-48)	-1.92869(11)	0	1.0 × 106	1.0 × 106	51 hrs (A100)
PixelCNN (20-48)	-1.91537(13)	0	1.4 × 106	1.4 × 106	66 hrs (A100)
PixelCNN (8-72)	-1.88536(19)	0	1.3 × 106	1.3 × 106	35 hrs (A100)
PixelCNN (9-72)	-1.89266(15)	0	1.4 × 106	1.4 × 106	40 hrs (A100)
PixelCNN (10-72)	-1.90718(15)	0	1.6 × 106	1.6 × 106	47 hrs (A100)
PixelCNN (11-72)	-1.92965(13)	0	1.7 × 106	1.7 × 106	52 hrs (A100)
PixelCNN (12-72)	-1.91186(14)	0	1.9 × 106	1.9 × 106	57 hrs (A100)
PixelCNN (S) (11-48)	-1.92915(12)	0	7.7 × 105	7.7 × 105	31 hrs (A100)
PixelCNN (S) (12-48)	-1.93112(11)	0	8.4 × 105	8.4 × 105	34 hrs (A100)
PixelCNN (S) (15-48)	-1.93233(9)	0	1.0 × 106	1.0 × 106	47 hrs (A100)
PixelCNN (S) (20-48)	-1.93058(10)	0	1.4 × 106	1.4 × 106	64 hrs (A100)
PixelCNN (S) (11-72)	-1.93211(10)	0	1.7 × 106	1.7 × 106	54 hrs (A100)
PixelCNN (S) (12-72)	-1.92812(12)	0	1.9 × 106	1.9 × 106	59 hrs (A100)
EW (7-48-8)	-1.94001(6)	1.2 × 106	5.0 × 105	1.7 × 106	46 hrs (A100) or 89 hrs (V100)
BW (7-48-70)	-1.93842(42)	8.9 × 105	5.0 × 105	1.4 × 106	136 hrs (V100)
Table 4: Energy per site  $\downarrow$ , tensor network parameters  $N_{\mathrm{params,TN}}$ , neural network parameters  $N_{\mathrm{params,NN}}$ , total parameters  $N_{\mathrm{params}}$  and runtime  $T_{\mathrm{training}}$  for  $10 \times 10$  system at  $J_2 = 0.5$  with various algorithms. Elementwise (EW) and blockwise (BW) are two constructions of ANTN with PixelCNN as the underlying ARNN. For MPS, the number in the parenthesis labels the bond dimension. For PixelCNN, the two numbers label the number of layers and hidden dimension respectively. For ANTN, the numbers label (in this order) the number of layers, hidden dimension, and bond dimension. For PixelCNN, we also test the sign rule (S) The algorithms in boldface are new experiments not previously shown in the main paper. Best energies for each  $J_2$  are also highlighted in boldface. The MPS with DMRG algorithm is run on CPUs, while the PixelCNN, elementwise (EW) and blockwise (BW) ANTN are trained using GPUs.
Additional information for  $12 \times 12$  system
Algorithms	Nparams,TN	Nparams,NN	Nparams	Ttraining
MPS (8)	1.8 × 104	0	1.8 × 104	< 1 hr
MPS (70)	1.3 × 106	0	1.3 × 106	< 1 hr
MPS (1024)	2.6 × 108	0	2.6 × 108	329 hrs
PixelCNN (7-48)	0	5.0 × 105	5.0 × 105	82 hrs
EW (7-48-8)	1.7 × 106	5.0 × 105	2.2 × 106	183 hrs
BW (7-48-70)	1.3 × 106	5.0 × 105	1.8 × 106	287 hrs
Table 5: Additional information for the number of parameters and runtime of the  $12 \times 12$  experiment in the paper for various algorithms. The runtime is average over different  $J_{2}$ 's. The MPS with DMRG algorithm is run on CPUs, while the PixelCNN, elementwise (EW), and blockwise (BW) ANTN are trained using V100 GPUs. The ANTN uses PixelCNN as the underlying ARNN. For MPS, the number in the parenthesis labels the bond dimension. For PixelCNN, the two numbers label the number of layers and hidden dimension respectively. For ANTN, the numbers label (in this order) the number of layers, hidden dimension, and bond dimension.
In Table 3, we perform 2 addition tests: a) pretrain pure PixelCNN by learning MPS results; b) train ANTN (elementwise) without DMRG initialization. We find that DMRG pretraining negatively affects the result, while ANTN without DMRG initialization performs as good as with DMRG initialization. This is understandable, as learning MPS results with complex sign structure is essentially the same task as learning a quantum state from shallow random circuit, which has been shown a difficult task for ARNN in Figure 2(b) of the original paper. Since the MPS result is an approximation of the actual ground state, learning this state actually makes a worse initialization for PixelCNN. (We note that the PixelCNN was originally initialized from the result of smaller system sizes, as we described in Appendix D under transfer learning.) On the other hand, the ANTN, despite not using DMRG initialization, still retains the ability to represent flexible sign structures. In addition, since DMRG is just an optimization algorithm to train MPS, the gradient-descent-based algorithm should still train the MPS part of ANTN reasonably well without DMRG initialization.
In Table 4, we test a series of ARNNs with different layers (up to 20 layers) and hidden dimensions (up to 72 hidden dimensions) such that the largest ARNN tested has more parameters than ANTN. We found that all the ARNN energies are not as good at ANTN (which is only based on 7 layers and 48 hidden dimensions), despite that ARNNs have more parameters and take a longer time to train using the same GPU (see new data below). Furthermore, we observe that an increase of parameters can help ARNN to improve energy in general, but this improvement hits a diminishing return potentially due to difficulties in optimization as the number of parameters increases. In addition, we further tested the (approximate) sign rule and found that it could help ARNN obtain better energies but still worse than ANTN without the sign rule. This indicates that the ANTN has the advantage of inheriting the flexible sign structure from MPS, avoiding the manual bias of adding the sign rule. The new results have provided strong support for our expectation, both from an expressivity perspective and from the fact that ANTN has a better physics inductive bias than ARNN for optimization. We have also added a plot for the new benchmark in the updated manuscript. Although the added layer improves the result to some extent, the PixelCNN fails to beat ANTN in terms of energy calculation. This is consistent with Theorem 5.3 that ANTN is more expressive than ARNN.
C.2 Comparison of Runtime and Number of Parameters
In Table 3, Table 4, Table 5 and Figure 4, we show the number of parameters and runtime of each algorithm. In summary, MPS has the most parameters, and is the slowest with DMRG algorithm, taking as long as 2 weeks. The ANTN is much more efficient compared to MPS while obtaining better energies, and it is comparable to pure ARNN while obtaining much better energies. Within the two types of ANTN, the elementwise construction, despite using more parameters, runs faster than the blockwise construction. We note that since TNs and ARNNs use the parameters very differently, a direct comparison of the two types of parameters may be unfair. Therefore, we list both the total number of parameters and the parameters for the individual parts (TN part and ARNN part).
D Experimental Details and Reproducibility
We note that while the following details can be used to reproduce our results, additional adjustments may be possible to further improve our results. The code associated with this paper is available at https://github.com/antn2023/ANTN
Implementation. The transformer (Vaswani et al., 2017) used in this work is a decoder-only transformer implemented in (Luo et al., 2020), and the PixelCNN used is the gated PixelCNN (Van den Oord et al., 2016) implemented in (Chen et al., 2022) but without the channelwise mask.
Hyperparameters. For the transformer in quantum state learning, we use 2 layers, 32 hidden dimensions, and 16 attention heads, whereas for the PixelCNN in variational Monte Carlo, we use 7 layers with dilations 1, 2, 1, 4, 1, 2, and 1 for each layer and 48 hidden dimensions. For additional experiments with more layers, we always use dilation of 1 for the layer 8 and beyond.
Initialization. Throughout the experiment, we initialize the weight matrices of the last fully connected layer of the ARNN and ANTN to 0, and initialize the bias according to the following rule:
- The bias of ARNN is always randomly initialized.
- For the random Bell state learning, the bias of ANTN is randomly initialized
- For the rest of the experiment, the bias of ANTN is set to 0.
The rest of the parameters of ARNN and ANTN are randomly initialized according to PyTorch (Paszke et al., 2019)'s default initialization scheme.
In addition, for the random Bell state learning experiment, we do not initialize the ANTN with MPS (and hence the random bias), whereas for the shallow random circuit learning experiment, we initialize the ANTN with an MPS that comes from explicit truncation of the full wavefunction, and for the variational Monte Carlo experiment, we initialize the ANTN with DMRG optimized MPS.
**Optimization.** Through the experiment, we use the Adam optimizer with an initial learning rate of 0.01. For quantum state learning experiments, we train the ARNN and ANTN with full basis for 2000 iterations, where the learning rate is halved at iterations 600, 1200, and 1800. In case the full basis cannot be computed in one pass due to GPU memory constraint, we divide the full basis into several batches and accumulates the gradient in each iteration. For variational Monte Carlo experiments, we train the ARNN and ANTN stochastically until the energy converges. During the training, we halve the learning rate at iterations 100, 500, 1000, 1800, 2500, and 4000. In each experiment, we choose the maximum batch size that can be fed into the GPU memory, while using accumulation steps to keep the effective batch size around 10000 throughout the training.
Transfer Learning. We in addition take advantage of transfer learning. For  $J_{2} = 0.2$ , 0.5, and 0.8, we begin from scratch with the  $4 \times 4$  system. Then, an  $N \times N$  system is initialized from an  $(N - 2) \times (N - 2)$  system by keeping all but the weights in the last layer. Then, for other  $J_{2}$ , the neural network is initialized from one of  $J_{2} = 0.2$ , 0.5, and 0.8 that is closest to the current  $J_{2}$  of the same system size by keeping all but the weights in the last layer. We find this allows the optimization to converge faster than always starting from scratch.
Computational Resources. The models are mainly trained using NVIDIA V100 GPUs with 32 GB memory and Intel Xeon Gold 6248 CPUs, with some models trained using NVIDIA A100 GPUs with 80 GB memory.