[Paper Review] DeepErwin

Michael Scherbela, Leon Gerard+1 more|arXiv (Cornell University)|Mar 17, 2023

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper introduces DeepErwin, a foundation wavefunction model that maps Hartree-Fock orbitals to high-accuracy neural network orbitals via a graph neural network, enabling transfer learning across molecules. Pre-training on diverse compounds allows minimal fine-tuning to achieve CCSD(T)-like accuracy 20x faster than standard training.

ABSTRACT

Deep neural networks have become a highly accurate and powerful wavefunction ansatz in combination with variational Monte Carlo methods for solving the electronic Schrödinger 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.

Motivation & Objective

• To overcome the high computational cost of training deep neural network wavefunctions from scratch for each new molecular system.
• To enable generalization of wavefunction models across different molecules and geometries through a parameter-efficient, size-invariant architecture.
• To establish a foundation wavefunction model pre-trained on diverse compounds, reducing the need for full re-optimization.
• To demonstrate that extensive pre-training leads to systematic improvements in accuracy and transferability.
• To outperform conventional high-accuracy methods like CCSD(T)-ccpVTZ with significantly reduced computational cost through fine-tuning.

Proposed method

• A graph neural network (GNN) processes nuclear coordinates and initial orbital descriptors to generate learned, correlated orbitals independent of system size.
• The model uses a message-passing architecture with symmetric and anti-symmetric components to ensure proper fermionic statistics.
• Electron orbitals are constructed via a learnable function combining electron embeddings and orbital descriptors, ensuring spin and spatial symmetry.
• The wavefunction is built as a Slater determinant of these learned orbitals, with symmetry enforced through anti-symmetric activation functions.
• A pre-training phase on small molecular fragments enables transfer to larger, chemically similar compounds via fine-tuning.
• The method decouples expensive GNN and orbital network evaluations per batch, reducing computational overhead.

Experimental results

Research questions

• RQ1 Can a single neural network wavefunction model be pre-trained across multiple molecules and geometries to enable fast, accurate predictions on new systems?
• RQ2 Does pre-training on small molecular fragments lead to improved accuracy and generalization on larger, chemically similar compounds?
• RQ3 Can the model achieve CCSD(T)-level accuracy with only minimal fine-tuning, reducing computational cost by orders of magnitude?
• RQ4 How does model size, data size, and pre-training steps affect the final accuracy and transferability of the wavefunction model?
• RQ5 How does the proposed method compare to concurrent approaches like GLOBE in terms of energy accuracy and generalization across diverse compounds?

Key findings

• Pre-training on small molecular fragments enables successful transfer to larger, chemically similar compounds with high accuracy in relative energy predictions.
• Fine-tuning the pre-trained DeepErwin model achieves CCSD(T)-ccpVTZ-level accuracy approximately 20 times faster than training a new model from scratch.
• Systematic scaling of model size, data size, and pre-training steps leads to consistent and substantial improvements in wavefunction accuracy.
• The model achieves lower absolute energies and higher accuracy in relative energy calculations than the concurrent method GLOBE, especially on out-of-distribution molecules.
• The architecture’s size-invariance allows weight sharing across molecules of different sizes, overcoming a key limitation of prior wavefunction ansatzes.
• The method maintains high accuracy across in-distribution and out-of-distribution molecules, demonstrating strong generalization capability.