Systematic improvement of neural network quantum states using Lanczos
Abstract
The quantum many-body problem lies at the center of the most important openchallenges in condensed matter, quantum chemistry, atomic, nuclear, and highenergy physics.
While quantum Monte Carlo, when applicable, remains the mostpowerful numerical technique capable of treating dozens or hundreds of degrees offreedom with high accuracy, it is restricted to models that are not afﬂicted by theinfamous sign problem.
A powerful alternative that has emerged in recent yearsis the use of neural networks as variational estimators for quantum states.
In thiswork, we propose a symmetry-projected variational solution in the form of linearcombinations of simple restricted Boltzmann machines.
This construction allowsone to explore states outside of the original variational manifold and increase therepresentation power with moderate computational effort.
Besides allowing one torestore spatial symmetries, an expansion in terms of Krylov states using a Lanczosrecursion offers a solution that can further improve the quantum state accuracy.
Weillustrate these ideas with an application to the Heisenberg J 1 − J 2 model on thesquare lattice, a paradigmatic problem under debate in condensed matter physics,and achieve state-of-the-art accuracy in the representation of the ground state.
Introduction
Understanding correlated quantum systems requires dealing with a large conﬁguration space: datasetsare comprised of all possible electronic conﬁgurations (cid:126)σ and cannot be stored in the memory of the largest supercomputer.
Hence, the quantum many-body problem can be interpreted as an “extremedata science” problem [13] from an information processing perspective.
In a quantum wave function,each electronic or spin conﬁguration has an associated complex amplitude ψ ( (cid:126)σ ) determined bysolving for the eigenvectors of the Hamiltonian operator.
In particular, if one is interested in thezero temperature properties of the system, the solution is given by the eigenvector with the smallesteigenvalue.
Finding the exact solution of a N quantum bit system with interactions requires solvingfor the eigenvectors of a 2 N × 2 N matrix.
Alternatively, one can formulate the calculation as anoptimization problem in which an “energy functional” E ( ψ ) has to minimized with respect to all the2 N complex amplitudes.
Since the number of conﬁgurations d grows exponentially with the number of degrees of freedom(electrons, spins), this problem quickly becomes intractable.
A solution consists of “compressing’ thewave function by proposing a suitable guess for the amplitudes based on some variational parameters(cid:126)α =
( α 1 , α 2 , · · · , α m ) .
Typically, a functional form ψ ( (cid:126)σ ) =
f ( (cid:126)σ, (cid:126)α ) based on some physicalintuition is utilized to represent the amplitude of given conﬁguration/state (cid:126)σ .
The optimal parametersα i are determined by solving the system of equations ∇ α E = 0 .
The objective of this solution is toachieve the lowest possible energy with a number of parameters m (cid:28) d .
Some relatively simple wave functions have enjoyed various degrees of success in the past, suchas those of the Jastrow type where the amplitudes can we written as pair products
f ( (cid:126)σ, (cid:126)α ) = 37 (cid:81) ij U ( α ij σ i σ j ) .
However, in recent years we have witnessed impressive developments based on theuse of neural network (NN) wave functions as variational estimators [4], which have jump started ofa new vibrant ﬁeld of research dubbed as “quantum machine learning”.
Notice that the optimizationof the wave function parameters now translates into the “training” of the NN by minimizing theenergy functional that becomes a cost function (we describe the training process below).
The powerof NN wave functions lies on their complex non-linear structure that provides them with remarkableexpressivity to represent arbitrary complex many-by states by, at the same time, being completelyagnostic to the physics.
Since restricted Boltzman machines(RBM) were originally used as a variational ansatz for ﬁnding theground state of the quantum many body systems [4], there has been a growing effort to investigateother forms of neural networks, including convolutional neural networks(CNN)[9, 23], recurrentneural networks(RNN)[19], graph networks[22], transformers[25], to mention a few.
Thus, neuralnetwork quantum states(NNQS) become the most appealing numerical alternative to treat quantummany body systems since they can be systematically improved by adding new layers or hiddenvariables, for instance.
In addition to the ground state search, the application of NNQS rangesfrom classical simulation of quantum circuits[1, 5, 32], calculation of spectral function[17, 18],thermodynamics simulation[16, 31], and quantum tomography[44].
Contributions In this work, we show how one can use a mathematically simple structure, arestricted Boltzmann machine (RBM), and yet obtain values of the ground state energy that beat allprevious estimates by a range a numerical methods, including using convolutional neural networks.
As we describe below, instead of increasing the number of layers or hidden variables, the solutionlies on considering linear combinations of RBMs.
The new wave function allows one to explorea much larger space of solutions.
In particular, one can use this construction to restore spatialsymmetries [40, 9, 28, 29].
In addition, we propose implementing a projection method based on a
Lanczos recursion using a “Krylov basis” of RBMs obtained by sequentially applying powers of the
Hamiltonian operator.
The paper is organized as follows:
In Sec.2.1 we describe the quantum many-body problem in thecontext of the Heisenberg model; in Sec.2.2 we summarize prior attempts to study this problemusing NNQS; in Sec.3 we review the basic formalism, including the structure of neural network wavefunctions, how to incorporate the symmetries of the problem into the quantum many-body state,and the numerical training procedure to optimize it.
In Sec.4 we present results of state of the artcalculations for the J 1 − J 2 Heisenberg model on the square lattice and compare to other numericaltechniques.
We ﬁnally close with a summary and conclusions.
The quantum many-body problem2.1 Model
In the following, we will focus on quantum spin problems where the degrees of freedom σ i canassume two possible values ± 1 / 2 (or “up” and “down”).
Similarly, one can think of them as generictwo-level systems or “qubits”.
In particular, we will benchmark our methods in the context of thespin 12 antiferromagnetic Heisenberg model with nearest and next nearest neighbor interactions, theso-called J 1 − J 2 model deﬁned by the Hamiltonian:
[Equation]
where (cid:126)S =
( ˆ S x , ˆ S y , ˆ S z ) are spin operators, the ﬁrst term runs over nearest neighboring sites (cid:104) ij (cid:105) ona square lattice and the second term runs over next nearest pairs (cid:104)(cid:104) ij (cid:105)(cid:105) along the diagonals of theplaquettes.
For convenience, in the following we set J 1 = 1 as the unit of energy.
In this problem,the number of possible conﬁgurations grows as d = 2 N .
However, the ground state wave functionlies on the sector with the same number of up and down spins, constraining our search to a smallersubset of states, albeit still exponentially large.
Without the J 2 term, the problem can be numerically solved for hundreds of spins using quantum
Monte Carlo (QMC) [38].
However, the method cannot be applied to problems with frustration since it is noticeably affected by the infamous sign problem[24].
In our case, this is due to the presence ofthe J 2 term that makes some transition probabilities ill deﬁned (negative).
The ground states of thismodel are well established in two extreme cases: at small J 2 /J 1 the system antiferromagneticallyorders with wave vector q = ( π, π ) ; at large J 2 /J 1 spins prefer columnar order q
= ( π, 0) , (0 , π ) ,in which they aligned antiparallel in one direction, but ferromagnetically in the other.
However, inthe maximally frustrated regime J 1 ∼ 0 .
5 J 2 , the system does not display any apparent order and thenature of this spin liquid state remains controversial despite signiﬁcant research effort over the pastthree decades[3, 6, 11, 10, 35, 39, 41, 34, 37, 26, 20, 15, 21, 45].
Therefore, we choose this Hamiltonian for two reasons: (i) it realizes a quantum spin liquid in aparameter regime near J 2 ∼ 0 .
5 J 1 and (ii) conventional Monte Carlo methods fail, making the modelan ideal testing ground to benchmark new techniques.
Variational Monte Carlo(VMC) provides asuitable alternative that can be scaled up to large two-dimensional systems without being affected bythe sign problem.
The quest for relatively simple but yet powerful variational states has focused onneural network states, which have shown a great deal of promise.
The complexity of the problem liesin the fact that many states with similar energy have very different physical properties.
Therefore, anaccurate representation of the ground state becomes the key to study the nature of the quantum phase.
2.2 Related work
Before the concept of NNQS became a popular new alternative for simulating many-body systems,the most successful numerical techniques to treat the 2D J 1 − J 2 model have been the density matrixrenormalization group (DMRG)[15], VMC based on a projected fermionic ansatz[20], and tensorproduct states[45].
Recently, some research has focused on improving the accuracy of NNQS byusing deep neural networks such as CNN[9] and group-CNN[36].
The idea of applying quantumnumber projection to recover the symmetries of the wave function[40, 46] has proven to be effectivein improving the performance of NNQS[9, 28, 29, 36].
In addition, other alternatives that enhancethe quality of the approximations consist of combining NNQS with Gutzwiller-projected fermionicwave functions[12], or pair-product wave functions[30].
Method
3.1 Neural Network Wave Function with symmetry
An RBM wave function takes a spin conﬁguration – a sequence of N values ± 1 / 2 – and returns acomplex coefﬁcient corresponding to the wave function amplitude.
In other words, it is a functionψ : {− 1 / 2 , +1 / 2 } N → IC .
This function is highly non-linear and is parametrized by biases (cid:126)a,(cid:126)b andweights W as:
[Equation]
In this expression, the number of “hidden variables” M is a tunable parameter.
While RBMs haveremained a simple example of a basic neural network for many decades, it was only recently thattheir potential as variational wave functions was appreciated [4].
In this case, unlike conventionalmachine learning applications, the biases and weights are complex valued.
It is possible to account for certain symmetries [27] of the problem directly within the internalmathematical structure of the RBM.
In particular:
• Spin ﬂip symmetry:
If the z -component of the total magnetization is zero ( (cid:80) i σ zi = 0 ), the
global spin ﬂip operation σ zi → − σ zi preserves this property.
Notice that since cosh ( x ) isan even function, we can easily restore the global ﬂip
symmetry in RBM wave function byremoving the “magnetic ﬁeld” terms associated to biases (cid:126)a,(cid:126)b in Eq.(2).
Thus, the RBMwave function coefﬁcients become:
[Equation]
• Translational symmetry:
In translationally invariant systems, the ground state of the Hamil-
tonian is expected to preserve the translational symmetry of the lattice.
By applying the
E
[Equation]
[Equation] [Equation] [Equation]
E
[Equation]
-   momentum projection, one can construct a variational wave function that preserved thesymmetry with a well deﬁned momentum K :
[Equation]
where the translation operator T R shift all particles by a distance R .
For the 2D squarelattice,
[Equation]
The operator T R will perform a translation by m steps in the ˆ x direction, and n steps inthe ˆ y direction.
The coefﬁcient ψ K now satisﬁes the translational symmetry at the cost ofrequiring a computation time N time larger.•
Lattice point symmetry:
As our target model is on the 2D square lattice, the point group
symmetries
, which consists of rotation and reﬂection operations
, may also be included:
[Equation]
where L is the symmetry operation in the C 4 v point group, and χ ( L ) is the character ofthe irreducible representation I for the symmetry operation L .
Since there are 8 operationsin the C 4 v point group, as shown in table 1, ψ K L is 8 times more expensive to calculatecompared to ψ K .
Notice that even though the computational cost of optimizing and evaluating observables with thesymmetrized wave function has increased, the resulting state has a much larger expressivity than theoriginal one, translating into a remarkable accuracy as we shall demonstrate.
We should highlighthere that the new states, by being linear combinations of RBMs, are no longer RBMs, and thereforeallow one to explore a much larger space outside the original manifold deﬁned by ψ s , Eq.(2).
3.2 Wave Function Optimization
The goal of the calculation is to minimize the cost function deﬁned by the expectation value of theenergy:
[Equation]
where the probability distribution is determined by the normalized wave function coefﬁcients
[Equation]
and the local energy is given by
[Equation]
By formulating the problem in probabilistic terms, one can resort to Metropolis-Hastings Markov
Chain Monte Carlo to evaluate the averages.
The sampling over the spin conﬁgurations (cid:126)σ is carriedout by randomly ﬂipping pairs of anti-aligned spins, and using von Neumann rejection according to atransition probability W = |(cid:104) (cid:126)σ new | ψ K L (cid:105)| 2 / |(cid:104) (cid:126)σ old
| ψ K L (cid:105)| 2 .
[Equation]
The wave function optimization can be implemented by a variety of methods.
Since the energylandscape is extremely complex, simple gradient descent tends to get trapped into metastable solutions.
More sophisticated strategies are usually employed, such as natural gradient descent or “stochasticreconﬁguration”[42].
In contrast to the "standard" natural gradient descent method, the Fubini-studymetric[33], which is the complex-valued form of Fisher information, is used to measure the "distance"between wave functions | ψ (cid:105) and | φ (cid:105) :
[Equation]
The procedure to update variational parameters using natural gradient descent is well describedin literature[4, 8, 30], and we hereby summarize it.
The optimization is done by minimizing the
Fubini-study metric between | e − dτH ψ ( θ ) (cid:105) and ψ ( θ + δθ ) (cid:105) where dτ is a small step in imaginarytime and can be viewed as learning rate in the training of neural network.
The optimal choice for δθis given by the solution of a system of equations:
[Equation]
(12) where the log derivative
O k = 1 ψ ( θ ) ∂ψ ( θ ) ∂θ and (cid:104)· · · (cid:105) means an average over samples.
We update theparameters by θ k
= θ (cid:48)
k + δθ k and repeat until convergence is reached.
3.3 Lanczos recursion
Using the symmetrized RBM wave function combined with the stochastic reconﬁguration method,a good approximation of the ground state can be achieved after hundreds or thousands iterations.
However, due to the limited representation power of neural network wave functions, and the errorsstemming from the Monte carlo sampling and the optimization method, the true ground state of the
Hamiltonian H can still differ signiﬁcantly from the variational one.
One possible way to increasethe expressivity of the wave function is to introduce additional hidden variables or layers.
However,an alternative to systematically improve the neural network wave function consists of applying amodiﬁed Lanczos recursion [14, 2, 20].
The procedure begins with a (normalized) trial wave functionψ 0 , which in our case is an initial guess for the ground state, ψ 0 = ψ K L .
Then, a new state ψ 1 isconstructed by applying the Hamiltonian on ψ 0 and subtracting the projection over ψ 0 in order topreserve orthogonality:
[Equation]
where (cid:104) H n (cid:105) =
(cid:104) ψ
| H n | ψ 0 (cid:105) .
Notice that ψ 1 is orthogonal to ψ 0 and also normalized.
In the usual
Lanczos method, this recursion can be continued such that a new complete orthogonal basis can beconstructed.
In this representation, the Hamiltonian will have tri-diagonal form.
However, we onlyuse ψ 0 and ψ 1 as our basis, and thus the Hamiltonian will be a 2 × 2 matrix.
The eigenvector ˜ ψ 0 that corresponds to the lowest eigenvalue ˜ E 0 of this matrix will be a betterapproximation of the true ground state of Hamiltonian compared to ψ 0 .
The lowest eigevalue and are
[Equation] [Equation]
where
[Equation] [Equation]
The eigenvector ˜ ψ 0 , being a linear combination of ψ 0 and ψ 2 , is the improved neural network wavefunction, and ˜ E 0 is the new improved variational energy.
By considering ˜ ψ 0 as the new trial wave
[Table]
function replacing ψ 0 , this method can be repeated to further improve the wave function.
The neuralnetwork wave function obtained during the Lanczos recursion can be generalized as
[Equation]
where p is maximum number of Lanczos steps, and β i is the wave function coefﬁcient correspondingto H
i | ψ 0 (cid:105) .
In this form, one can easily identify the wave function as an expansion in a Krylov basis.
In practice, taking into account the fact that the computational complexity increases dramaticallywith increasing p , only a few steps can be calculated for large quantum many body system.
In thisstudy, and for illustration purposes, we shall consider only the p = 1 or p = 2 cases.
3.4 Implementation details
In this work, we focus on the 2D J 1 J 2 Heisenberg model on L × L square lattices where L is an evennumber.
For the neural network, we use ψ K L in all simulations and consider three different valuesfor the number hidden variables M consisting of 2, 2.5 and 3 times of the number of spins N =
L 2in the system.
The parameters W in the RBM are initialized to be randomly chosen random numberswith a uniform distribution between [ − 0 .
for both real and imaginary parts.
The groundstate can belong to the A 1 or B 1 irreducible representations of the C 4 v point group, depending on thevalue of J 2 /J 1 .
In our calculations we consider both cases near the transition between the spin liquidphase and the columnar phase with K =
( π, 0) , i. e. for J 2 /J 1 ≥ 0 .
Due to the large number of parameters and the numerical noise in sampling, we implement theconjugate gradient method to solve the system of equations, Eq.(12).
To stabilize the method, weintroduce a ridge parameter λ
= 10 − 6 .
For each training step, we collect 10000 samples to evaluateaverages as mentioned in Sec. 3.2 including the variational energy and log derivatives.
Since theadjacent states in the Markov chain are highly correlated, the number of the skipped states betweensamples N skip is chosen according to this relation N skip = 5 × 1 . 0 /r , where r is the acceptancerate in the previous training step.
The typical value for N skip is from 30 to 100.
As for evaluation,we collect 2 × 10 5 samples to calculate the average and statistical error.
The learning rate usedin the training ranges from 5 × 10 − 4 to 3 × 10 − 2 .
Once we observe that variational energy is notdecreasing, a smaller learning rate(half of previous one) is used instead.
For large L , to save trainingtime, we initialize the parameters W in ψ K L using the parameters trained by means of the cheaperwave function ψ s .
All simulations are performed using Eigen and Intel MKL on Intel E5-2680v4 and
AMD Rome 7702 CPU nodes.
Results
4.1 Comparison with Exact Diagonalization
We benchmark the accuracy of the neural network wave functions for the ground state mainly onthe 6 × 6 and 10 × 10 square lattices with periodic boundary conditions.
For the 6 × 6 lattice, the
[Table] [Table]
S( π, π )(RBM) 2.06(3) 1.56(2) 1.18(2)
S( π, 0 )(RBM) 0.186(1) 0.191(2) 0.200(2)
J 1 − J 2 is numerically soluble by enumerating the possible spin conﬁgurations, constructing the
Hamiltonian matrix, and explicitly solving the eigenvalue problem [39].
Once the ground state (orits variational approximation) is obtained, the wave function can be used to calculate other physicalquantities besides the energy.
Here, for illustration, we compute the spin structure factor, that deﬁnesthe sublattice magnetization squared for a ﬁnite system
[Equation]
where the wave vector q determines spatial structure of the magnetic order.
Notice that in all thetables shown here, we display the results times a factor N for readability.
We ﬁrst focus on the symmetrized RBM wave function without the Lanczos optimization, and westart by comparing the ground state energy for a 6 × 6 lattice as a function of J 2 /J 1 , as shown in Fig.1.
In this ﬁgure we calculate the relative error as | E nn − E exact | / | E exact | using the exact ground stateenergy from Ref.
We also include the relative error of the ground state energy obtained usinga convolutional neural network wave function from Ref.[9].
While the relative error of the CNNs are in order of 10 − 3 , our RBM wave function achieves an accuracy of 10 − 4 in the frustrated regime.
Even comparing other recent works using CNNs[43, 36, 7], our RBM wave function still outperformsthe CNN wave function.
Besides the ground state energy, the spin structure factor computed fromoptimized wave functions agree very well with the exact solution as shown in Fig.
2, where thedifferences are smaller than the symbol size, and in data table 2.
4.2 Comparison with state-of-the-art quantum Monte Carlo
For larger lattices, the problem is numerically intractable.
However, as mentioned before, it can besolved using QMC[38] for J 2 = 0 .
Thus, for the case without frustration we can compare with QMCresults for several different lattice sizes.
From table 3, we can see that even on the 10 × 10 lattice theenergy difference is about 3 × 10 − 5 , showing the extraordinary accuracy of our RBM wave function.
For the frustrated case, J 2 (cid:54) = 0 , we compare to other methods, such as those obtained with CNNwave functions as well as results using the density matrix renormalization groump(DMRG) methodwith SU (2) symmetry from Ref.
[15] and VMC using an Abrikosov-fermion mean ﬁeld with a
Z 2 gauge structure from Ref.
From the data tables 4 and 5, we observe that our RBM wavefunction outperform the CNN wave function again in the entire range of J 2 /J 1 .
In the frustratedregime, comparisons with VMC and DMRG using all the data available in literature demonstrate thatthe RBM wave functions still yield competitive ground state energies except at J 2 /J 1 = 0 .
55 where
DMRG yields a lower value.
4.3 Lanczos optimization
Since the most interesting regime lies around the maximally frustrated point J 2 ∼ 0 .
5 J 1 , we choose 3different values of J 2 /J 1 using 6 × 6 and 10 × 10 lattices and perform a few Lanczos steps to furtherimprove the ground state energy.
From data table 2 and 4, we see that the Lanczos steps are veryeffective regardless of the system size.
Remarkably, by performing p = 1 Lanczos steps, we obtainbetter energy at J 2 /J 1 = 0 .
55 for the 10 × 10 lattice that improves signiﬁcantly the best availabledata using state-of-the-art DMRG, as shown in data table 4.
Besides, compared to the "RBM+PP"results[30], which is generally considered as the start-of-the-art NNQS method, we obtain slightlylower variational energy at J 2 = 0 . 55 on a 6 × 6 lattice while for a 10 × 10 lattice at J 2 = 0 . 5 ,their variational energy is 8 × 10 − 4 lower than ours.
Additionally, with the help of the Lanczosrecursion, a better estimate of the energy can be obtained by carrying out a variance extrapolationas illustrated in Ref.
We also try to improve the estimation of spin structure factor using Lanczos, but the Monte Carlo sampling error makes the improvement not obvious.
[Table]
Energy(RBM) -0.586411(9) -0.71429(1)
S( π, π )(Exact) 0.
S( π, π )(RBM) 0.0779(4)
S( π, 0 )(Exact) 1.
S( π, 0 )(RBM) 1.43(2) 1.53(2) 1.57(2) 1.59(2) 10 ×
J 2 = 0 .
J 2 = 0 . J 2 = 0 .
J 2 = 0 . J 2 = 0 .
Energy(CNN)
Energy(RBM) -0.671519(4)
S( π, π )(RBM) 5.38(6) 4.24(5) 2.59(4) 0.408(3) 0.1594(9)
S( π, 0 )(RBM) 0.1684(9) 0.1702(8) 0.180(2) 0.99(2) 2.33(7) 10 ×
J 2 = 0 .
J 2 = 1 . J 2 = 1 .
J 2 = 1 .
Energy(CNN)
Energy(RBM) -0.57404(2) -
[Table]
Conclusion
Neural network wave functions hold a great deal of promise due to their ability to compress complexquantum many body states within a relatively simple mathematical structure that, owed the its nonlinearity, can encode and exponentially large amount of information with polynomial resources.
Inparticular, RBM wave functions, initially deemed too simple, can be used as building blocks forsystematically improved wave functions.
These improved states obey the internal symmetries of themodel and the point group symmetries of the lattice.
In addition, they may contain contributions fromstate living in a “tangent space” to the original RBM manifold.
This tangent vectors are spanned interms of powers of the Hamiltonian and form a Krylov basis.
We have demonstrated that we can achieve state-of-the-art accuracy that improves previous resultsusing convolutional neural networks with a minimal amount of extra computational cost compared tosimple RBMs.
The combination of Lanczos and symmetrization offer a effective solution to problemspreviously beyond reach of the most powerful numerical techniques and provide the means to bypassthe sign problem.
These ideas can seamlessly translate to other areas of research ranging frommaterials science to quantum chemistry.
Besides, our variational solution can be adopted to calculatethe excitation spectrum of a quantum many-body system[17, 18], providing valuable information thatcan be directly compared to experiments.
Limitations The computational cost of a single training step scales as O ( N sample × MN 2 ) , wherethe number of hidden variables M is usually proportional to the system size N .
Thus, the computationtime of the calculation may be a bottleneck for its application on larger lattices.
Besides, the Lanczoscorrection will be smaller and smaller as p increases in the thermodynamic limit.
Negative Societal Impact
Our work presents the theoretical simulation of the quantum many-bodyproblems without any foreseeable negative societal impacts.