⚛️ quantum physics

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

This paper introduces Proximal Wavefunction Optimization (PWO), a trust-region reinforcement learning algorithm that enhances the stability and scalability of training autoregressive Neural Quantum States by clipping probability ratios and phase increments, enabling efficient optimization of models up to 1.5 billion parameters.

Original authors: Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes, Eliška Greplová, Thomas Spriggs, Aaron Courville, Anna Dawid

Published 2026-07-03

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes, Eliška Greplová, Thomas Spriggs, Aaron Courville, Anna Dawid

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Finding the Perfect Arrangement

Imagine you have a massive, complex puzzle made of billions of tiny magnets (quantum particles). Your goal is to arrange them in a specific pattern that uses the least amount of energy possible. In physics, this is called finding the "ground state."

For a long time, scientists have used Neural Quantum States (NQS) to solve this. Think of an NQS as a super-smart AI robot that tries to guess the correct arrangement of magnets. The better the robot guesses, the closer it gets to the true answer.

However, teaching this robot has been tricky. The paper introduces a new way to train the robot that is faster, more stable, and allows it to handle puzzles that are thousands of times bigger than before.

The Problem: The "Wobbly" Training

To teach the robot, scientists use a method called "training." Imagine you are trying to teach a dog to sit.

Old Method (Adam): You give the dog a treat every time it sits, but you don't care how much it wiggles before sitting. Sometimes, the dog gets confused and starts spinning wildly, making it hard to learn.
Strict Method (Stochastic Reconfiguration): You try to be very precise, measuring exactly how much the dog's posture changes with every move. This is very accurate, but it requires so much math and calculation that the training process is incredibly slow and the computer often crashes (like a calculator trying to divide by zero).

The paper says: "We need a method that is as fast as the first one but as stable as the second one."

The Solution: The "Trust Region" (PWO)

The authors realized that training these quantum robots is mathematically very similar to Reinforcement Learning (RL), the same technology used to teach AI to play video games like Super Mario or StarCraft.

In RL, there is a famous technique called Proximal Policy Optimization (PPO). Imagine a trainer who says to the dog: "You can move, but don't move too far away from where you were a moment ago. If you move too wildly, I'll stop you." This keeps the training steady and prevents the dog from getting confused.

The authors created a new algorithm called Proximal Wavefunction Optimization (PWO). It applies this "don't move too far" rule to the quantum robot.

How PWO works with two "channels":
Quantum waves have two parts: Amplitude (how strong the wave is) and Phase (the timing or "color" of the wave).

The Amplitude Channel: The algorithm clips the changes in strength. If the robot tries to change its guess too drastically, the algorithm pulls it back, ensuring it stays close to its previous "good guess."
The Phase Channel: The algorithm does the same for the timing. It prevents the robot from spinning the wave's "phase" too quickly, which would cause the math to break.

The Results: Faster and Stronger

The team tested this new method on several difficult quantum puzzles (called spin systems).

Speed: On standard puzzles, PWO found the solution much faster than the old methods. While other methods took 30 minutes to reach a certain level of accuracy, PWO did it in about 5 minutes.
Stability: On the hardest puzzles (where magnets are "frustrated" and can't agree on a pattern), the old methods often crashed or gave up. PWO kept going steadily.
Scale: The biggest breakthrough was scaling. The authors took a massive AI model (a Large Language Model with 1.5 billion parameters, similar to the ones that power chatbots) and used it to solve a quantum puzzle.
- Analogy: Imagine using a supercomputer designed to write novels to solve a simple math problem. Previous methods couldn't handle such a huge machine without breaking. PWO allowed them to fine-tune this giant model successfully, proving that quantum simulation can now use the same massive tools as modern AI.

The Takeaway

This paper bridges two worlds: Quantum Physics and Reinforcement Learning. By realizing that training a quantum AI is like training a game-playing AI, the authors borrowed a "trust region" trick. This trick stops the AI from making wild, confusing moves, allowing it to learn complex quantum patterns faster and on a scale never seen before.

In short: They figured out how to teach a quantum robot to walk without tripping, allowing it to run much faster and tackle much bigger mountains than ever before.

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

Problem Statement

Neural Quantum States (NQS) offer a scalable framework for approximating quantum many-body wavefunctions, particularly when utilizing autoregressive models that enable exact, independent sampling from the Born distribution. However, the optimization of these models remains a significant bottleneck. Existing methods face a sharp trade-off:

First-order optimizers (e.g., Adam): Computationally efficient and scalable but ignore the geometry of the variational wavefunction, often leading to unstable or inaccurate convergence in NQS applications.
Stochastic Reconfiguration (SR) and variants (e.g., minSR): Geometrically principled (approximating natural-gradient descent) but computationally prohibitive for large networks due to the need to solve large, often ill-conditioned linear systems. They also suffer from numerical fragility in complex regimes.

While autoregressive NQS resolve sampling bottlenecks, stable and scalable optimization for large autoregressive wavefunction models remains an open challenge.

Methodology

The authors propose Proximal Wavefunction Optimization (PWO), a trust-region algorithm inspired by Proximal Policy Optimization (PPO) from Reinforcement Learning (RL). The core methodology rests on a formal equivalence between variational energy minimization in NQS and policy-gradient optimization in RL.

RL-VMC Equivalence: The paper establishes that for stoquastic Hamiltonians, the gradient of the variational energy can be expressed as an advantage-weighted policy-gradient update over the Born distribution. In this mapping:
- Spin configurations correspond to actions.
- Centered local energies ( $E_{loc}(s) - \mathbb{E}[E_{loc}]$ ) correspond to advantages.
- The Born distribution corresponds to the policy.
Dual-Channel Optimization: Unlike standard PPO which handles a single probability distribution, NQS wavefunctions are complex-valued, requiring separate handling of amplitude and phase. PWO introduces two clipped surrogate objectives:
- Amplitude Channel: Uses a clipped probability ratio ( $r_\theta = P_\theta / P_{\theta_{old}}$ ) to constrain updates, mirroring standard PPO. This ensures the updated wavefunction remains close to the reference in terms of probability mass.
- Phase Channel: Introduces a novel clipped surrogate that constrains the wrapped phase increment ( $\phi_\theta$ ) rather than a probability ratio. This prevents abrupt rotations of the wavefunction while following the imaginary part of the Variational Monte Carlo (VMC) gradient.
Algorithmic Implementation: PWO performs multiple inner updates ( $K$ ) on a frozen batch of samples drawn from the reference distribution $P_{\theta_{old}}$ . It reuses these samples across updates by employing importance weighting (with stop-gradient operators for the phase channel) to correct for distribution mismatches. This avoids explicit matrix inversion, retaining the scalability of first-order methods while enforcing trust-region constraints.

Key Contributions

Formal Connection: The authors formally connect variational energy minimization to advantage policy-gradient RL, demonstrating that the NQS gradient admits an advantage-weighted form over the Born distribution.
PWO Algorithm: Introduction of Proximal Wavefunction Optimization, a trust-region algorithm that clips probability-ratio changes in the amplitude channel and phase increments in the phase channel.
Theoretical Guarantees: Proof that the PWO surrogate satisfies a trust-region improvement bound and that the first-order update matches the exact VMC gradient direction when on-policy.
Scalability Demonstration: Successful fine-tuning of a 1.5 billion-parameter RWKV-7 Large Language Model (LLM) as an autoregressive NQS, demonstrating optimization at a scale over three orders of magnitude beyond prior work.

Experimental Results

The authors evaluated PWO against Adam, minSR, and SPRING across several benchmarks:

1D Spin Systems (Ising and J1–J2):
- Ising Model: PWO converged to a relative error of $10^{-7}$ in approximately 5 minutes, significantly faster than minSR (~30 minutes) and Adam.
- Frustrated J1–J2 Chain: PWO maintained stable convergence with a decreasing V-score. In contrast, minSR was numerically unstable (6 out of 10 runs produced NaNs), and Adam plateaued at a much higher error.
2D Frustrated J1–J2 Model: On a $10 \times 10$ lattice, PWO reached lower energy regimes faster than Adam and maintained a lower variance-based error signal (V-score) within the same wall-clock budget.
Scalability: As model size increased (from "tiny" to "medium" NQS), PWO demonstrated consistent stability and monotonic energy improvement. Adam degraded in stability and efficiency for larger models, while minSR failed to reach competitive energy values.
Large-Scale Fine-Tuning: The 1.5B-parameter RWKV-7 model fine-tuned with PWO achieved lower final relative errors and V-scores than Adam, maintaining stable training throughout the budget.

Significance and Claims

The paper claims that PWO bridges the gap between the scalability of first-order optimizers and the stability of geometric methods (like SR) by leveraging the trust-region principles of RL.

Stability in Difficult Regimes: The authors assert that PPO-style proximal objectives provide the necessary stability to train expressive NQS on difficult, frustrated quantum Hamiltonians where traditional methods fail or become unstable.
Sample Efficiency: By enabling sample reuse across multiple updates without explicit matrix inversion, PWO improves wall-clock convergence speed.
Scale: The successful application to a 1.5B-parameter model suggests that NQS optimization can now enter the regime of modern large-scale sequence modeling. The authors posit that this opens the door to investigating whether scale alone can overcome current barriers in quantum many-body simulation, aligning with the "bitter lesson" in AI.

The work does not claim that LLMs are the optimal inductive bias for 1D Hamiltonians, but rather uses the RWKV-7 model to demonstrate the efficacy of the proximal objective at unprecedented scales.