Towards a Foundation Model for Neural Network Wavefunctions

Deep neural networks have become a highly accurate and powerful wavefunction ansatz in combination with variational Monte Carlo methods for solving the electronic Schr\"odinger equation. However, despite their success and favorable scaling, these methods are still computationally too costly for wide adoption. A significant obstacle is the requirement to optimize the wavefunction from scratch for each new system, thus requiring long optimization. In this work, we propose a novel neural network ansatz, which effectively maps uncorrelated, computationally cheap Hartree-Fock orbitals, to correlated, high-accuracy neural network orbitals. This ansatz is inherently capable of learning a single wavefunction across multiple compounds and geometries, as we demonstrate by successfully transferring a wavefunction model pre-trained on smaller fragments to larger compounds. Furthermore, we provide ample experimental evidence to support the idea that extensive pre-training of a such a generalized wavefunction model across different compounds and geometries could lead to a foundation wavefunction model. Such a model could yield high-accuracy ab-initio energies using only minimal computational effort for fine-tuning and evaluation of observables.

[1]  Stephan Gunnemann,et al.  Generalizing Neural Wave Functions , 2023, ICML.

[2]  Weizhong Fu,et al.  Towards the ground state of molecules via diffusion Monte Carlo on neural networks , 2022, Nature Communications.

[3]  F. Noé,et al.  Electronic excited states in deep variational Monte Carlo , 2022, Nature Communications.

[4]  W. Foulkes,et al.  Discovering Quantum Phase Transitions with Fermionic Neural Networks , 2022, Physical review letters.

[5]  Ingrid von Glehn,et al.  A Self-Attention Ansatz for Ab-initio Quantum Chemistry , 2022, ICLR.

[6]  Weiluo Ren,et al.  Interatomic force from neural network based variational quantum Monte Carlo. , 2022, The Journal of chemical physics.

[7]  Stephan Gunnemann,et al.  Sampling-free Inference for Ab-Initio Potential Energy Surface Networks , 2022, ICLR.

[8]  P. Marquetand,et al.  Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need? , 2022, NeurIPS.

[9]  Simon L. Batzner,et al.  The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials , 2022, ArXiv.

[10]  Lisa Anne Hendricks,et al.  Training Compute-Optimal Large Language Models , 2022, ArXiv.

[11]  Ji Chen,et al.  Ab initio calculation of real solids via neural network ansatz , 2022, Nature communications.

[12]  A. Bhowmik,et al.  Neural network ansatz for periodic wave functions and the homogeneous electron gas , 2022, Physical Review B.

[13]  Stephan Günnemann,et al.  Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions , 2021, ICLR.

[14]  P. Marquetand,et al.  Solving the electronic Schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks , 2021, Nature Computational Science.

[15]  Lu Yuan,et al.  Florence: A New Foundation Model for Computer Vision , 2021, ArXiv.

[16]  Michael Gastegger,et al.  SE(3)-equivariant prediction of molecular wavefunctions and electronic densities , 2021, NeurIPS.

[17]  N. Tubman,et al.  Simulations of state-of-the-art fermionic neural network wave functions with diffusion Monte Carlo , 2021, 2103.12570.

[18]  Ilya Sutskever,et al.  Learning Transferable Visual Models From Natural Language Supervision , 2021, ICML.

[19]  David Pfau,et al.  Better, Faster Fermionic Neural Networks , 2020, ArXiv.

[20]  Mark Chen,et al.  Language Models are Few-Shot Learners , 2020, NeurIPS.

[21]  Kristof T. Schütt,et al.  A deep neural network for molecular wave functions in quasi-atomic minimal basis representation. , 2020, The Journal of chemical physics.

[22]  Timothy C. Berkelbach,et al.  Recent developments in the PySCF program package. , 2020, The Journal of chemical physics.

[23]  F. Noé,et al.  Deep-neural-network solution of the electronic Schrödinger equation , 2019, Nature Chemistry.

[24]  David Pfau,et al.  Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks , 2019, Physical Review Research.

[25]  Kristof T. Schütt,et al.  Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions , 2019, Nature Communications.

[26]  R. Walker,et al.  Publisher , 2019, Definitions.

[27]  E Weinan,et al.  Solving many-electron Schrödinger equation using deep neural networks , 2018, J. Comput. Phys..

[28]  B. Clark,et al.  Variational optimization in the AI era: Computational Graph States and Supervised Wave-function Optimization , 2018, 1811.12423.

[29]  David M. Ceperley,et al.  Towards the solution of the many-electron problem in real materials: equation of state of the hydrogen chain with state-of-the-art many-body methods , 2017, 1705.01608.

[30]  Matthias Troyer,et al.  Solving the quantum many-body problem with artificial neural networks , 2016, Science.

[31]  Roger B. Grosse,et al.  Optimizing Neural Networks with Kronecker-factored Approximate Curvature , 2015, ICML.

[32]  Paul G. Mezey,et al.  A fast intrinsic localization procedure applicable for ab initio and semiempirical linear combination of atomic orbital wave functions , 1989 .

[33]  W. K. Hastings,et al.  Monte Carlo Sampling Methods Using Markov Chains and Their Applications , 1970 .

[34]  S. F. Boys,et al.  Canonical Configurational Interaction Procedure , 1960 .