[Paper Review] A simple linear algebra identity to optimize Large-Scale Neural Network Quantum States

Riccardo Rende, Luciano Loris Viteritti+3 more|arXiv (Cornell University)|Oct 9, 2023

Quantum many-body systems45 references4 citations

TL;DR

This paper introduces a novel application of a linear algebra identity to enable exact Stochastic Reconfiguration (SR) optimization in large-scale Neural Network Quantum States (NNQS), overcoming the traditional $P \times P$ matrix inversion bottleneck by reducing it to an $M \times M$ problem. The method achieves state-of-the-art ground-state energy for the $J_1$-$J_2$ Heisenberg model at $J_2/J_1 = 0.5$ on a $10 \times 10$ lattice using a Deep Vision Transformer with 267,720 parameters and only 6,000 Monte Carlo samples.

ABSTRACT

Neural-network architectures have been increasingly used to represent quantum many-body wave functions. These networks require a large number of variational parameters and are challenging to optimize using traditional methods, as gradient descent. Stochastic Reconfiguration (SR) has been effective with a limited number of parameters, but becomes impractical beyond a few thousand parameters. Here, we leverage a simple linear algebra identity to show that SR can be employed even in the deep learning scenario. We demonstrate the effectiveness of our method by optimizing a Deep Transformer architecture with $3 imes 10^5$ parameters, achieving state-of-the-art ground-state energy in the $J_1$-$J_2$ Heisenberg model at $J_2/J_1=0.5$ on the $10 imes10$ square lattice, a challenging benchmark in highly-frustrated magnetism. This work marks a significant step forward in the scalability and efficiency of SR for Neural-Network Quantum States, making them a promising method to investigate unknown quantum phases of matter, where other methods struggle.

Motivation & Objective

• Overcome the computational infeasibility of Stochastic Reconfiguration (SR) in large-scale Neural Network Quantum States (NNQS) where the number of parameters $P$ exceeds the number of Monte Carlo samples $M$.
• Enable exact SR optimization for deep architectures like Deep Transformers with hundreds of thousands of parameters, which are otherwise intractable with standard SR due to $P \times P$ matrix inversion.
• Achieve high-accuracy ground-state energy estimates for strongly frustrated quantum spin systems, such as the $J_1$-$J_2$ Heisenberg model at $J_2/J_1 = 0.5$, where other methods struggle.
• Demonstrate that high parameter counts in NNQS do not require large sample sizes $M$ for accurate SR, challenging the conventional belief that $M \gg P$ is necessary.
• Extend the applicability of SR to complex, sign-structure-free quantum systems by removing assumptions on wave function sign patterns.

Proposed method

• Leverage a simple linear algebra identity to transform the $P \times P$ matrix inversion in standard SR into an equivalent $M \times M$ matrix inversion, where $M$ is the number of Monte Carlo samples.
• Apply the identity to the Fisher information matrix in the SR framework, enabling exact computation of the inverse without approximations.
• Use the transformed $M \times M$ matrix to compute the SR update direction for the neural network parameters, maintaining the geometric optimization benefits of SR.
• Implement the method within a variational Monte Carlo framework using JAX for differentiable neural network computation and mpi4jax for distributed parallelization.
• Combine the modified SR update with a real-valued Deep Vision Transformer (ViT) architecture followed by a complex-valued fully connected output layer to represent complex-valued quantum states.
• Restore physical symmetries (translational, rotational, reflectional, spin parity) incrementally during optimization, observing consistent energy decreases.

Experimental results

Research questions

• RQ1 Can Stochastic Reconfiguration be applied exactly to large-scale NNQS with $P \gg M$ without approximations?
• RQ2 Does the proposed linear algebra identity enable efficient and stable optimization of deep neural networks with hundreds of thousands of parameters in quantum many-body systems?
• RQ3 Can this method achieve state-of-the-art ground-state energy estimates for highly frustrated quantum spin models like the $J_1$-$J_2$ Heisenberg model at $J_2/J_1 = 0.5$?
• RQ4 Is it possible to achieve high accuracy with SR using only $M \approx 6,000$ samples even when $P \approx 267,720$?
• RQ5 How does the incremental restoration of physical symmetries affect the variational energy during optimization in the absence of a sign prior?

Key findings

• The proposed method enables exact Stochastic Reconfiguration for NNQS with up to 267,720 parameters using only 6,000 Monte Carlo samples, bypassing the $P \times P$ inversion bottleneck.
• The method achieves a ground-state energy of $-0.49575(3)$ per site for the $10 \times 10$ $J_1$-$J_2$ Heisenberg model at $J_2/J_1 = 0.5$, setting a new state-of-the-art result.
• Energy consistently decreases when restoring physical symmetries (translational, rotational, reflectional, spin parity) in the Deep ViT architecture, confirming the method's stability and physical consistency.
• The optimization is robust and insensitive to initial seeds, with consistent results across multiple runs, indicating high reliability.
• The method performs well without assuming a sign structure for the ground state, making it applicable to a broad class of quantum systems where such assumptions are unavailable.
• The approach challenges the conventional requirement of $M \gg P$ for accurate SR, demonstrating that $M \approx 6,000$ is sufficient even for $P \approx 267,720$.