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

Neural quantum states (NQS) offer a powerful, scalable way to approximate complex quantum many-body wavefunctions using neural networks. While autoregressive models allow for exact, independent sampling from the Born distribution, their optimization remains a significant challenge. Standard first-order optimizers like Adam are scalable but often unstable, while second-order methods like stochastic reconfiguration (SR) are geometrically principled but computationally expensive and numerically fragile. This paper asks whether reinforcement learning (RL) principles can provide a more stable, scalable optimization framework for NQS.

The authors establish a formal mathematical connection between variational energy minimization and the policy-gradient objective in RL. By treating spin configurations as actions and centered local energies as advantages, they reformulate NQS training as a trust-region optimization problem. They introduce Proximal Wavefunction Optimization (PWO), which adapts the Proximal Policy Optimization (PPO) algorithm to the quantum setting. PWO uses clipped probability-ratio changes for the amplitude channel and clipped phase-increment surrogates for the phase channel, allowing for stable updates and sample reuse without the need for explicit matrix inversion.

PWO demonstrates superior performance across several benchmarks, including 1D and 2D Ising and frustrated J1–J2 spin systems. It consistently outperforms Adam, minSR, and SPRING in both wall-clock convergence speed and numerical stability. Notably, while second-order methods like minSR often suffer from numerical instability in frustrated regimes, PWO remains robust. The authors further demonstrate the scalability of their approach by fine-tuning a 1.5-billion-parameter RWKV-7 model, showing that PWO can successfully optimize NQS at a scale three orders of magnitude larger than previous work.