Variational Monte Carlo on a Budget -- Fine-tuning pre-trained Neural Wavefunctions

TL;DR

This paper introduces a pre-trained deep learning model for quantum chemistry that accurately predicts molecular wavefunctions and energies without full optimization, significantly reducing computational costs and enabling rapid analysis of new molecules.

Contribution

The authors develop a pre-trained DL-VMC model that requires minimal fine-tuning for new molecules, outperforming traditional methods in accuracy and efficiency.

Findings

Zero-shot accuracy improved by two orders of magnitude.

Few fine-tuning steps suffice for accurate relative energies.

Model outperforms CCSD(T)-2Z in absolute energies.

Abstract

Obtaining accurate solutions to the Schr\"odinger 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.