One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

2026-07-02

AI Summary

This work addresses the longstanding challenge in neural quantum state (NQS) optimization of simultaneously achieving scalability, stability, and geometric awareness. By reformulating variational energy minimization as an advantage policy gradient problem grounded in the Born distribution, the authors introduce—inspired by trust-region methods from reinforcement learning—a novel Proximal Wavefunction Optimization (PWO) algorithm. PWO avoids explicit matrix inversion by clipping amplitude probability ratios and constraining phase increments, enabling efficient and stable training. Integrating autoregressive NQS architectures with advantage-based policy gradients and trust-region constraints, the method significantly outperforms established optimizers such as Adam, minSR, and SPRING across multiple one- and two-dimensional spin systems. Notably, PWO successfully optimizes an RWKV-7 model with 1.5 billion parameters, surpassing prior NQS scales by over three orders of magnitude.

Abstract

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 $J_1$-$J_2$ one- and two-dimensional spin systems, PWO improves stability and wall-clock convergence over Adam, minSR, and SPRING. Finally, we fine-tune a $1.5$B-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.

Problem

• Neural Quantum States
• Autoregressive Models
• Variational Energy Minimization
• Optimization
• Quantum Many-Body

Innovation

• Neural Quantum States
• Reinforcement Learning
• Trust-Region Optimization
• Proximal Wavefunction Optimization
• Autoregressive Models