[Paper Review] A Self-Attention Ansatz for Ab-initio Quantum Chemistry

Ingrid von Glehn, James S. Spencer+1 more|arXiv (Cornell University)|Nov 24, 2022

Machine Learning in Materials Science30 citations

TL;DR

The Psiformer is a self-attention based wavefunction Ansatz for ab-initio quantum chemistry that improves ground-state energies, especially for larger molecules, compared to FermiNet and FermiNet+SchNet.

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ödinger 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.

Motivation & Objective

• Develop a self-attention based neural network Ansatz for solving the many-electron Schrödinger equation without external data.
• Improve accuracy of ground-state energy calculations for molecular systems compared to existing neural-network Ansatzes.
• Demonstrate scalability of attention-based architectures to larger molecules.
• Incorporate physical constraints such as cusp conditions via Jastrow factors and envelope functions.

Proposed method

• Define Psi_theta as exp(J_theta) times a determinant sum, with J_theta a Jastrow-like pre-factor and Phi_theta^k learned via self-attention networks.
• Use a single-stream multi-headed self-attention architecture to model electron-nuclear interactions input features with spin concatenated into inputs.
• Employ input rescaling to handle large-separation geometries and enforce boundary conditions via an exponential envelope; apply residual connections and a tanh nonlinearity.
• Train with variational quantum Monte Carlo using Rayleigh quotient optimization and Metropolis-Hastings sampling; optimize with Kronecker-factored Approximate Curvature (KFAC).
• Compare Psiformer to FermiNet and FermiNet+SchNet on small and large molecular benchmarks, including benzene, to assess accuracy gains.

Experimental results

Research questions

• RQ1 Can a self-attention based neural network provide a competitive or superior Ansatz for the many-electron Schrödinger equation in ab-initio quantum chemistry?
• RQ2 Does the Psiformer improve ground-state energy accuracy relative to existing neural-network Ansätze, particularly for larger molecular systems?
• RQ3 How does incorporating spin, cusp handling via Jastrow, and attention-based electron-nuclear coupling affect learning and stability across system sizes?
• RQ4 To what extent can self-attention architectures scale to molecules with dozens to dozens of electrons compared to prior methods?

Key findings

• Psiformer achieves substantially improved energies over FermiNet and FermiNet+SchNet, especially on larger systems.
• On small molecules, the large Psiformer reaches within about 1 kcal/mol of CCSD(T)/CBS for some systems like ozone.
• For 42–74 electron molecules (benzene to CCl4), Psiformer outperforms FermiNet and even surpasses some published DMC energies, with notable gains on benzene and CCl4.
• Psiformer energies improve with system size, showing up to ~75 kcal/mol (100 mHa) improvements over conventional baselines on large molecules.
• LayerNorm and other training choices influence optimization, with Psiformer generally maintaining superior accuracy across configurations.