PWO:神经量子态的临近波函数优化算法

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

精选理由

这篇把强化学习的信任域优化用在神经量子态上,新算法PWO比Adam稳,还能用1.5B的RWKV-7跑,规模大了三个数量级。

AI 摘要

该论文将变分能量最小化重新表述为Born分布上的策略梯度问题,并提出信任域算法PWO(Proximal Wavefunction Optimization)。PWO分别裁剪振幅通道的概率比变化和相位通道的相位增量,避免了显式矩阵求逆,并可重用样本进行多次更新。在Ising和受挫J1-J2一维/二维自旋系统上,PWO的稳定性和收敛速度优于Adam、minSR和SPRING。作者还使用1.5B参数的RWKV-7模型微调,将NQS优化规模提升到此前工作的千倍以上。

原文 · arXiv cs.LG

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

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.