DOE PAGES title logo U.S. Department of Energy
Office of Scientific and Technical Information

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

Journal Article · · Journal of Computational Physics
ORCiD logo [1];  [2];  [3]
  1. University of California, Berkeley, CA (United States)
  2. University of California, Berkeley, CA (United States). Simons Institute for the Theory of Computing
  3. University of California, Berkeley, CA (United States); Lawrence Berkeley National Laboratory (LBNL), Berkeley, CA (United States)

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.

Research Organization:
Lawrence Berkeley National Laboratory (LBNL), Berkeley, CA (United States)
Sponsoring Organization:
USDOE; USDOE Office of Science (SC), Advanced Scientific Computing Research (ASCR)
Grant/Contract Number:
AC02-05CH11231; SC0023112
OSTI ID:
2434086
Journal Information:
Journal of Computational Physics, Journal Name: Journal of Computational Physics Vol. 516; ISSN 0021-9991
Publisher:
ElsevierCopyright Statement
Country of Publication:
United States
Language:
English

References (31)

Zur Quantentheorie der Molekeln journal January 1927
A Randomized Kaczmarz Algorithm with Exponential Convergence journal April 2008
Randomized Kaczmarz solver for noisy linear systems journal April 2010
Sketch-Based Empirical Natural Gradient Methods for Deep Learning journal August 2022
Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation journal November 2022
Paved with good intentions: Analysis of a randomized block Kaczmarz method journal January 2014
NetKet: A machine learning toolkit for many-body quantum systems journal July 2019
A Blocked Linear Method for Optimizing Large Parameter Sets in Variational Monte Carlo journal May 2017
Ab initio calculation of real solids via neural network ansatz journal December 2022
Deep-neural-network solution of the electronic Schrödinger equation journal September 2020
Neural-network quantum states for ultra-cold Fermi gases journal May 2024
An accurate analytic potential function for ground-state N2 from a direct-potential-fit analysis of spectroscopic data journal October 2006
Optimization of quantum Monte Carlo wave functions by energy minimization journal February 2007
Weak binding between two aromatic rings: Feeling the van der Waals attraction by quantum Monte Carlo methods journal July 2007
Multiconfiguration wave functions for quantum Monte Carlo calculations of first‐row diatomic molecules journal July 1996
An accelerated linear method for optimizing non-linear wavefunctions in variational Monte Carlo journal January 2020
DeepQMC: An open-source software suite for variational optimization of deep-learning molecular wave functions journal September 2023
Ground-state correlation energies for atomic ions with 3 to 18 electrons journal May 1993
Message-passing neural quantum states for the homogeneous electron gas journal July 2024
Generalized Lanczos algorithm for variational quantum Monte Carlo journal June 2001
Optimizing large parameter sets in variational quantum Monte Carlo journal January 2012
Discovering Quantum Phase Transitions with Fermionic Neural Networks journal January 2023
Optimization of Ground- and Excited-State Wave Functions and van der Waals Clusters journal July 2001
Alleviation of the Fermion-Sign Problem by Optimization of Many-Body Wave Functions journal March 2007
Geometry of learning neural quantum states journal May 2020
Ab initio solution of the many-electron Schrödinger equation with deep neural networks journal September 2020
Rayleigh-Gauss-Newton optimization with enhanced sampling for variational Monte Carlo journal August 2022
Quantum Monte Carlo simulations of solids journal January 2001
Solving the quantum many-body problem with artificial neural networks journal February 2017
Randomized Extended Kaczmarz for Solving Least Squares journal January 2013
Natural Gradient Works Efficiently in Learning journal February 1998