Enhanced NQS via Annealed Gradient Descent

The Challenge of Subspace Trapping

Neural quantum states (NQS) use neural networks to represent complex many-body wave functions. While expressive, their training is often limited by the stochastic nature of Variational Monte Carlo (VMC). The authors identify a specific failure mechanism: subspace trapping. In this scenario, configurations that are critical to the ground state are assigned low probabilities by the model, causing them to be rarely sampled. Because these configurations are absent from the sampling batches, the optimizer receives no gradient feedback to correct their amplitudes. This creates a self-reinforcing cycle where the model collapses into a restricted subspace of the Hilbert space, converging to a metastable state that is numerically stable but energetically inferior to the true ground state.

Annealed Gradient Descent (AGD)

To address this, the authors propose annealed gradient descent (AGD). Unlike standard optimization that relies solely on the current model distribution, AGD introduces a temperature-dependent annealing factor. During the early stages of training, AGD artificially increases the relative weight of sampled low-probability configurations while tempering the influence of high-probability ones. This forces the model to maintain a broader exploration of the configuration space. As training progresses, the annealing factor is gradually reduced, allowing the optimizer to recover the standard variational gradient for final energy refinement.

This approach provides a lightweight, scalable solution to a fundamental bottleneck in quantum many-body optimization. Existing methods like stochastic reconfiguration or natural gradient descent focus on the geometry of the parameter space, but they remain susceptible to the sampling-induced mode collapse described here. By modifying the gradient update rule rather than the sampling distribution or the neural architecture, AGD achieves competitive performance without the prohibitive computational costs associated with massive sample sizes or complex curvature-aware updates. The authors demonstrate that this method enables compact neural networks to achieve chemical accuracy in molecular systems and effectively navigate the energy landscapes of $J_1$-$J_2$ lattice models.