A Self-Attention Ansatz for Ab-initio Quantum Chemistry

TL;DR

This paper introduces the Wavefunction Transformer, a self-attention neural network architecture that improves the accuracy of solving the many-electron Schrödinger equation in quantum chemistry, especially for larger molecules.

Contribution

The Wavefunction Transformer (Psiformer) is a novel self-attention based neural network that enhances first-principles quantum calculations without external data, outperforming previous models like FermiNet and PauliNet.

Findings

Significantly improved ground state energy calculations for large molecules

Demonstrated the effectiveness of self-attention in modeling electron correlations

Achieved qualitative leaps in accuracy over previous neural network approaches

Abstract

We present a novel neural network architecture using self-attention, the Wavefunction Transformer (Psiformer), which can be used as an approximation (or Ansatz) for solving the many-electron Schr\"odinger equation, the fundamental equation for quantum chemistry and material science. This equation can be solved from first principles, requiring no external training data. In recent years, deep neural networks like the FermiNet and PauliNet have been used to significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here we show that the Psiformer can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy in chemical calculations on larger systems.