New Insights and Algorithms for Training Neural Network Wavefunctions

Neural network wavefunctions have shown impressive results for simulating strongly correlated quantum systems, but they are difficult to train due to ill-conditioned and non-convex loss landscapes. We describe a new theoretical framework for understanding these optimization challenges, analyzing existing algorithms, and developing new ones. We use the framework to provide the first analysis of natural gradient (aka stochastic reconfiguration or SR) and Gauss-Newton methods in the practical regime when the number of parameters greatly exceeds the size of the mini-batch. Our analysis explains the recently proposed Subsampled PRojected-Increment Natural Gradient descent (SPRING) algorithm and suggests new ways to regularize and accelerate the subsampled Gauss-Newton method. It also reveals a surprising effect wherein Gauss-Newton methods only provide benefits when the mini-batch size surpasses a certain problem-dependent threshold. A key piece of our theoretical framework is a new model problem that has the potential to be broadly useful as a proving ground for developing, analyzing, and testing new optimizers for neural network wavefunctions and related problems.