CIFAR AI Chair: 1.5B-Parameter RWKV-7 Model Achieves Scale in NQS Optimization

Researchers are applying a new approach to optimizing neural quantum states (NQS), utilizing a 1.5 billion-parameter RWKV-7 model to scale these approximations of quantum many-body wavefunctions beyond previous limitations. The team, spanning Mila Quebec AI Institute, Applied Quantum Algorithms, and Delft University of Technology, addresses challenges with existing optimization methods like stochastic reconfiguration, which they describe as costly and numerically fragile in large models. To overcome these hurdles, they frame variational energy minimization as an advantage policy-gradient problem, drawing a connection between reinforcement learning and quantum physics. This theoretical approach led to the development of Proximal Wavefunction Optimization (PWO), a trust-region algorithm that avoids explicit matrix inversion and reuses samples, improving stability and convergence across complex spin systems.

A 1.5 billion-parameter language model, typically used for text generation, is now being repurposed to tackle challenging problems in quantum physics. Researchers demonstrated the application of the RWKV-7 architecture to optimize Neural Quantum States (NQS), potentially changing how quantum systems are modeled and understood. Existing methods, like stochastic reconfiguration, become computationally prohibitive as system size increases. To circumvent this, they developed Proximal Wavefunction Optimization (PWO), a new trust-region algorithm that avoids explicit matrix inversion and reuses samples across multiple updates, offering a significant advantage over methods hampered by autocorrelation and slow mixing. The researchers state that PWO avoids explicit matrix inversion, reuses samples across multiple updates, and combines the scalability of first-order optimization with theoretical guarantees. Across simulations of Ising and frustrated spin systems, PWO demonstrably improved both stability and convergence speed when compared to established optimization algorithms, allowing for the exploration of more complex quantum systems with greater accuracy and efficiency.

Current approaches to approximating quantum many-body wavefunctions with Neural Quantum States (NQS) increasingly rely on autoregressive models, prized for their ability to perform exact, independent sampling from the Born distribution and circumvent limitations of Markov Chain Monte Carlo methods. Researchers have demonstrated the connection between variational energy minimization and policy-gradient reinforcement learning, revealing that the NQS gradient can be expressed as an advantage-weighted form over the Born distribution. They then fine-tuned a 1.5 billion-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work, highlighting a growing synergy between machine learning and quantum physics and leveraging techniques from disparate fields to tackle fundamental challenges in computational quantum mechanics.

Researchers are demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work. The team, based at institutions including Mila Quebec AI Institute and Leiden University, identified a gap in optimization principles for NQS. While autoregressive NQS offer advantages in sampling efficiency, their optimization has remained a significant hurdle; existing approaches present a trade-off between speed and accuracy. First-order optimizers like Adam, though scalable, ignore the geometry of the variational wavefunction and can converge unstably or inaccurately in NQS applications, according to the researchers. This allows PWO to clip probability-ratio changes, effectively controlling the step size and preventing drastic updates that can destabilize training. A 1.5 billion-parameter RWKV-7 model, a large language model repurposed for NQS optimization, was fine-tuned on a one-dimensional Ising model. Comparative tests against Adam, minSR, and SPRING on standard benchmarks, including one- and two-dimensional Ising and J1-J2 spin systems, revealed that PWO consistently improves both stability and convergence speed, suggesting a promising path toward more efficient and reliable NQS training.

Variational Energy Minimization as Policy-Gradient RL

The pursuit of simulating complex quantum systems has taken an unexpected turn, with researchers leveraging the architecture of large language models to optimize neural quantum states (NQS). This approach is not simply about applying more computational power; it’s a fundamental reframing of the optimization problem itself, drawing parallels between quantum mechanics and reinforcement learning. Researchers have demonstrated NQS optimization at a scale over three orders of magnitude beyond prior work by fine-tuning a 1.5 billion-parameter RWKV-7 model. This advancement stems from a novel theoretical insight: the researchers demonstrate that variational energy minimization, the core process in NQS, can be mathematically understood as an advantage policy-gradient problem over the Born distribution. This connection allows them to apply trust-region optimization techniques, commonly used in reinforcement learning, to NQS training, resulting in Proximal Wavefunction Optimization (PWO), a new algorithm designed to address the limitations of existing methods.

This suggests a pathway towards tackling increasingly complex quantum systems, potentially unlocking new insights in materials science, drug discovery, and fundamental physics. The team’s work highlights a growing trend of cross-disciplinary innovation, where techniques from seemingly disparate fields are combined to push the boundaries of scientific computation.

Scaling neural quantum state optimization has become feasible with a 1.5 billion-parameter RWKV-7 model, a feat previously beyond reach. Researchers affiliated with institutions including Mila Quebec AI Institute, Université de Montréal, Leiden University, Delft University of Technology, and CIFAR, have demonstrated this significant leap in the ability to model complex quantum systems. This work culminated in the successful fine-tuning of a 1.5 billion-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.

The pursuit of scalable methods for approximating quantum many-body wavefunctions has increasingly focused on neural quantum states (NQS), but optimizing these models presents significant challenges. Testing PWO on both Ising and frustrated J1-J2 one- and two-dimensional spin systems revealed substantial improvements in both stability and convergence compared to established optimizers like Adam, minSR, and SPRING. The team successfully fine-tuned a 1.5 billion-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.