Variational Monte Carlo on a Budget -- Fine-tuning pre-trained Neural Wavefunctions
This work addresses the computational bottleneck for researchers in quantum chemistry by enabling faster and more accurate molecular energy predictions with minimal fine-tuning.
The paper tackles the high computational cost of deep-learning-based Variational Monte Carlo (DL-VMC) in quantum chemistry by proposing a pre-trained model that achieves zero-shot absolute energies outperforming CCSD(T)-2Z and reduces fine-tuning steps for relative energies, improving zero-shot accuracy by two orders of magnitude.
Obtaining accurate solutions to the Schrödinger equation is the key challenge in computational quantum chemistry. Deep-learning-based Variational Monte Carlo (DL-VMC) has recently outperformed conventional approaches in terms of accuracy, but only at large computational cost. Whereas in many domains models are trained once and subsequently applied for inference, accurate DL-VMC so far requires a full optimization for every new problem instance, consuming thousands of GPUhs even for small molecules. We instead propose a DL-VMC model which has been pre-trained using self-supervised wavefunction optimization on a large and chemically diverse set of molecules. Applying this model to new molecules without any optimization, yields wavefunctions and absolute energies that outperform established methods such as CCSD(T)-2Z. To obtain accurate relative energies, only few fine-tuning steps of this base model are required. We accomplish this with a fully end-to-end machine-learned model, consisting of an improved geometry embedding architecture and an existing SE(3)-equivariant model to represent molecular orbitals. Combining this architecture with continuous sampling of geometries, we improve zero-shot accuracy by two orders of magnitude compared to the state of the art. We extensively evaluate the accuracy, scalability and limitations of our base model on a wide variety of test systems.