# A Kaczmarz-inspired approach to accelerate the optimization of neural network wavefunctions

arXiv: https://arxiv.org/abs/2401.10190 · HTML: https://mlsanity.com/paper/2401.10190/

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 43 references that are in the map.
- First posted on arXiv 18 Jan 2024.

## Builds on (most cited first, 12 of 12)

- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [Deep neural network solution of the electronic Schr"odinger equation](https://mlsanity.com/paper/1909.08423/) (2019)
- [Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks](https://mlsanity.com/paper/1905.10961/) (2019)
- [A Self-Attention Ansatz for Ab-initio Quantum Chemistry](https://mlsanity.com/paper/2211.13672/) (2022)
- [Efficient Subsampled Gauss-Newton and Natural Gradient Methods for Training Neural Networks](https://mlsanity.com/paper/1906.02353/) (2019)
- [Gradient Descent on Neurons and its Link to Approximate Second-Order Optimization](https://mlsanity.com/paper/2201.12250/) (2022)
- [Gold-standard solutions to the Schr"odinger equation using deep learning: How much physics do we need?](https://mlsanity.com/paper/2205.09438/) (2022)
- [Discovering Quantum Phase Transitions with Fermionic Neural Networks](https://mlsanity.com/paper/2202.05183/) (2022)
- [Geometry of learning neural quantum states](https://mlsanity.com/paper/1910.11163/) (2019)
- [Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation](https://mlsanity.com/paper/2112.03491/) (2021)
- [Rayleigh-Gauss-Newton optimization with enhanced sampling for variational Monte Carlo](https://mlsanity.com/paper/2106.10558/) (2021)

## Built on by (most cited first, 26 of 26)

- [Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization](https://mlsanity.com/paper/2505.12149/) (2025)
- [Principled Approaches for Extending Neural Architectures to Function Spaces for Operator Learning](https://mlsanity.com/paper/2506.10973/) (2025)
- [Curvature-Aware Optimization for High-Accuracy Physics-Informed Neural Networks](https://mlsanity.com/paper/2604.05230/) (2026)
- [Neural Pfaffians: Solving Many Many-Electron Schr"odinger Equations](https://mlsanity.com/paper/2405.14762/) (2024)
- [Accurate Ab-initio Neural-network Solutions to Large-Scale Electronic Structure Problems](https://mlsanity.com/paper/2504.06087/) (2025)
- [Functional Neural Wavefunction Optimization](https://mlsanity.com/paper/2507.10835/) (2025)
- [Worth Their Weight: Randomized and Regularized Block Kaczmarz Algorithms without Preprocessing](https://mlsanity.com/paper/2502.00882/) (2025)
- [A Sketch-and-Project Analysis of Subsampled Natural Gradient Algorithms](https://mlsanity.com/paper/2508.21022/) (2025)
- [Spin-Adapted Neural Network Wavefunctions in Real Space](https://mlsanity.com/paper/2511.01671/) (2025)
- [Natural gradient descent with momentum](https://mlsanity.com/paper/2604.15554/) (2026)
- [Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State](https://mlsanity.com/paper/2603.14515/) (2026)
- [A Dirac-Frenkel-Onsager principle: Instantaneous residual minimization with gauge momentum for nonlinear parametrizations of PDE solutions](https://mlsanity.com/paper/2605.00284/) (2026)
- [Stochastic Reconfiguration as Statistical Filtering for Overparameterized Neural Quantum States](https://mlsanity.com/paper/2609.23334/) (2026)
- [Ground state phases of the two-dimension electron gas with a unified variational approach](https://mlsanity.com/paper/2405.19397/) (2024)
- [Application of Langevin Dynamics to Advance the Quantum Natural Gradient Optimization Algorithm](https://mlsanity.com/paper/2409.01978/) (2024)
- [Neural-Quantum-States Impurity Solver for Quantum Embedding Problems](https://mlsanity.com/paper/2509.12431/) (2025)
- [Fermionic neural Gibbs states](https://mlsanity.com/paper/2512.04663/) (2025)
- [PROMA: Projected Microbatch Accumulation for Reference-Free Proximal Policy Updates](https://mlsanity.com/paper/2601.10498/) (2026)
- [Predicting magnetism with first-principles AI](https://mlsanity.com/paper/2602.09093/) (2026)
- [Randomized Advantage Transformation (RAT): Computing Natural Policy Gradients via Direct Backpropagation](https://mlsanity.com/paper/2605.18591/) (2026)
- [Dirac-Frenkel dynamics with inertia for nonlinearly parametrized solutions of evolution problems](https://mlsanity.com/paper/2606.24769/) (2026)
- [One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective](https://mlsanity.com/paper/2607.02292/) (2026)
- [Enhanced Neural Quantum State via Annealed Gradient Descent](https://mlsanity.com/paper/2607.18865/) (2026)
- [Energy Manifold Natural Gradient Descent: Riemannian Optimization for Neural PDE Solvers](https://mlsanity.com/paper/2607.22004/) (2026)
- [Backpropagated Output Momentum: Relocating Optimizer History from Parameters to Task Space](https://mlsanity.com/paper/2609.36738/) (2026)
- [NucleoNet: A Nuclear-Aware Determinant Ansatz for Multi-Geometry Neural Variational Monte Carlo](https://mlsanity.com/paper/2610.12985/) (2026)

## Abstract (arXiv)

Neural network wavefunctions optimized using the variational Monte Carlo method have been shown to produce highly accurate results for the electronic structure of atoms and small molecules, but the high cost of optimizing such wavefunctions prevents their application to larger systems. We propose the Subsampled Projected-Increment Natural Gradient Descent (SPRING) optimizer to reduce this bottleneck. SPRING combines ideas from the recently introduced minimum-step stochastic reconfiguration optimizer (MinSR) and the classical randomized Kaczmarz method for solving linear least-squares problems. We demonstrate that SPRING outperforms both MinSR and the popular Kronecker-Factored Approximate Curvature method (KFAC) across a number of small atoms and molecules, given that the learning rates of all methods are optimally tuned. For example, on the oxygen atom, SPRING attains chemical accuracy after forty thousand training iterations, whereas both MinSR and KFAC fail to do so even after one hundred thousand iterations.