Organizations: National Laboratory of Solid State Microstructures and Department of Physics, Nanjing University, Nanjing 210093, China · Department of Physics, University of Massachusetts, Amherst, Massachusetts 01003, USA · Collaborative Innovation Center of Advanced Microstructures, Nanjing University, Nanjing 210093, China · Jiangsu Physical Science Research Center, Nanjing 210093, China · Hefei National Laboratory, Hefei 230088, People’s Republic of China · A. Alikhanyan National Science Laboratory, Br. Alikhanian 2, Yerevan 0036, Armenia · Institute of Theoretical Physics, Faculty of Physics, University of Warsaw, Warsaw, Poland
Neural networks provide expressive representations for scientific computing. However, even sufficiently expressive networks can suffer training failure in weak-gradient regimes, limiting their practical use in quantum many-body physics and ab initio quantum chemistry. Here we derive an unbiased direct gradient estimator and introduce the adaptive minimum-variance phase (AMVP) estimator for neural-network variational optimization. By improving the signal-to-noise ratio of weak gradients, these methods enable reliable scientific calculations where training previously failed, while substantially reducing computational cost. The framework enables compact networks to outperform larger and fine-tuned default standard-estimator models with over an order of magnitude less GPU time on correlated flux models, and ultimately exceed the density matrix renormalization group (DMRG) accuracy. It further achieves chemical accuracy in N2 bond breaking and, for the first time, in heavy-element I2 with explicit spin-orbit coupling. These results demonstrate that gradient-estimator design expands the capabilities of neural-network variational methods for accurate scientific computing.
Figures & tables
Figure 1: Estimator design changes the stochastic optimization dynamics. (a) The standard, direct, and adaptive minimum-variance phase (AMVP) estimators share the same population gradient but exhibit different finite-sample fluctuations. (b) The standard estimator first aggregates the imaginary local-energy contributions before coupling them to the phase score, whereas the direct estimator retains connection-resolved phase differences that remove common-mode fluctuations. AMVP adaptively combines unbiased estimators according to their stochastic properties. (c) Different estimator noise structures lead to different optimization trajectories. The standard estimator can leave optimization trapped in metastable regimes, whereas the direct and AMVP estimators allow neural states to better exploit their available capacity and approach target states.
Figure 2: Gradient-estimator choice overcomes a phase-optimization barrier. (a) On a 50-rung flux ladder, the standard estimator remains trapped in a high-error regime despite variations in learning rate and network capacity. In contrast, the direct and adaptive minimum-variance phase (AMVP) estimators allow the same neural-network representation to reach subpercent errors, with median errors of 0.50% and 0.43% , respectively. (b) The chiral current profile obtained with direct and AMVP agrees with the DMRG reference ( jc/jcDMRG≈1.01 ), whereas the standard estimator shows a marked deviation ( jc/jcDMRG≈0.79 ). (c) Training trajectories and computational cost. Both a direct-trained ResNet and a compact AMVP-trained MLP reach energy errors below 0.5%, whereas the standard-trained ResNet remains in a higher-error regime. The inset compares allocated A100 GPU times for AMVP-trained MLP and standard-trained ResNet runs: 0.36 GPU h and 4.61 GPU h, respectively. (d) Energy errors for an 8×8 flux square with periodic boundary conditions. Energies obtained with AMVP estimators are lower than the finite- χ DMRG energies obtained up to χ=2048 , whereas the standard estimator remains in a higher-error regime. (e) Cumulative fraction of runs that have reached subpercent energy error.
Figure 3: Estimator-limited trainability of molecular wavefunctions. (a) FCI errors along the N 2 bond stretch for the jointly optimized amplitude–phase MLP. The lower strip gives the independently calculated excitation gap from the target-singlet ground state to the first wrong-spin state. (b) FCI errors for I 2 at R=2.666A˚ with explicit spin–orbit coupling under joint amplitude–phase optimization. (c) Batch-size dependence of fixed-amplitude phase optimization for I 2 at R=2.666A˚ with explicit spin–orbit coupling. The upper and lower panels show the FCI error and phase infidelity 1−Fphase , respectively. (d) FCI errors for fixed-amplitude phase-learning controls on H 2 O at three geometries and LiH. (e) Fraction of paired runs within chemical accuracy for C 2 at R=1.24A˚ across MLP widths and batch sizes. Each cell reports success counts in the order standard → direct/AMVP; the shading gives the mean direct–AMVP success fraction. In (b,d), the left-to-right order is standard, AMVP, direct. Small symbols denote individual runs; large symbols show medians; shaded bands and vertical bars indicate interquartile ranges.
Figure 4: Matched-state phase-gradient statistics in quantum many-body physics and molecular calculations. (a) Variance ratios Vstd/Vm for the direct and AMVP estimators along the training process of an 8×8 flux torus. (b) Shared-norm batch SNR for the same flux-torus states, using an identical debiased population-signal estimate for all estimators. (c) Exact phase-gradient norm for the N 2 fixed-amplitude phase-channel diagnostic at R=1.70A˚ . (d) Rows give the estimator that produced the state, columns the checkpoint in thousands of updates; the gains do not depend on which estimator produced the state, showing that the variance difference is a property of the estimator itself rather than of the states. (e) Corresponding gains in descent-aligned SNR relative to the standard estimator. (f) Exact-eigenstate control for the periodic honeycomb J1 – J2 model. At the higher-energy eigenstate, the standard score-function phase-force estimate vanishes for every finite batch, and AMVP consequently selects λ∗=0 . The direct estimator instead has nonzero finite-batch realizations despite having the same zero population mean, allowing a stochastic displacement toward lower energy. (g,h) Evolution from the same excited eigenstate: energy decrease (g) and exact phase-force norm (h). AMVP esc denotes the escape-aware extension with a temporary lower bound on the direct mixing weight. Curves and bands in (a,b,c,g,h) show medians and interquartile ranges; heatmap cells in (d,e) show medians.
Variance reduction develops methods that make Monte Carlo estimates, gradients, and experimental comparisons more precise with limited data or computation.
Complex neural quantum states are difficult to optimize when their wavefunction phase carries gauge, chiral, fermionic, or topological structure. We show that the major failure mode is not only ansatz expressivity, but the Monte Carlo estimator used to learn this phase. For separated amplitude-phase states, differentiating the local energy at fixed samples gives a different unbiased estimator of the same variational Monte Carlo phase force, without changing the objective. We further extend the construction to coupled two-head networks by keeping the amplitude-gradient contribution and applying the direct derivative only to the phase path. An adaptive minimum-variance mixture interpolates between standard and direct estimators during training. Across flux ladders, chiral chains, two-dimensional flux cylinders, an interacting fermion ladder, shared-network controls, and a fractional quantum Hall benchmark, the resulting estimators reduce phase-gradient variance, suppress seed failures, and often move multi-percent standard-gradient plateaus to sub-percent accuracy.
Yi-Ran Xue, Rui Wang, Baigeng Wang +1
*National Laboratory of Solid State Microstructures and Department of Physics, Nanjing University, Nanjing 210093, China. · Department of Physics, University of Massachusetts, Amherst, Massachusetts 01003, USA. · Collaborative Innovation Center of Advanced Microstructures, Nanjing University, Nanjing 210093, China. +3
Neural quantum states offer expressive representations of quantum many-body wave functions, yet their practical accuracy can be limited by stochastic optimization rather than representational capacity. Here we identify a finite-sample instability, termed subspace trapping, in which physically important configurations become strongly underestimated, remain absent from successive sampling batches and receive insufficient gradient feedback. This self-reinforcing loss of sampled support can confine optimization to an effective subspace and produce apparently stationary states above the true ground state energy. To address this problem, we introduce annealed gradient descent (AGD), a sampling-aware update with annealing factor that temporarily increases the relative contribution of sampled low-probability configurations while limiting the dominance of high-probability ones. We establish the connection between finite-sample support loss and effective subspace optimization, and then evaluate the method across molecular systems, one and two-dimensional J1-J2 models. Annealed gradient descent suppresses metastable trapping, preserves physically relevant configurations and enables compact neural quantum states to attain chemical accuracy and competitive state-of-the-art performance. These results establish AGD as a lightweight complement to expressive neural architectures, improved sampling strategies for scalable quantum many-body optimization.
Shiwei Zhou, Yiming Huang, Xiao Yuan +1
Institute of High Energy Physics, Chinese Academy of Sciences, Beijing 100049, China · China Center of Advanced Science and Technology, Beijing 100190, China · Center on Frontiers of Computing Studies, School of Computer Science, Peking University, Beijing 100871, China
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 J1-J2 one- and two-dimensional spin systems, PWO improves stability and wall-clock convergence over Adam, minSR, and SPRING. Finally, we fine-tune a 1.5B-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.
Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes +4
Mila Quebec AI Institute, Université de Montréal, Montréal, Canada · Applied Quantum Algorithms ⟨aQaL⟩, LIACS & LION, Leiden University, Leiden, Netherlands · QuTech and Kavli Institute of Nanoscience, Delft University of Technology, Delft, Netherlands +1