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

TL;DR

This paper introduces SPRING, a novel optimizer inspired by Kaczmarz, which accelerates the training of neural network wavefunctions in quantum chemistry, outperforming existing methods in efficiency and accuracy.

Contribution

The paper presents SPRING, a new optimizer combining stochastic reconfiguration and Kaczmarz methods, significantly improving the speed and accuracy of neural network wavefunction optimization.

Findings

SPRING outperforms MinSR and KFAC in small atom and molecule tests.

SPRING achieves chemical accuracy on oxygen atom with fewer iterations.

SPRING demonstrates faster convergence compared to existing optimizers.

Abstract

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.