Momentum Stability and Adaptive Control in Stochastic Reconfiguration

TL;DR

This paper analyzes the stability of momentum-based stochastic reconfiguration methods in neural quantum state optimization, providing theoretical insights and proposing a robust, tuning-free adaptive algorithm called PRIME-SR.

Contribution

It clarifies the mechanisms of momentum sensitivity in SPRING, establishes convergence conditions, and introduces PRIME-SR, a spectral-informed adaptive method for improved stability.

Findings

SPRING's effectiveness is highly sensitive to the momentum parameter .

Counterexamples show  can cause divergence due to uncontrolled growth.

PRIME-SR achieves comparable performance to tuned SPRING with enhanced robustness.

Abstract

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 $\mu$. The original sensitivity of $\mu$ and the instability observed at $\mu=1$, have remained unclear. In this work, we clarify the distinct mechanisms governing the regimes $\mu<1$ and $\mu=1$. We establish convergence guarantees for $0\le\mu<1$ under mild assumptions, and construct counterexamples showing that $\mu=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.