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

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.

• The paper introduces PWO, a novel trust-region optimization algorithm that recasts NQS training as a constrained policy-gradient problem with theoretical improvement guarantees.
• It demonstrates superior performance over optimizers like Adam, minSR, and SPRING in terms of stability, convergence speed, and scalability across quantum benchmarks.
• Empirical evaluations show PWO achieving relative errors as low as 10⁻⁷ and robust energy variance suppression in 1D, 2D, and billion-parameter systems.

Proximal Wavefunction Optimization: A Trust-Region Perspective on Neural Quantum State Training

Introduction and Theoretical Foundations

Neural quantum states (NQS) have emerged as expressive ansätze for variational quantum many-body problems, leveraging neural networks to approximate ground-state wavefunctions with high accuracy intractable for classical methods. Among NQS architectures, autoregressive models are distinguished by enabling exact, independent sampling from the Born distribution, thereby overcoming limitations inherent to the Markov Chain Monte Carlo (MCMC) approaches prevalent in traditional variational Monte Carlo (VMC).

Despite these sampling advantages, optimization of large autoregressive NQS models remains challenging. First-order optimizers such as Adam provide scalability but are agnostic to wavefunction geometry, frequently resulting in instability, slow convergence, or suboptimal solutions as the Hilbert space dimension grows. In contrast, stochastic reconfiguration (SR) and its variants (e.g., minSR, SPRING) are geometric but computationally prohibitive and numerically fragile for large-scale neural models. This optimization-sampling decoupling has motivated research into more principled, efficient training methods for NQS.

The paper "One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective" (2607.02292) introduces a new optimization paradigm for autoregressive NQS, leveraging a formal equivalence between variational quantum energy minimization and policy-gradient reinforcement learning (RL). Building on this conceptual bridge, the authors propose Proximal Wavefunction Optimization (PWO), a trust-region optimization algorithm structurally analogous to Proximal Policy Optimization (PPO) in RL, but tailored to the quantum variational context. PWO addresses instability and inefficiency in large-scale NQS training, providing a mathematically grounded and computationally scalable first-order optimizer with theoretical improvement guarantees.

Variational Quantum Energy Minimization as Advantage Policy Gradient

The key theoretical insight underpinning PWO is a formal equivalence between the variational energy minimization objective for NQS and the advantage-weighted policy-gradient objective in RL. Specifically, for stoquastic Hamiltonians and fixed-phase wavefunctions, the gradient of the variational energy with respect to the parameters $\boldsymbol{\theta}$ of the neural network ansatz reduces to

$\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$

where $E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s})$ is the local energy and $\mathcal{P}_{\boldsymbol{\theta}}$ is the Born probability. Here, the centered local energy acts as a surrogate for the RL advantage, and spin configurations correspond to discrete actions under the policy defined by the wavefunction squared modulus. This equivalence forms the backbone for importing trust-region algorithms from RL to quantum state optimization.

Proximal Wavefunction Optimization (PWO) Algorithm

PWO recasts NQS training as a constrained trust-region policy optimization problem. It inherits the sample-efficient, monotonic-improvement properties of PPO but extends them to complex-valued quantum wavefunctions. The algorithm consists of two main components:

• Amplitude Channel: Probability-ratio clipping is used to limit the change in the amplitude of the wavefunction between updates, ensuring conservative steps and reliable sample reuse.
• Phase Channel: Since the VMC gradient for complex-valued NQS includes an imaginary part governing the phase, PWO introduces direct phase increment clipping, preventing abrupt wavefunction rotations while maintaining first-order consistency with the complex-valued VMC direction.

This trust-region construction is motivated by the observation that, inside a "trust region," the surrogate loss landscape agrees with the true variational energy landscape, enabling reliable surrogate-based updates and safe sample reuse.

Figure 1: Trust-region intuition for PWO. Inside the trust region, the surrogate landscape gives a reliable local improvement direction; outside the trust region, the surrogate and true landscapes may disagree substantially.

Theoretical guarantees provided for PWO include first-order consistency with the original VMC gradient and an explicit improvement certificate for energy minimization under amplitude and phase trust regions, generalizing monotonic-improvement results from RL (see, e.g., TRPO, PPO). This enables rigorous, finite-update performance bounds and ensures robustness when scaling to deep architectures.

Empirical Evaluation: Benchmarking PWO Against Adam, minSR, and SPRING

The paper empirically evaluates PWO in standard and challenging NQS benchmarks:

• 1D Transverse-Field Ising Model: On a 12-site periodic chain, PWO achieves relative error $10^{-7}$ in approximately 5 minutes, compared to 30 minutes for minSR and significantly longer for Adam and SPRING. PWO also achieves rapid suppression of energy variance, indicating fast convergence to an eigenstate rather than merely a variational minimizer.

Figure 2: Comparison of PWO, Adam, minSR, and SPRING on the transverse-field Ising model over 10 random seeds. All methods were run with 1024 samples on a single NVIDIA L40S GPU.

• Heisenberg $J_1$--$J_2$ Chain (Majumdar–Ghosh Point): On strongly frustrated Hamiltonians, minSR suffers numerical instability (60% NaN runs), while PWO maintains stability and achieves relative error $10^{-7}$ in 15 minutes. Adam and SPRING both plateau at much higher error.

Figure 3: Comparison of PWO, Adam, minSR, and SPRING on the Heisenberg $J_1$--$J_2$ chain over 10 random seeds. MinSR is highly unstable on this Hamiltonian, with 6 out of 10 runs producing NaNs. All methods were run with 1024 samples on a single NVIDIA L40S GPU.

• 2D $\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$0--$\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$1 Square Lattice: On a $\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$2 2D Heisenberg model, PWO descends rapidly to a lower-energy regime and maintains low variance relative to Adam for the same wall-clock budget, demonstrating effectiveness beyond 1D chains and in the presence of high frustration and large Hilbert spaces.

Figure 4: Two-dimensional frustrated $\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$3--$\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$4 Heisenberg model on the $\grad_{\boldsymbol{\theta}} E[\psi_{\boldsymbol{\theta}}] = \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} \left[ \left( E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) - \mathbb{E}_{\mathbf{s} \sim \mathcal{P}_{\boldsymbol{\theta}}} [ E^{\mathrm{loc}}_{\boldsymbol{\theta}}(\mathbf{s}) ] \right) \grad_{\boldsymbol{\theta}} \log \mathcal{P}_{\boldsymbol{\theta}}(\mathbf{s}) \right],$5 lattice. Left: mean real energy. Right: V-score. PWO reaches lower energies faster than Adam and maintains a lower variance-based error signal over the same wall-clock budget.

Scalability to Large Models and Wall-Clock Efficiency

The scalability of PWO is further substantiated by experiments demonstrating consistent monotonic improvement in energy with increasing NQS size. PWO outpaces both minSR and Adam in wall-clock efficiency, especially for medium and large models. The speedup is attributed not only to algorithmic sample reuse but also to avoiding explicit matrix inversion, which is the computational bottleneck for geometric approaches like minSR.

Figure 5: Wall-clock scaling comparison across model sizes and optimization methods. Boxplots show the interquartile mean relative error over seeds, with boxes indicating the interquartile range and lines indicating the min and max. Results are grouped by model size and wall-clock time, and run on a single NVIDIA A100 GPU.

Billion-Parameter Fine-Tuning: RWKV-7 as NQS

PWO is applied to fine-tune a 1.5B-parameter RWKV-7 LLM as an NQS on the 1D Ising model—exceeding previous NQS training scale by over three orders of magnitude. PWO maintains stability and achieves lower final error and variance compared to Adam, indicating robust optimization even in the billion-parameter regime.

Figure 6: Fine-tuning curves of a 1.5B-parameter RWKV7LLM on the 1-D Ising Model.

Figure 7: Individual-seed fine-tuning curves for the 1.5B-parameter RWKV-7 neural quantum state on the transverse-field Ising model. Each curve corresponds to one random seed, with relative error shown on the left and V-score on the right. PWO consistently remains stable across seeds and reaches lower final error and variance than the Adam baseline, indicating that the proximal objective improves robustness even in the billion-parameter regime.

This empirical result demonstrates that PWO unlocks NQS training at modern deep learning scale, motivating further exploration of large-scale expressive models for quantum simulation.

Practical and Theoretical Implications

Practical Implications: PWO represents a significant advance for the scalable, stable optimization of neural quantum states, particularly for expressive autoregressive models. By combining trust-region benefits with first-order scalability and sample-efficient updates, PWO enables practical NQS training on complex, frustrated Hamiltonians and with billion-parameter neural architectures. This opens new directions for simulating quantum systems with machine learning methods at unprecedented scale and accuracy.

Theoretical Implications: The formal equivalence between VMC and advantage-weighted policy-gradient RL established in the paper unifies two previously parallel literatures and creates a systematic pathway for importing RL optimization theory—trust regions, monotonic improvement bounds, surrogate objectives—into quantum variational learning. The energy improvement certificates derived in this framework provide strong justification for PWO's stability and resilience to hyperparameter changes, especially as model and system size increase.

Future Directions: The work suggests a number of avenues for further exploration, including application of trust-region and other RL-derived optimization methods to fermionic NQS, adaptive trust-region scheduling for more aggressive optimization, and systematic study of architectural scaling (the "bitter lesson") for quantum many-body simulation. Additionally, the demonstrated feasibility of fine-tuning LLMs as NQS points towards a potential synthesis between mainstream sequence modeling advances and physics-inspired neural quantum simulation.

Conclusion

"One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective" introduces a robust, scalable, and theoretically justified algorithm—Proximal Wavefunction Optimization (PWO)—for training autoregressive NQS. By formalizing the connection between VMC and policy-gradient RL, and instantiating a PPO-derived trust-region surrogate adapted to the amplitude and phase channels of complex quantum wavefunctions, PWO achieves superior numerical stability and wall-clock convergence relative to both first-order and geometric baselines. Empirical results across standard and challenging quantum benchmarks, and at scales up to 1.5B parameters, support the claims.

PWO, together with the underlying RL-VMC equivalence, bridges previously isolated research areas and inaugurates new strategies for highly scalable, expressive neural quantum simulation. Future research will likely build on these ideas, exploring both architectural and optimization frontiers in the simulation of quantum many-body systems using deep learning.