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

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.

arXiv commentFinal proof for Journal of Computational Physics

References in corpus (19)

• Solving the Quantum Many-Body Problem with Artificial Neural Networks
• Deep neural network solution of the electronic Schrödinger equation
• Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks
• Generalized Lanczos Algorithm for Variational Quantum Monte Carlo
• Paved with Good Intentions: Analysis of a Randomized Block Kaczmarz Method
• Randomized Extended Kaczmarz for Solving Least-Squares
• Optimization of quantum Monte Carlo wave functions by energy minimization
• Weak binding between two aromatic rings: feeling the van der Waals attraction by quantum Monte Carlo methods
• Randomized Kaczmarz solver for noisy linear systems
• Optimizing large parameter sets in variational quantum Monte Carlo
• NetKet: A Machine Learning Toolkit for Many-Body Quantum Systems
• Optimization of ground and excited state wavefunctions and van der Waals clusters
• Discovering Quantum Phase Transitions with Fermionic Neural Networks
• Ab initio calculation of real solids via neural network ansatz
• Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks
• Message-Passing Neural Quantum States for the Homogeneous Electron Gas
• DeepQMC: an open-source software suite for variational optimization of deep-learning molecular wave functions
• An accelerated linear method for optimizing non-linear wavefunctions in variational Monte Carlo
• Rayleigh-Gauss-Newton optimization with enhanced sampling for variational Monte Carlo

Cited by in corpus (8)

• Ground state phases of the two-dimension electron gas with a unified variational approach
• Second-order optimisation strategies for neural network quantum states
• Design principles of deep translationally-symmetric neural quantum states for frustrated magnets
• Improving neural network performance for solving quantum sign structure
• Simple Fermionic backflow states via a systematically improvable tensor decomposition
• Expressivity of determinantal ansatzes for neural network wave functions
• Assessing Orbital Optimization in Variational and Diffusion Monte Carlo
• Application of Langevin Dynamics to Advance the Quantum Natural Gradient Optimization Algorithm