Solving the electronic Schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks

Abstract

Accurate numerical solutions for the Schrödinger equation are of utmost importance in quantum chemistry. However, the computational cost of current high-accuracy methods scales poorly with the number of interacting particles. Combining Monte Carlo methods with unsupervised training of neural networks has recently been proposed as a promising approach to overcome the curse of dimensionality in this setting and to obtain accurate wavefunctions for individual molecules at a moderately scaling computational cost. These methods currently do not exploit the regularity exhibited by wavefunctions with respect to their molecular geometries. Inspired by recent successful applications of deep transfer learning in machine translation and computer vision tasks, we attempt to leverage this regularity by introducing a weight-sharing constraint when optimizing neural network-based models for different molecular geometries. That is, we restrict the optimization process such that up to 95 percent of weights in a neural network model are in fact equal across varying molecular geometries. We find that this technique can accelerate optimization when considering sets of nuclear geometries of the same molecule by an order of magnitude and that it opens a promising route towards pre-trained neural network wavefunctions that yield high accuracy even across different molecules.

Wave function Ansatz (but Periodic) Networks and the Homogeneous Electron Gas

“…In testing and in other works [2,24,26], KFAC shows clear advantages in optimisation over more standard techniques such as Adam [3,54]. However, the algorithm requires careful balancing of the hyperparameters (learning rate, damping, and norm constraint) to achieve optimal performance, which consumes time in development and tweaking of model design.…”

Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions

“…To investigate PESNet's accuracy and training time benefit, we compare it to FermiNet , PauliNet (Hermann et al, 2020), and DeepErwin (Scherbela et al, 2021) on diverse systems ranging from 3 to 28 electrons. Note, the concurrently developed DeepErwin was only recently released as a pre-print and still requires separate models and training for each configuration.…”

“…This procedure is done for the hydrogen rectangle, the hydrogen chain, and the nitrogen molecule. For H + 4 and cyclobutadiene, we train on discrete sets of geometries from the literature (Scherbela et al, 2021;Kinal & Piecuch, 2007).…”

“…To address this issue, weight-sharing has been proposed to reduce the time per training, but this was initially limited to non-fermionic systems (Yang et al, 2020). In a concurrent work, Scherbela et al (2021) extend this idea to electronic wave functions. However, their DeepErwin model still requires separate models for each geometry, does not account for symmetries and achieves lower accuracy, as we show in Section 4.…”

“…Here, we study the training times for all systems from the previous section. We compare the official JAX implementation of FermiNet , the official PyTorch implementation of PauliNet (Hermann et al, 2020), the official TensorFlow implementation of DeepErwin (Scherbela et al, 2021), and our JAX (Bradbury et al, 2018) implementation of PESNet. We use the same hyperparameters as in the experiments or the defaults from the respective works.…”