Stochastic Reconfiguration as Statistical Filtering for Overparameterized Neural Quantum States

Tak Hur

quant-phSep 20, 2026

Stochastic Reconfiguration as Statistical Filtering for Overparameterized Neural Quantum States

Authors: Tak Hur

Organizations: Department of Applied Mathematics, University of Waterloo 200 University Avenue West, Waterloo, Ontario N2L 3G1, Canada

Abstract

Stochastic reconfiguration (SR) is the standard optimizer for neural quantum states (NQS), but modern NQS often have far more parameters than Monte Carlo samples. We show that in this regime the diagonal shift is more than a numerical stabilizer. It acts as a statistical filter for finite-sample generalization. At a fixed wave function, SR is ridge regression from tangent features to the centered local energy. Its residual is the expressivity gap, the part of imaginary-time evolution outside the current tangent space. This gap is orthogonal to the tangent space in population, but finite batches make it act as noise that SR can overfit. The shift therefore balances shrinkage of useful update directions against variance from fitting sampled residuals. Exact diagnostics on a $4\times4$ Heisenberg graph separate two effects of overparameterization. Larger tangent spaces help when they reduce the expressivity gap, but they can hurt when they overfit a fixed gap. In a $100$ -site transverse-field Ising family trained with a foundation NQS, validation risk is U-shaped in the shift while variance decreases, matching the noisy-ridge model. This view leads to multi-shift SR (MS-SR), which averages independent ridge solves at data-adaptive shifts to form a richer, lower-variance spectral filter. Checkpoint-local experiments show that MS-SR lowers validation risk and update variance relative to the fixed-shift SR baseline. We further compare MS-SR and SR in paired online training continuations, with independent endpoint energy evaluations and a separate update-cost benchmark.

Figures & tables

Figure 1: Historical shift in the parameter-to-sample ratio used by VMC/NQS calculations. The dashed line marks

P=N_{s}

. Marker shape denotes domain: circles are spin/lattice NQS and diamonds are electronic-structure NQS. Sources: [ 2 , 21 , 20 , 4 , 32 , 10 , 17 , 8 , 30 , 18 , 3 , 26 ] .

Figure 2: Exact regression diagnostics on the $4\times 4$ Heisenberg graph specified in Appendix B , at $N_{s}=4096$ and $\lambda=10^{-9}$ . Panels compare empirical and population-optimal amplitude tangent predictions. Gap variance and excess risk refer to this regression. Matched noise exposes ViT overfitting at comparable residual variance; matched state shows lower gap variance and excess risk for the ViT.

Figure 3: Large-scale shared-checkpoint test of the noisy-ridge mechanism on the periodic

L=100

transverse-field Ising model with

h

ranging from

0.8

1.2

with foundational NQS. A fixed set of late-training checkpoints is reused while only the solve-time shift

\lambda

is varied.

\widehat{\mathcal{E}}_{\mathrm{val}}

is U-shaped, while

\mathcal{V}_{\mathrm{mb}}

decreases monotonically with

\lambda

, matching the two predictions of Section 2 .

Figure 4: Checkpoint-local ablations for the ingredients of MS-SR at

K=4

. Paired bars show raw

\hat{\mathcal{E}}_{\mathrm{val}}

and

\mathcal{V}_{\mathrm{mb}}

averaged over fixed source checkpoints with one standard error. Standard SR uses

\lambda_{\mathrm{train}}=10^{-4}

. Both systems show bagged SR at

\lambda_{\mathrm{train}}

and at

\lambda_{\mathrm{tuned}}

, selected per checkpoint by an eight-point bagged-SR residual sweep using the evaluation data. All multi-shift variants use the same NTK quantile grid.

Figure 5: Online performance results for MS-SR against SR. (a) Five paired continuations from a single shared SR checkpoint. Curves show means of binned traces with

\pm 1

across-run sample standard deviation. The nominal horizontal axis counts candidate SR updates: one per SR step and four per MS-SR step. Diamonds at 500 mark independent endpoint means. (b) Endpoint energies for the same five paired continuations as in (a). Bold marks the lower energy.

\Delta=(E_{\mathrm{MS\text{-}SR}}-E_{\mathrm{SR}})/N

. Individual

\pm

values are combined Monte Carlo standard errors in units of

10^{-5}

. (c) A separate TFIM benchmark on one H100 80GB GPU at

K=4

, showing medians over three paired device/seed settings for the NTK-quantile MS-SR implementation. Protocols and computational accounting are given in Appendices F.1 and C.4 .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

\lambda_{j}\leftarrow\operatorname{Quantile}(\{s_{i}\},q_{j}),\qquad j=1,\dots,K.

Appendix

Algorithm 1 Multi-Shift SR (MS-SR): one optimization step

TFIM update rule	$10^{3}\widehat{\mathcal{E}}_{\mathrm{val}}$	$10^{3}\mathcal{V}_{\mathrm{mb}}$
Single SR, $\lambda_{\mathrm{single}}$	$9.70\pm 2.43$	$4.71\pm 1.24$
Bagged SR, $\lambda_{\mathrm{train}}$	$9.33\pm 2.05$	$0.61\pm 0.13$
Bagged SR, $\lambda_{\mathrm{tuned}}$	$6.17\pm 1.49$	$1.24\pm 0.48$
MS-SR, NTK quantiles	$6.80\pm 1.59$	$1.39\pm 0.29$

Appendix

Table 1: TFIM baseline comparison: mean

\pm

standard deviation across ten checkpoints (20 resampling repeats per checkpoint). Bagged SR and MS-SR use four independent batches and four solves per candidate shift. Here

\lambda_{\mathrm{train}}=10^{-4}

and

\lambda_{\mathrm{single}}=10^{-6}

throughout. Bagged-residual tuning selects

10^{-7}

at step

1100

and

10^{-6}

elsewhere. Both tuned references use the evaluation data for selection; the additional shift-search cost is excluded from the four-solve count.

Figure 6: Fixed-shift bagged SR on TFIM (

K=4

). Points are means over ten checkpoint-level diagnostics; error bars are standard errors over those checkpoints. The dashed line and shaded band show the MS-SR mean and standard error on the same checkpoints. The training shift is

10^{-4}

. Minimizing the bagged residual selects

10^{-7}

at step

1100

and

10^{-6}

at the other nine checkpoints. The complete sweep and checkpoint standard deviations are retained in the released summary data.

Method	$K=2$	$K=3$	$K=4$	$K=5$
MS-SR	$8.54/2.28$	$7.16/1.74$	$6.80/1.39$	$6.38/1.27$
Bagged SR, $\lambda_{\mathrm{train}}$	$9.93/1.22$	$9.53/0.82$	$9.33/0.61$	$9.21/0.49$
Bagged SR, $\lambda_{\mathrm{tuned}}$	$7.33/2.34$	$6.56/1.56$	$6.17/1.17$	$5.93/0.94$

Appendix

Table 2: Sensitivity to

K

on TFIM. Each entry is

\widehat{\mathcal{E}}_{\mathrm{val}}/\mathcal{V}_{\mathrm{mb}}

, with both values in units of

10^{-3}

, averaged over ten checkpoint-level diagnostics. Here

\lambda_{\mathrm{train}}=10^{-4}

and

\lambda_{\mathrm{tuned}}=10^{-6}

. The sampling and solve budget grows linearly with

K

Quantile grid	$\widehat{\mathcal{E}}_{\mathrm{val}}$	$\mathcal{V}_{\mathrm{mb}}$
$\{0.9,0.7,0.4,0.1\}$ (default)	$6.802\pm 1.595$	$1.389\pm 0.291$
$\{0.95,0.8,0.5,0.2\}$	$6.938\pm 1.625$	$1.469\pm 0.283$
$\{0.8,0.6,0.3,0.05\}$	$6.574\pm 1.568$	$1.369\pm 0.321$

Appendix

Table 3: Sensitivity to the

K=4

NTK quantile grid on TFIM. Values are in units of

10^{-3}

and show means

\pm

standard deviations across ten checkpoint-level diagnostics, using

20

resampling repeats per checkpoint.

$\lambda_{\mathrm{train}}$	Method	$\widehat{\mathcal{E}}_{\mathrm{val}}$	$\mathcal{V}_{\mathrm{mb}}$
$10^{-3}$	Single SR, $\lambda_{\mathrm{single}}$	$36.832\pm 6.731$	$16.348\pm 1.470$
	Bagged SR, $\lambda_{\mathrm{single}}$	$24.590\pm 6.990$	$4.063\pm 0.369$
	MS-SR	$24.726\pm 4.623$	$5.384\pm 0.759$
$10^{-5}$	Single SR, $\lambda_{\mathrm{single}}$	$2.107\pm 0.621$	$0.700\pm 0.179$
	Bagged SR, $\lambda_{\mathrm{single}}$	$1.580\pm 0.482$	$0.174\pm 0.044$
	MS-SR	$1.575\pm 0.470$	$0.256\pm 0.060$

Appendix

Table 4: TFIM diagnostics on trajectories trained with alternative diagonal shifts. Values are in units of

10^{-3}

and show means

\pm

standard deviations across the three checkpoint-level diagnostics. Bagged SR and MS-SR use

K=4

Experiment	Manuscript use	Worker	Wall time	GPU-hours
4 by 4 small experiments	Figures 2 and 7	$1\times$ H100 80GB GPU	$\approx 30$ h	$\approx 30$
TFIM-FNQS training	Figure 3 and 4 checkpoints	$8\times$ H100 80GB GPUs	$\approx 1.9$ h	$\approx 15$
TFIM single-shift sweep	Figure 3	$8\times$ H100 80GB GPUs	$\approx 25.5$ h	$\approx 204$
$J_{1}$ – $J_{2}$ ViT training	Figure 4 checkpoints	$8\times$ H100 80GB GPUs	$\approx 15.1$ h	$\approx 121$
$J_{1}$ – $J_{2}$ MS-SR ablations	Figure 4	$8\times$ H100 80GB GPUs	$\approx 46$ h	$\approx 368$
TFIM MS-SR ablations	Figure 4	$8\times$ H100 80GB GPUs	$\approx 3.0$ h	$\approx 24$

Appendix

Table 5: Compute-resource accounting for the original experiment campaigns. The listed entries sum to approximately

762

H100 GPU-hours; this is a historical subtotal rather than the total for all camera-ready experiments.

Figure 7: Online update tracking in exact

4\times 4

systems. SR, SPRING, and MS-SR are trained separately from the same initialization. At each step we compare the stochastic update with a full-summation population SR reference at shift

10^{-4}

. We plot the full-trajectory log-ratio

\Delta_{\log r}:=\log_{10}(\overline{r}_{\mathrm{MS}\text{-}\mathrm{SR}}/\min\{\overline{r}_{\mathrm{SR}},\overline{r}_{\mathrm{SPRING}}\})

. Negative values favor MS-SR. Boxes aggregate random seeds, sample sizes, and learning rates for each Hamiltonian–ansatz pair. Iteration counts are matched; MS-SR uses four candidate solves plus spectrum and weight-fit overhead. This variant smooths its stacking weights across iterations.

System	RBM, $\alpha=2$	RBM, $\alpha=4$	ViT
TFIM	$4/10$	$8/11$	$12/12$
$J_{1}$ - $J_{2}$	$11/11$	$10/10$	$8/9$
SSM	$12/12$	$10/11$	$9/10$

Appendix

Table 6: Per-panel counts for Figure 7 . Each entry is the number of configurations with lower MS-SR relative QGT-norm error than both SR and SPRING, divided by the number of completed matched configurations in that panel. The denominators sum to

96

Overparameterized Neural Network

See all 13 papers

Jul 2, 2026

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.

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

Neural Quantum States Second-Order Optimization Policy Gradient Methods Trust-Region Optimization Quantum Many-Body Physics

Ask PDF

May 13, 2026

Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo

Neural-network quantum states have emerged as a powerful variational framework for quantum many-body systems, with recent progress often driven by massively parallel architectures such as transformers. Recurrent neural network quantum states, however, are frequently regarded as intrinsically sequential and therefore less scalable. Here we revisit this view by showing that modern recurrent architectures can support fast, accurate, and computationally accessible neural quantum state simulations. Using autoregressive recurrent wave functions together with recent advances in parallelizable recurrence, we develop variational ansätze, called parallel scan recurrent neural quantum states (PSR-NQS), which can be trained efficiently within variational Monte Carlo in one and two spatial dimensions. We demonstrate accurate benchmark results and show that, with iterative retraining, our approach reaches two-dimensional spin lattices as large as

52\times52

while remaining in agreement with available quantum Monte Carlo data. Our results establish recurrent architectures as a practical and promising route toward scalable neural quantum state simulations with modest computational resources.

Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel +2

Department of Physics and Astronomy, University of California, Davis, California 95616, USA · Department of Applied Mathematics, University of Waterloo, Waterloo, ON N2L 3G1, Canada · Vector Institute, Toronto, Ontario, M5G 0C6, Canada +1

Neural Quantum States Variational Monte Carlo Recurrent Neural Networks Quantum Many-Body Physics

Ask PDF

Jul 21, 2026

Enhanced Neural Quantum State via Annealed Gradient Descent

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

J_1

J_2

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 Gradient Descent Quantum Many-Body Physics

Ask PDF

Abstract

Overparameterized Neural Network

Explore similar work

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

Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo

Enhanced Neural Quantum State via Annealed Gradient Descent