Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo
Authors: Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel, Richard T. Scalettar, Ehsan Khatami
Organizations: Department of Physics and Astronomy, University of California, Davis, California 95616, USA · Department of Applied Mathematics, University of Waterloo, Waterloo, ON N2L 3G1, Canada · Vector Institute, Toronto, Ontario, M5G 0C6, Canada · Department of Physics and Astronomy, San Jos´e State University, San Jos´e, California 95192, USA
Abstract
Neural-network quantum states have emerged as a powerful variational framework for quantum many-body systems, with recent progress often driven by massively parallel architectures such as transformers. Recurrent neural network quantum states, however, are frequently regarded as intrinsically sequential and therefore less scalable. Here we revisit this view by showing that modern recurrent architectures can support fast, accurate, and computationally accessible neural quantum state simulations. Using autoregressive recurrent wave functions together with recent advances in parallelizable recurrence, we develop variational ansätze, called parallel scan recurrent neural quantum states (PSR-NQS), which can be trained efficiently within variational Monte Carlo in one and two spatial dimensions. We demonstrate accurate benchmark results and show that, with iterative retraining, our approach reaches two-dimensional spin lattices as large as 52×52 while remaining in agreement with available quantum Monte Carlo data. Our results establish recurrent architectures as a practical and promising route toward scalable neural quantum state simulations with modest computational resources.
Neural quantum states (NQS) provide a flexible and scalable framework for approximating quantum many-body wavefunctions. Among NQS parameterizations, autoregressive models are especially attractive because they enable exact, independent sampling from the Born distribution, avoiding the autocorrelation and mixing issues of Markov chain methods. Yet their optimization remains comparatively underexplored: Adam is a scalable method but ignores function space geometry, while stochastic reconfiguration is principled but costly and numerically fragile in large models. To address this gap, we show that variational energy minimization can be viewed as an advantage policy-gradient problem over the Born distribution, motivating trust-region optimization for NQS training. We introduce Proximal Wavefunction Optimization (PWO), a principled trust-region algorithm that clips probability-ratio changes in the amplitude channel and phase increments in the phase channel. PWO avoids explicit matrix inversion, reuses samples across multiple updates, and combines the scalability of first-order optimization with theoretical guarantees. Across Ising and frustrated J1-J2 one- and two-dimensional spin systems, PWO improves stability and wall-clock convergence over Adam, minSR, and SPRING. Finally, we fine-tune a 1.5B-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.
Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes +4
In this work, we extend the class of previously introduced non-Euclidean neural quantum states (NQS) which consists only of Poincare hyperbolic GRU, to new variants including Poincare RNN as well as Lorentz RNN and Lorentz GRU. In addition to constructing the new non-Euclidean hyperbolic NQS ansatzes, we generalize the results of our earlier work regarding the definitive outperformances delivered by hyperbolic Poincare GRU NQS when benchmarked against their Euclidean counterparts in the Variational Monte Carlo (VMC) experiments involving the Heisenberg J1J2 and J1J2J3 models. Here, using larger systems consisting of 100 spins, we find that all four hyperbolic RNN/GRU NQS variants always outperform their respective Euclidean counterpart with the same architecture. In our experiments, among the four hyperbolic NQS, Lorentz RNN stands out in particular because despite having almost three times fewer parameters, it is capable of surpassing the more complex Poincare GRU and Lorentz GRU to emerge as the best overall hyperbolic NQS ansatz on many instances involving different J2 and (J2,J3) couplings. Given the findings from this work showing that the four newly constructed hyperbolic RNN/GRU NQS ansatzes are able to outperform the well-established Euclidean RNN/GRU NQS in Heisenberg spin models, we establish the utility and efficiency of the hyperbolic Poincare RNN/GRU and Lorentz RNN/GRU NQS for future variational studies of quantum many-body systems, especially those exhibiting a hierarchical structure in the form of the different degrees of nearest-neighbor interactions.
Neural quantum states offer expressive representations of quantum many-body wave functions, yet their practical accuracy can be limited by stochastic optimization rather than representational capacity. Here we identify a finite-sample instability, termed subspace trapping, in which physically important configurations become strongly underestimated, remain absent from successive sampling batches and receive insufficient gradient feedback. This self-reinforcing loss of sampled support can confine optimization to an effective subspace and produce apparently stationary states above the true ground state energy. To address this problem, we introduce annealed gradient descent (AGD), a sampling-aware update with annealing factor that temporarily increases the relative contribution of sampled low-probability configurations while limiting the dominance of high-probability ones. We establish the connection between finite-sample support loss and effective subspace optimization, and then evaluate the method across molecular systems, one and two-dimensional J1-J2 models. Annealed gradient descent suppresses metastable trapping, preserves physically relevant configurations and enables compact neural quantum states to attain chemical accuracy and competitive state-of-the-art performance. These results establish AGD as a lightweight complement to expressive neural architectures, improved sampling strategies for scalable quantum many-body optimization.