SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhuo Chen $^{123}$  Laker Newhouse $^{4*}$  Eddie Chen $^{3*}$  Di Luo $^{1235\dagger}$  Marin Soljačić $^{13}$
<sup>1</sup>NSF AI Institute for Artificial Intelligence and Fundamental Interactions
$^{2}$ Center for Theoretical Physics, Massachusetts Institute of Technology
$^{3}$ Department of Physics, Massachusetts Institute of Technology
$^{4}$ Department of Mathematics, Massachusetts Institute of Technology
$^{5}$ Department of Physics, Harvard University
{chenzhuo, lakern, ezchen, diluo, soljacic}@mit.edu
Quantum many-body physics simulation has important impacts on understanding fundamental science and has applications to quantum materials design and quantum technology. However, due to the exponentially growing size of the Hilbert space with respect to the particle number, a direct simulation is intractable. While representing quantum states with tensor networks and neural networks are the two state-of-the-art methods for approximate simulations, each has its own limitations in terms of expressivity and inductive bias. To address these challenges, we develop a novel architecture, Autoregressive Neural TensorNet (ANTN), which bridges tensor networks and autoregressive neural networks. We show that Autoregressive Neural TensorNet parameterizes normalized wavefunctions, allows for exact sampling, generalizes the expressivity of tensor networks and autoregressive neural networks, and inherits a variety of symmetries from autoregressive neural networks. We demonstrate our approach on quantum state learning as well as finding the ground state of the challenging 2D  $J_{1} - J_{2}$  Heisenberg model with different systems sizes and coupling parameters, outperforming both tensor networks and autoregressive neural networks. Our work opens up new opportunities for quantum many-body physics simulation, quantum technology design, and generative modeling in artificial intelligence.
Quantum many-body physics is fundamental to our understanding of the universe. It appears in high energy physics where all the fundamental interactions in the Standard Model, such as quantum electrodynamics (QED) and quantum chromodynamics (QCD), are described by quantum mechanics. In condensed matter physics, quantum many-body physics has led to a number of rich phenomena and exotic quantum matters, including superfluids, superconductivity, the quantum Hall effect, and topological ordered states (Girvin & Yang, 2019). As an application, quantum many-body physics is crucial for new materials design. The electronic structure problem and chemical reactions in quantum chemistry are governed by quantum many-body physics. The recent development of quantum computers is also deeply connected to quantum many-body physics. A multi-qubit quantum device is intrinsically a quantum many-body system, such that progress on quantum computer engineering is tied to our understanding of quantum many-body physics (Preskill, 2021).
*equal contribution
† corresponding author: diluo@mit.edu
All information of a closed quantum many-body system is captured by the wavefunction, whose properties are described by the famous Schrödinger equation. An important tool to study and understand quantum many-body physics is to classically simulate the wavefunction. However, the wavefunction is a high dimensional function in Hilbert space, whose dimension grows exponentially with the number of particles. For example, for qubit systems where each qubit has two degrees of freedom), the wavefunction of 300 qubits will have dimension  $2^{300}$  which is larger than the number of atoms in the observable universe. Furthermore, the Schrödinger equation is a complex-valued
high-dimensional equation, which is challenging to solve or simulate in general.
Figure 1: Diagrammatic representation of autoregressive neural network (ARNN), tensor network (TN) and our Autoregressive Neural TensorNet (ANTN).
A number of algorithms have been developed to simulate quantum many-body physics, including quantum Monte Carlo, tensor networks, neural network quantum states, and quantum computation. In particular, computing the ground state of quantum many-body systems is of great interest. One important approach is the variational principle, which provides an upper bound for the ground state energy. To apply the variational principle successfully, one must design an ansatz that can represent and optimize the wavefunction efficiently. Tensor networks and neural network quantum states (Carleo & Troyer, 2017) are the two main state-of-the-art methods that can be applied with the variational principle for quantum many-body simulation. However, tensor networks usually suffer from an expressivity issue in systems with more than one dimension, while neural network quantum states usually lack inductive bias from the underlying physics structure and have sign structure challenges in the representation.
In this paper, we develop a novel architecture, Autoregressive Neural TensorNet (ANTN), to bridge neural network quantum states and tensor networks, achieving the best of both worlds. In particular, our contributions are threefold:
- Develop ANTN with two variants called "elementwise" and "blockwise," which each naturally generalize the two state-of-the-arts ansatzes, tensor networks (TN) and autoregressive neural networks (ARNN), to provide proper inductive bias and high expressivity.
- Prove that ANTN is normalized with exact sampling, has generalized expressivity of TN and ARNN, and inherits multiple symmetries from ARNN.
- Demonstrate our methods on quantum state learning and variationally finding the ground state of the challenging 2D  $J_{1}$ - $J_{2}$  Heisenberg model, outperforming both TN and ARNN.
Tensor Networks (TN) represent high-dimensional wavefunctions using low-rank tensor decomposition, notably matrix product state (MPS) (Vidal, 2003, 2004), PEPS (Verstraete & Cirac, 2004), and MERA (Vidal, 2007). They capture the entanglement structure of physical systems and, with algorithms like the density matrix renormalization group (DMRG) (White, 1992), are used for state simulations and real-time dynamics. However, their expressivity can be limited in systems with more than one dimension and systems with high entanglement. In machine learning, TN appears as tensor train (Oseledets, 2011) and CP decomposition methods.
Neural Network Quantum State (NNQS) leverages neural networks for high dimensional wavefunction representation (Carleo & Troyer, 2017). It's been demonstrated that many quantum states can be approximated or represented by NNQS (Sharir et al., 2021; Gao & Duan, 2017; Lu et al., 2019; Levine et al., 2019a; Luo et al., 2021b,a; Deng et al., 2017; Huang & Moore, 2021; Vieijra et al., 2020), and it has yielded state-of-the-art results in computing quantum system properties (Gutiérrez & Mendl, 2020; Schmitt & Heyl, 2020; Vicentini et al., 2019; Yoshioka & Hamazaki, 2019; Hartmann & Carleo, 2019; Nagy & Savona, 2019; Luo et al., 2021b, 2022a). Key advancements in NNQS include ARNN that improves sample efficiency and gradient estimation, and the development of
neural networks that adhere to the underlying symmetries of quantum systems (Luo et al., 2021b; Hibat-Allah et al., 2020; Choo et al., 2019; Luo & Clark, 2019; Hermann et al., 2019; Pfau et al., 2020; Luo et al., 2021a, 2022b, 2021b; Chen et al., 2022). Despite its potential, NNQS faces challenges such as the lack of physics prior and the sign structure problem. While there are attempts to integrate NNQS with TN, including matrix product state with neural network backflow (Lami et al., 2022) and generalizing MPS to RNN (Wu et al., 2022), the former can not produce normalized wavefunctions, while the latter only builds on tensor contractions that lacks the nonlinear activation functions of standard neural network wavefunctions.
3.1 Quantum Preliminaries
In this work, we focus on qubit systems in quantum many-body physics. The wavefunction or quantum state  $\psi$  of the system is a normalized function  $\psi : \mathbb{Z}_2^n \to \mathbb{C}$  with  $\sum_{\boldsymbol{x}} |\psi(\boldsymbol{x})|^2 = 1$ , where  $n$  is the system size. The input to the wavefunction is an  $n$ -bit string  $\boldsymbol{x} \in \{0,1\}^n$ . Therefore, the wavefunction  $\psi$  can be viewed as a complex-valued vector of size  $2^n$ , with Dirac notation  $\langle \psi|$  and  $|\psi\rangle$  correspond to a conjugate row vector and a column vector respectively, and  $\psi(\boldsymbol{x}) = \langle \boldsymbol{x}|\psi\rangle$ . Because the size of the wavefunction grows exponentially with the system size  $n$ , a direct computation quickly becomes intractable as the system size increases. The goal of NNQS is to design a compact architecture that can approximate and optimize the wavefunction efficiently.
3.2 Quantum State Learning
Given a quantum state  $|\phi \rangle$ , we are often interested in finding  $|\psi_{\theta}\rangle$  that is closest to  $|\phi \rangle$  given the variational ansatz. In quantum mechanics, the closeness of two quantum states is measured by the quantum fidelity  $F = |\langle \phi |\psi_{\theta}\rangle |^2 = |\sum_x\phi^* (\pmb {x})\psi_\theta (\pmb {x})|^2$  where  $^{*}$  refers to complex conjugation. Quantum fidelity satisfies  $0\leq F\leq 1$ , with  $F = 1$  corresponding to exact match, and  $F = 0$  orthogonal quantum states. Finding  $F_{\mathrm{max}}$  for a given ansatz allows us to quantify how good the ansatz can be used to approximate that state. In practice, minimizing  $-\log F$  is usually better than maximizing  $F$  itself. This can be achieved by enumerating all the basis  $x$  in small systems and stochastically in large systems such as quantum state tomography (see Appendix A.1).
3.3 Variational Monte Carlo
For a given quantum system with  $n$  qubits, the Hamiltonian  $\hat{\mathcal{H}}$  can be written as a Hermitian matrix of size  $2^{n} \times 2^{n}$ . The ground state energy  $E_{g}$  of the system is the smallest eigenvalue of  $\hat{\mathcal{H}}$  and the ground state is the corresponding eigenvector. For large system sizes, finding the ground state directly is usually impossible. In this case, the variational principle in quantum mechanics provides an upper bound on the ground state energy  $E_{g}$ . For all (normalized) wavefunctions  $|\psi\rangle$ , it is evident that  $E_{g} \leq \langle \psi | \hat{\mathcal{H}} | \psi \rangle$ . Therefore, finding the  $|\psi_{\theta}\rangle$  that minimizes  $E_{\theta} = \langle \psi_{\theta} | \hat{\mathcal{H}} | \psi_{\theta} \rangle$  gives the lowest upper bound of the ground state energy. For large system sizes, we can stochastically evaluate and minimize
$$
\langle \psi_ {\theta} | \hat {\mathcal {H}} | \psi_ {\theta} \rangle = \sum_ {\boldsymbol {x x ^ {\prime}}} | \psi_ {\theta} (\boldsymbol {x}) | ^ {2} \frac {\mathcal {H} _ {\boldsymbol {x} , \boldsymbol {x} ^ {\prime}} \psi_ {\theta} \left(\boldsymbol {x} ^ {\prime}\right)}{\psi_ {\theta} (\boldsymbol {x})} = \mathbb {E} _ {\boldsymbol {x} \sim | \psi_ {\theta} | ^ {2}} \frac {\sum_ {\boldsymbol {x} ^ {\prime}} \mathcal {H} _ {\boldsymbol {x} , \boldsymbol {x} ^ {\prime}} \psi_ {\theta} \left(\boldsymbol {x} ^ {\prime}\right)}{\psi_ {\theta} (\boldsymbol {x})}, \tag {1}
$$
where  $\mathcal{H}_{\boldsymbol{x},\boldsymbol{x}'}$  refers to the matrix element of  $\hat{\mathcal{H}}$  and we interpret  $|\psi_{\theta}(\boldsymbol{x})|^2$  as a probability distribution. The summation over  $\boldsymbol{x}'$  can be efficiently computed given  $\boldsymbol{x}$  since the Hamiltonian is usually sparse. The gradient  $\nabla_{\theta}E_{\theta}$  can also be calculated stochastically in a similar fashion (see Appendix A.2).
4.1 Preliminaries
Autoregressive Neural Network Wavefunction. The autoregressive neural network (ARNN) (Fig. 1 left) parameterizes the full probability distribution as a product of conditional probability distributions
$$
p (\boldsymbol {x}) = \prod_ {i = 1} ^ {n} p \left(x _ {i} \mid \boldsymbol {x} _ {<   i}\right), \tag {2}
$$
where  $\pmb{x} = (x_{1},\dots,x_{n})$  is a configuration of  $n$  qubits and  $\pmb{x}_{< i} = (x_1,\dots,x_{i - 1})$  is any configuration before  $x_{i}$ . The normalization of the full probability can be guaranteed from the normalization of individual conditional probabilities. The ARNN also allows for exact sampling from the full distribution by sampling sequentially from the conditional probabilities (see Appendix B.1). The autoregressive constructions for probabilities can be easily modified to represent quantum wavefunction by (1) replacing  $p(x_{i}|x_{< i})$  with a complex-valued conditional wavefunction  $\psi (x_i|x_{< i})$  and (2) using the following normalization condition for the conditional wavefunctions:  $\sum_{x}|\psi (x_{i}|x_{< i})|^{2} = 1$  (Sharir et al., 2020; Hibat-Allah et al., 2020; Luo et al., 2021b). Similar to the case of probabilities, ARNN automatically preserves the normalization of wavefunctions and allows for exact sampling (see Appendix B.1). Because of this, ARNN is often more efficient when training and computing various quantum observables compared to other generative neural networks.
Matrix Product State (MPS). MPS (also known as tensor train) is a widely used TN architecture to study quantum many-body physics (Vidal, 2003, 2004). The MPS defines a wavefunction using  $n$  rank-3 tensors  $M_{x_i}^{\alpha_i - 1\alpha_i}$  with  $\alpha_{i - 1}$  (or  $\alpha_{i}$ ) the index of the left (or right) bond dimensions and  $x_{i}$  the configuration of the  $i$ th qubit. Then, the full wavefunction is generated by first choosing a particular configuration  $\pmb{x} = (x_{1},\dots,x_{n})$  and then contracting the tensors selected by this configuration (Fig. 1 right) as
$$
\psi (\boldsymbol {x}) = \sum_ {\alpha_ {1}, \dots , \alpha_ {n - 1}} M _ {x _ {1}} ^ {\alpha_ {0} \alpha_ {1}} \dots M _ {x _ {n}} ^ {\alpha_ {n - 1} \alpha_ {n}} = \sum_ {\boldsymbol {\alpha}} \prod_ {i = 1} ^ {n} M _ {x _ {i}} ^ {\alpha_ {i - 1} \alpha_ {i}}, \tag {3}
$$
where the left-and-right-most bond dimensions are assumed to be  $\mathcal{D}(\alpha_0) = \mathcal{D}(\alpha_{n + 1}) = 1$  so are not summed.
4.2 Autoregressive Neural TensorNet (ANTN)
Elementwise and Blockwise Construction for ANTN. Both ARNN and MPS are powerful ansatzes for parameterizing quantum many-body wavefunctions. Albeit very expressive, ARNN lacks the physics prior of the system of interest. In addition, since wavefunctions are complex-valued in general, learning the sign structure can be a nontrivial task (Westerhout et al., 2020). MPS, on the other hand, contains the necessary physics prior and can efficiently represent local or quasi-local sign structures, but its expressivity is severely limited. The internal bond dimension needs to grow exponentially to account for a linear increase in entanglement. It turns out that MPS representation is not unique in that many MPS actually represent the same wavefunction; and if we choose a particular constraint (without affecting the expressivity), MPS allows for efficient evaluation of conditional probability and exact sampling (Ferris & Vidal, 2012) (see Appendix B.2). Because both these ansatzes allow for exact sampling, it is natural to combine them to produce a more powerful ansatz. Therefore, we develop the Autoregressive Neural TensorNet (ANTN) (Fig. 1 (middle)). In the last layer of the ANTN, instead of outputting the conditional wavefunction  $\psi(x_i | x_{<i})$ , we output a conditional wavefunction tensor  $\tilde{\psi}^{\alpha_{i-1}\alpha_i}(x_i | x_{<i})$  for each site. Defining the left partially contracted tensor up to qubit  $j$  as  $\tilde{\psi}_L^{\alpha_j}(x_{\leq j}) \coloneqq \sum_{\alpha < j} \prod_{i=1}^{j} \tilde{\psi}^{\alpha_{i-1}\alpha_i}(x_i | x_{<i})$ , we can define the (unnormized) marginal probability distribution as
$$
q \left(\boldsymbol {x} _ {\leq j}\right) := \sum_ {\alpha_ {j}} \tilde {\psi} _ {L} ^ {\alpha_ {j}} \left(\boldsymbol {x} _ {\leq j}\right) \tilde {\psi} _ {L} ^ {\alpha_ {j}} \left(\boldsymbol {x} _ {\leq j}\right) ^ {*}, \tag {4}
$$
where  $^*$  denotes complex conjugation. Then, the (normalized) conditional probability can be obtained as  $q(x_{j}|\pmb{x}_{< j}) = q(\pmb{x}_{\leq j}) / \sum_{x_{j}}q(\pmb{x}_{\leq j})$  . We construct the overall wavefunction by defining both its amplitude and phase according to
$$
\psi (\boldsymbol {x}) := \sqrt {q (\boldsymbol {x})} e ^ {i \phi (\boldsymbol {x})}, \tag {5}
$$
with  $q(\pmb{x}) \eqqcolon \prod_{i=1}^{n} q(x_i | \pmb{x}_{<i})$  and the phase  $\phi(\pmb{x}) \eqqcolon \operatorname{Arg} \sum_{\alpha} \prod_{i=1}^{n} \tilde{\psi}^{\alpha_{i-1}, \alpha_i}(x_i | \pmb{x}_{<i})$ . In other words, we define the amplitude of the wavefunction through the conditional probability distributions and define the phase analogous to the standard MPS.
We develop two different constructions for ANTN that differ in the last layer on how to construct conditional wavefunction tensors.
The elementwise ANTN is given by
$$
\tilde {\psi} ^ {\alpha_ {i - 1} \alpha_ {i}} \left(x _ {i} \mid \boldsymbol {x} _ {<   i}\right) = M _ {x _ {i}} ^ {\alpha_ {i - 1} \alpha_ {i}} + f _ {N N} \left(x _ {i}, \alpha_ {i - 1}, \alpha_ {i} \mid \boldsymbol {x} _ {<   i}\right), \tag {6}
$$
where  $f_{NN}(x_i, \alpha_{i-1}, \alpha_i | x_{<i})$  is the complex-valued output. The blockwise ANTN is given by
$$
\tilde {\psi} ^ {\alpha_ {i - 1} \alpha_ {i}} \left(x _ {i} \mid \boldsymbol {x} _ {<   i}\right) = M _ {x _ {i}} ^ {\alpha_ {i - 1} \alpha_ {i}} + f _ {N N} \left(x _ {i} \mid \boldsymbol {x} _ {<   i}\right), \tag {7}
$$
where the complex-valued output  $f_{NN}(x_i | \boldsymbol{x}_{<i})$  is broadcasted over  $\alpha$ . This results in a lower complexity that allows us to use a larger maximum bond dimension.
Transformer and PixelCNN Based ANTN. Our construction above is general and can be applied to any standard ARNN. Depending on the application, we can use different ARNN architectures. In this work, we choose the transformer and PixelCNN depending on the specific tasks. The transformer (Vaswani et al., 2017) used here is similar to a decoder-only transformer implemented in (Luo et al., 2020), and The PixelCNN we use is the gated PixelCNN (Van den Oord et al., 2016) implemented in (Chen et al., 2022).
MPS Initialization. Since our ANTN generalizes from TN, both the elementwise and the blockwise can take advantage of the optimized MPS from algorithms such as DMRG (White, 1992) as an initialization. In practice, we can initialize the TN component with the optimized DMRG results of the same bond dimension (similar to (Lami et al., 2022; Wu et al., 2022)). The MPS Initialization can also be thought of as a pretraining of the ANTN. This is a nice feature since it provides a good sign structure and inductive bias from the physics structure, which does not exist in the conventional ARNN.
Limitations. In this work, we only integrated MPS into the ANTN, where MPS may not be the best TN in various settings. Besides MPS, many other TNs also allow efficient evaluation of conditional probabilities and exact sampling, such as MERA (Vidal, 2007) and PEPS (Verstraete & Cirac, 2004), or cylindrical MPS for periodic boundary conditions. In fact, the recently developed TensorRNN (Wu et al., 2022) can also be viewed as a type of TN and can be integrated into our construction for future work. In addition, while the ANTN construction is general in terms of the base ARNN used, our current study only focuses on transformer and PixelCNN. Lastly, as shown later in Sec. 5.1, the current implementation of ANTN has an additional sampling overhead that is linear in system size which can be avoided.
5 Theoretical Results
5.1 Exact Sampling and Complexity Analysis of ANTN
Theorem 5.1. Autoregressive Neural TensorNet wavefunction is automatically normalized and allows for exact sampling.
Proof. This is a direct consequence that we defined the amplitude of the wavefunction through normalized conditional probability distributions  $q(x_{i}|\pmb{x}_{< i})$ . (See Appendix B.1 for the detailed sampling procedure.)
Complexity Analysis. We first note that for MPS, the number of parameters and computational complexity for evaluating a bitstring  $\pmb{x}$  scales as  $\mathcal{O}(n\chi^2)$ , where  $n$  is the number of particles and  $\chi = \mathcal{D}(\alpha)$  is the (maximum) bond dimension of MPS. The sampling complexity scales as  $\mathcal{O}(n^2\chi^2)$ . The DMRG algorithm has a computational complexity of  $\mathcal{O}(n\chi^3)$  (White, 1992). The number of parameters and computational complexity of ARNN depends on the specific choice. Assuming the ARNN has  $n_{\mathrm{ARNN}}$  parameters and a single pass (evaluation) of the ARNN has a computational complexity of  $c_{\mathrm{ARNN}}$ , then the sampling complexity of the ARNN is  $\mathcal{O}(nc_{\mathrm{ARNN}})$  in our current implementation. The ANTN based on such ARNN would have  $\mathcal{O}(n_{\mathrm{ARNN}} + n\chi^2 h_{\mathrm{dim}})$  parameters for the elementwise construction, and  $\mathcal{O}(n_{\mathrm{ARNN}} + n\chi^2 + nh_{\mathrm{dim}})$  parameters for the blockwise construction. The computational complexity for evaluating and sampling bitstrings scales as  $\mathcal{O}(n^s (c_{\mathrm{ARNN}} + n\chi^2 h_{\mathrm{dim}}))$  for elementwise construction and  $\mathcal{O}(n^s (c_{\mathrm{ARNN}} + n\chi^2))$  for blockwise construction with  $s = 0$  for evaluation and  $s = 1$  for sampling.
We note that the complexity analysis above is based on our current implementation, where the sampling procedure for all algorithms has an  $\mathcal{O}(n)$  overhead compared to evaluation. It is possible to remove it by storing the partial results during the sampling procedure.
Then, each gradient update of the variational Monte Carlo algorithm requires sampling once and evaluating  $N_{\mathrm{con}}$  times where  $N_{\mathrm{con}}$  is the number of connected (non-zero)  $\pmb{x}^{\prime}$  s in  $\mathcal{H}_{\pmb{x},\pmb{x}^{\prime}}$  given  $\pmb{x}$ , which usually scales linearly with the system size. Then, the overall complexity is  $\mathcal{O}(N_s(N_{\mathrm{con}}c_{\mathrm{eva}} + c_{\mathrm{samp}}))$  with  $N_{s}$  the batch size,  $c_{\mathrm{eva}}$  the evaluation cost and  $c_{\mathrm{samp}}$  the sampling cost. Usually, the first part dominates the cost.
The difference in computational complexities and number of parameters between elementwise ANTN and blockwise ANTN implies that blockwise ANTN is usually more economical than elementwise ANTN for the same hidden dimsnieon  $h_{\mathrm{dim}}$  and bond dimension  $\chi$ . Therefore, for small bond dimensions, we use the elementwise ANTN for a more flexible parameterization with a higher cost. In contrast, for large bond dimensions, we use the blockwise ANTN, which saves computational complexity and improves the initial performance (with DMRG initialization) of the blockwise ANTN at the cost of less flexible modifications from the ARNN. We also note that compared to the state-of-the-art MPS simulation, even for our blockwise ANTN, a small bond dimension usually suffices. For example, in the later experimental section, we use bond dimension 70 for blockwise ANTN while the best MPS results use bond dimension 1024, which has much more parameters than our ANTN. In general, our ANTN has much fewer parameters compared to MPS due to the  $\mathcal{O}(\chi^2)$  parameters scaling which dominates at large bond dimension.
5.2 Expressivity Results of ANTN
Theorem 5.2. Autoregressive Neural TensorNet can have volume law entanglement, which is strictly beyond the expressivity of matrix product states.
Proof. We proved in Thm. 5.3 that ANTN can be reduced to an ARNN, which has been shown to have volume law entanglement that cannot be efficiently represented by TN (Sharir et al., 2021; Levine et al., 2019b). Hence, ANTN can represent states that cannot in general be represented by MPS efficiently; thus it has strictly greater expressivity.
Theorem 5.3. Autoregressive Neural TensorNet has generalized expressivity over tensor networks and autoregressive neural networks.