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

One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective
Jul 2, 2026Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes +4Neural Quantum States Second-Order Optimization
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
Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo
May 13, 2026Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel +2Neural Quantum States Variational Monte Carlo
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
Enhanced Neural Quantum State via Annealed Gradient Descent
Jul 21, 2026Shiwei Zhou, Yiming Huang, Xiao Yuan +1Neural Quantum States Gradient Descent
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

Jul 2, 2026J/K move · Enter open · S save

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

Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes, Eliška Greplová+3

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

Open paper AskPDF arXiv

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.

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