Predicting quantum many-body dynamics with transferable neural networks

TL;DR

This paper introduces a recurrent neural network framework that efficiently predicts the dynamics of 1D quantum many-body systems, demonstrating transferability and accuracy without detailed Hamiltonian knowledge.

Contribution

It presents a simple SRU-based transfer learning approach capable of predicting quantum dynamics from a single initial state, reducing computational costs and requiring minimal system information.

Findings

Accurately predicts 1D Ising model dynamics

Demonstrates transferability to larger systems

Achieves predictions with constant computational complexity

Abstract

Machine learning (ML) architectures such as convolutional neural networks (CNNs) have garnered considerable recent attention in the study of quantum many-body systems. However, advanced ML approaches such as transfer learning have seldom been applied to such contexts. Here we demonstrate that a simple recurrent unit (SRU) based efficient and transferable sequence learning framework is capable of learning and accurately predicting the time evolution of one-dimensional (1D) Ising model with simultaneous transverse and parallel magnetic fields, as quantitatively corroborated by relative entropy measurements and magnetization between the predicted and exact state distributions. At a cost of constant computational complexity, a larger many-body state evolution was predicted in an autoregressive way from just one initial state, without any guidance or knowledge of any Hamiltonian. Our work paves the way for future applications of advanced ML methods in quantum many-body dynamics only with knowledge from a smaller system.

Tables

Table 1. Table 1: Comparsion of time consumption (seconds) between ED and SRU-based methods. Three independent runs are performed to generate sequences of different batch sizes and the average time consumption is reported. B ​ S 𝐵 𝑆 BS denotes the batch size, and the bold black font indicates that our SRU-based method is superior to ED with increasing spin number and batch size.

	$BS=1$		$BS=64$		$BS=128$		$BS=256$
System	ED	Ours	ED	Ours	ED	Ours	ED	Ours
2-spin	0.015	$0.425$	1.1	$0.74$	2.3	$0.83$	$4.6$	$0.69$
3-spin	0.035	$0.425$	2.2	$0.74$	4.4	$0.69$	$8.7$	$0.71$
4-spin	0.059	$0.425$	3.8	$0.77$	7.6	$0.66$	$15.1$	$0.72$
5-spin	0.271	$0.425$	17.5	$0.82$	34.9	$0.76$	$69.7$	$0.79$
6-spin	0.556	$0.425$	35.2	$0.71$	70.5	$0.66$	$141.4$	$0.98$
7-spin	1.15	$0.425$	73.1	$0.80$	146.3	$0.91$	$292.5$	$2.18$

Equations

$$H=-\sum _i^{N-1}{\sigma }_i^z\otimes {\sigma }_{i+1}^z-h\sum _i^N{\sigma }_i^z-g\sum _i^N{\sigma }_i^x.$$

$$\Psi =\prod _i^N\otimes {\psi }_i=\prod _i^N\otimes \left(\begin{array}{c}{\varphi }_i^{\uparrow }\\ {\varphi }_i^{\downarrow }\end{array}\right),\dim \Psi =2^N.$$

$$|\Psi (t)⟩=\exp \text{quantity}(-\mathrm{i}\frac{H}{\hslash }t)|\Psi (0)⟩=V\exp \text{quantity}(-\mathrm{i}\frac{E}{\hslash }t)V^{-1}|\Psi (0)⟩,$$

$$D_{KL}(P||Q)=\sum _{x}P(x)\log (\frac{P(x)}{Q(x)}).$$

$$P_{n,t,x}=\frac{|M_{n,t,x}^{\text{ED}}|}{\sum _{x=1}^{2^N}|M_{n,t,x}^{\text{ED}}|}$$

$$Q_{n,t,x}=\frac{|M_{n,t,x}^{\text{SRU}}|}{\sum _{x=1}^{2^N}|M_{n,t,x}^{\text{SRU}}|}$$

$$D_{KL}(P||Q)(t)=\frac{1}{1000}\sum _{n=1}^{1000}\sum _{x=1}^{2^N}{P}_{n,t,x}\log \frac{P_{n,t,x}}{Q_{n,t,x}},$$

Full text

I Introduction

Machine learning (ML) approaches, particularly neural networks (NNs), have achieved great success in solving real-world industrial and social problems 1, such as image recognition2, high level image synthesis and style transfer3, human-like raw speech generator4, producing original melodious MIDI notes5, neural machine translation 6. Inspired by its widespread applicability, ML was soon adopted by condensed matter physicists in the modeling of quantum many-body behavior and phase transition discovery 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17. Compared to so many advances in computer vision18, speech processing19, and natural language processing 20, it is natural to ask if recent progress in these more sophisticated ML architectures can benefit or even revolutionize the modeling of quantum systems. For instance, can quantum many-body dynamics be “learned” through transferable learning21, 22.

Thus, the main objective of this work is to demonstrate the novel application of NNs in the transferable learning and prediction of the evolution of a many-body wavefunction, an otherwise computationally intensive task that has not been solved by generative models. Focusing on static problems, it is proven that deep NNs like restricted Boltzmann Machine (RBM) can represent most physical states24, and a recent work based on very deep and large CNNs shows the ability to circumvent the need for Markov Chain sampling on two-dimensional interacting spin model of larger systems25. Lately, physical properties of spin Hamiltonians are reproduced by deep Boltzmann Machine (DBM), as an alternative to the standard path integral26. Our approach is fundamentally in contrast with conventional approaches in computing many-body dynamics: instead of evolving the wavefunction explicitly with the Hamiltonian, which becomes prohibitively slow and impractical as the number of spin variables increases, we directly predict the dynamical wavefunction from the initial state by propagating it with an efficient and transferable framework based on unified spin encoding, chain encoding and SRU27 module. With the same level of parallelism as feed-forward CNNs and scalable context-dependent capacity of recurrent connections, our proposed framework are naturally suited for learning many-body systems with unified parameters, although they have never been harnessed for exact quantum state evolution, in our scenario, a 1D Ising model with both parallel and transverse magnetic field.

Inspired by end-to-end training28 and domain adaptation29, 30, we specialize to the many-body dynamics of a 1D Ising chain with transverse and parallel magnetic fields. Comparison with exact conventionally computed results with up to seven spins reveals high predictive accuracy, as quantified by the relative entropy as well as magnetization. Indeed, our SRU-propagated wavefunction shows a strong grasp of the periodicity in the time evolution, despite being unaware of the Hamiltonian that sets the energy (inverse periodicity) scale. Encouraged by circumventing the problem of exponential computational complexity through unified encoding mechanisms and parallel recurrent connections, we hope that such encouraging results from our pioneering transferable learning appoach will inspire further applications of transferable learning methods to build a shared model suited for quantum systems with vast spin variables.

II Dynamics on a 1D Ising chain

We consider a 1D Ising spin chain composed of $N$ spin variables with local transverse ($g$) and parallel ($h$) magnetic fields, described by the Hamiltonian

[TABLE]

where $\sigma^{x}$ and $\sigma^{z}$ denotes the Pauli matrices, and $i$ denotes the spin variable. When the magnetic field is parallel ($g=0$) or transverse ($h=0$), the Hamiltonian is exactly solvable. However, when $g\neq 0$ and $h\neq 0$, the dynamics of $N$ spins must be numerically computed in the $2^{N}$-dimensional many-body Hilbert space spanned by direct product states $\Psi$ of single-spin wavefunctions $\psi_{i}$:

[TABLE]

Wavefunction dynamics can be exactly computed through unitary time evolution of the Hamiltonian

[TABLE]

where $E=V^{-1}HV$ is the diagonal eigenenergy matrix.

The $2^{N}$-dimensional $N$-body wave function $\ket{\Psi(t)}$ quickly becomes expensive to compute as $N$ increases. We propose a ML approach with spin encoding layer, chain encoding layer, SRU layers, and spin decoding layer, which instead attemps to predict its time evolution based on prior knowledge of the time evolution behavior of known training states. This training (learning) only has to be performed once for the relatively inexpensive prediction of any number of initial states. Importantly, the training and prediction process captures solely the intrinsic evolution patterns of the wavefunctions, and does not involve any explicit knowledge about the Hamiltonian. From the ML perspective, this dynamical state evolution problem can be regarded as a straightforward sequence generation problem 31. Moreover, as we shall explain, our SRU-based framework is transferable.

III The transferable NN approach

We next outline the broad principles behind our NN approach of predicting quantum state evolution, with details in 32. Here we choose a NN composed of a spin encoding layer, a chain encoding layer, SRU layers and a spin decoding layer (Fig. 1b). The vanilla SRU NN with peephole connections (Fig.1a) substitutes inherent matrix multiplication with parallelizable element-wise multiplication operations ($\odot$ in Fig. 1a) associated with $c_{t-1}$, hence the calculation of $f_{t}$ doesn’t have to wait until the whole $c_{t-1}$ is updated. With the help of spin encoding and decoding layers, the amount of trained parameters is fixed, and thus the complexity has an upper bound instead of increase exponentially.

Our procedure occurs in two main stages: the training stage and the inference stage. In the training stage, we first “train” or optimize the weight parameters of our SRU-based framework by feeding it with a large number of training sequences, which are the time-evolved wavefunction data of $10^{4}$ randomly chosen initial $2$-spin to $7$-spin state sequences sampled over 500 timesteps, obtained via conventional exact diagonalization (ED). The SRU-based framework is fully optimized by Adam optimization algorithm 33 to minimize the mean squared error between the ED-evolved and SRU-evolved states at all timesteps in a mini-batch 32.

Following the training stage is the inference stage, when the SRU-based framework is ready for predicting the evolution of arbitrarily given initial states. As sketched in Fig. 1d, the initial many-body state $\ket{\Psi(t=0)}$ enters the leftmost block at $t=0$, then processed by a spin encoding layer, a chain encoding layer, and two fully-connected layers, and its output is propagated as input state to the next block with hidden layers $h_{t}$. The output of each block denotes a new quantum state at a certain timestep. The combination of memory cell $c_{t}$ and hidden output $h_{t}$ serves to implement effective context-dependent behaviors. As illustrated in Fig. 1(a) and further elaborated in 32, context-dependent information kept in memory cell $c_{t}$ is modified by its previous value $c_{t-1}$, new input $x_{t}$ interacted with forget gate and skip gate at that timestep, as well as “hidden” information on $h_{t-1}$ from the previous SRU cell. Based on the already optimized SRU-based framework, the final predicted quantum state as a function of time would be generated from one fully-connected layer and the spin decoding layer as shown in Fig. 1b.

IV Comparison between exact and SRU-based evolutions

We report very encouraging agreements between wavefunctions evolved by $e^{-iHt/\hbar}$ as computed by ED, and wavefunction evolutions as predicted by our SRU-based framework. As for the 1D Ising model, we set the local transverse magnetic field $g$ to be $-1.05$, parallel magnetic field $h$ to be $0.5$ and $\Delta t$, the time interval to be $0.002$, and keep this setting constant for all computation. We find that the maximum energy eigenvalue is about $0.1\gg 0.002$, proving that the time interval we choose is small enough. The number of spin variables studied ($2$ to $7$) decides the cost of exactly computing the $10^{4}$ different time evolutions over 0.2 second (100 timesteps) prior to training the network, since the time complexity of ED method is $\mathcal{O}(2^{n})$. The training and inference loss of different systems is shown in Fig. 2.

As a concrete demonstration, we visually illustrate the comparison for the evolution of a typical state from 2-spin to 7-spin in Fig. 3. These states are evolved from arbitrarily chosen initial states from the test set. Saliently, the evolution predicted by the SRU-based model accurately reproduces that from exact computations at the beginning $100$ timesteps. To confirm that this agreement is not just due to a fortuitous choice of component, we look at the evolution across all components of the same states in Fig. 4.

To further quantify the agreement of SRU and ED wavefunction evolutions, we compute the relative entropy (Kullback–Leibler divergence)34 of their distributions over 1000 test wavefunctions sequences. For discrete probability distributions $P$ and $Q$, the relative entropy is defined as

[TABLE]

Given ED-computed wavefunction coefficient vectors $M^{\textrm{ED}}$ and SRU-predicted coefficient vectors $M^{\textrm{SRU}}$, the $P$ and $Q$ variables take values

[TABLE]

at time $t$ and basis vector $x$, where $n$ labels the test sequence. Hence the mean relative entropy (MRE) at each timestep $t$ is

[TABLE]

and measures the amount of information lost when the distribution $Q$ from SRU predictions is used to represent the distribution $P$ from ED results. The smaller the value of $D_{KL}(P||Q)(t)$, the more accurate is their agreement.

In Fig. 5, we show how the MRE varies with time during the generation of test sequences. We find that in all six systems, the order of relative entropy is always within $0.03$. Evidently, with the increase of timesteps, the relative entropy generally shows an upward trend and increases linearly with timesteps (see 32), which is caused by the accumulation of errors in the process of conditional generation without any external guidance, though already suppressed dropout layers. To quantify our model’s performance by a physical variable, we draw the magnetization intensity calculated from both predicted (SRU) and simulated (ED) wavefunctions in Fig. 6, which have a nice agreement. Specifically, for smaller-sized systems, such as 2-spin and 3-spin, the predicted magnetization intensity has a very nice agreement with simulated one. With the increase of spin variables, our SRU-based framework has a performance drop due to the exponentially increased computation complexity. Meanwhile, with the increase of timesteps, the difference between predicted and simulated magnetization intensity also becomes larger, which is due to the error accumulation during the autoregressive generation without any external guidance.

Owing to its unified encoding and parallelism, our SRU-based NN is becoming increasingly more superior over the ED method in terms of efficiency, as the number of spins and batch size increase. Table 1 summarizes the results. When the number of spins gets larger, e.g. $6$ and $7$, the advantage of our SRU-based framework on inference speed becomes more and more obvious, that, is attributed to its constant computational complexity. In addition, when we enlarge the batch size to $256$ for $7$ spins, our model demonstrates a speed $130$ times faster than the ED-based method.

After obtaining base model trained with datasets of $2$ to $7$ spins by $300$ epoches, we may continue to finetune it with the dataset of the $8$-spin system. To make a comparison, we also train it from scratch. The results in Fig. 7(a) shows that the validation loss by finetuning base model is much lower than training from scratch, demonstrating that our NN has already learned transferable features from smaller systems. The MRE of 8-spin system is shown in Fig. 7(b).

V Conclusion

In this work, we have successfully applied a transferable NN approach based on SRU networks to approximate the state evolution of dynamic quantum many-body systems with high accuracy and superior scalability. Our work encourages future applications of advanced ML methods in quantum many-body dynamics in a Hamiltonian-agnostic manner. One possibility is to predict the behavior of large and inhomogeneous systems lack of traing data by just learning from a smaller-sized system[33]. Applications of these advancements in ML to quantum many-body problems are left to future work.

Acknowledgements

Xiao Zhang thanks Yingfei Gu, Meng Cheng, Yi Zhang for discussions. Xiao Zhang is supported by the National Natural Science Foundation of China (Grant No. 11874431), the National Key R & D Program of China (Grant No. 2018YFA0306800) and the Guangdong Science and Technology Innovation Youth Talent Program (Grant No. 2016TQ03X688). Shuo Yang is supported by NSFC (Grant No. 11804181), the National Key R & D Program of China (Grant No. 2018YFA0306504) and the Research Fund Program of the State Key Laboratory of Low-Dimensional Quantum Physics (Grant No. ZZ201803).