Momentum Stability and Adaptive Control in Stochastic Reconfiguration
For practitioners of variational Monte Carlo with neural network wavefunctions, this work provides theoretical understanding and a practical solution to the instability caused by momentum in stochastic reconfiguration optimization.
The paper analyzes the sensitivity of the momentum parameter in SPRING for VMC optimization, proving convergence for μ<1 and divergence for μ=1, and proposes PRIME-SR, a tuning-free adaptive method that matches optimally tuned SPRING's performance with improved robustness.
Variational Monte Carlo (VMC) combined with expressive neural network wavefunctions has become a powerful route to high-accuracy ground-state calculations, yet its practical success hinges on efficient and stable wavefunction optimization. While stochastic reconfiguration (SR) provides a geometry-aware preconditioner motivated by imaginary-time evolution, its Kaczmarz-inspired variant, subsampled projected-increment natural gradient descent (SPRING), achieves state-of-the-art empirical performance. However, the effectiveness of SPRING is highly sensitive to the choice of a momentum-like parameter $μ$. The original sensitivity of $μ$ and the instability observed at $μ=1$, have remained unclear. In this work, we clarify the distinct mechanisms governing the regimes $μ<1$ and $μ=1$. We establish convergence guarantees for $0\leμ<1$ under mild assumptions, and construct counterexamples showing that $μ=1$ can induce divergence via uncontrolled growth along kernel-related directions when the step-size is not summable. Motivated by these theoretical insights and numerical observations, we further propose \textit{Principal Range Informed MomEntum SR} (PRIME-SR), a tuning-free momentum-adaptive SR method based on effective spectral dimension and subspace overlap. PRIME-SR achieves performance comparable to optimally tuned SPRING while significantly improving robustness in VMC optimization.