# A Self-Attention Ansatz for Ab-initio Quantum Chemistry

arXiv: https://arxiv.org/abs/2211.13672 · HTML: https://mlsanity.com/paper/2211.13672/

## Key facts

- Cited by 49 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 59 references that are in the map.
- First posted on arXiv 24 Nov 2022.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [Learning to Plan Chemical Syntheses](https://mlsanity.com/paper/1708.04202/) (2017)
- [Deep neural network solution of the electronic Schr"odinger equation](https://mlsanity.com/paper/1909.08423/) (2019)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping](https://mlsanity.com/paper/2110.01765/) (2021)
- [Ab-initio quantum chemistry with neural-network wavefunctions](https://mlsanity.com/paper/2208.12590/) (2022)
- [Better, Faster Fermionic Neural Networks](https://mlsanity.com/paper/2011.07125/) (2020)
- [Gold-standard solutions to the Schr"odinger equation using deep learning: How much physics do we need?](https://mlsanity.com/paper/2205.09438/) (2022)
- [Discovering Quantum Phase Transitions with Fermionic Neural Networks](https://mlsanity.com/paper/2202.05183/) (2022)
- [Sampling-free Inference for Ab-Initio Potential Energy Surface Networks](https://mlsanity.com/paper/2205.14962/) (2022)
- [Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation](https://mlsanity.com/paper/2112.03491/) (2021)

## Built on by (most cited first, 30 of 49)

- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [PixArt-Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation](https://mlsanity.com/paper/2403.04692/) (2024)
- [Generalizing Neural Wave Functions](https://mlsanity.com/paper/2302.04168/) (2023)
- [A Kaczmarz-inspired approach to accelerate the optimization of neural network wavefunctions](https://mlsanity.com/paper/2401.10190/) (2024)
- [Accurate Computation of Quantum Excited States with Neural Networks](https://mlsanity.com/paper/2308.16848/) (2023)
- [An ab initio foundation model of wavefunctions that accurately describes chemical bond breaking](https://mlsanity.com/paper/2506.19960/) (2025)
- [Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schr"odinger Equation](https://mlsanity.com/paper/2307.07050/) (2023)
- [Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Cost](https://mlsanity.com/paper/2412.08541/) (2024)
- [Variational Monte Carlo on a Budget -- Fine-tuning pre-trained Neural Wavefunctions](https://mlsanity.com/paper/2307.09337/) (2023)
- [Neural Pfaffians: Solving Many Many-Electron Schr"odinger Equations](https://mlsanity.com/paper/2405.14762/) (2024)
- [Accurate Ab-initio Neural-network Solutions to Large-Scale Electronic Structure Problems](https://mlsanity.com/paper/2504.06087/) (2025)
- [Towards a Foundation Model for Neural Network Wavefunctions](https://mlsanity.com/paper/2303.09949/) (2023)
- [Ab-initio variational wave functions for the time-dependent many-electron Schr"odinger equation](https://mlsanity.com/paper/2403.07447/) (2024)
- [Transferable Neural Wavefunctions for Solids](https://mlsanity.com/paper/2405.07599/) (2024)
- [Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo](https://mlsanity.com/paper/2307.08214/) (2023)
- [Are queries and keys always relevant? A case study on Transformer wave functions](https://mlsanity.com/paper/2405.18874/) (2024)
- [Highly Accurate Real-space Electron Densities with Neural Networks](https://mlsanity.com/paper/2409.01306/) (2024)
- [Large Electron Model: A Universal Ground State Predictor](https://mlsanity.com/paper/2603.02346/) (2026)
- [Functional Neural Wavefunction Optimization](https://mlsanity.com/paper/2507.10835/) (2025)
- [First-Principles AI finds crystallization of fractional quantum Hall liquids](https://mlsanity.com/paper/2602.03927/) (2026)
- [QERNEL: a Scalable Large Electron Model](https://mlsanity.com/paper/2604.26018/) (2026)
- [Convergence of variational Monte Carlo simulation and scale-invariant pre-training](https://mlsanity.com/paper/2303.11602/) (2023)
- [On Representing Electronic Wave Functions with Sign Equivariant Neural Networks](https://mlsanity.com/paper/2403.05249/) (2024)
- [nabla2DFT: A Universal Quantum Chemistry Dataset of Drug-Like Molecules and a Benchmark for Neural Network Potentials](https://mlsanity.com/paper/2406.14347/) (2024)
- [Is attention all you need to solve the correlated electron problem?](https://mlsanity.com/paper/2502.05383/) (2025)
- [Ab-initio simulation of excited-state potential energy surfaces with transferable deep quantum Monte Carlo](https://mlsanity.com/paper/2503.19847/) (2025)
- [Spin-Adapted Neural Network Wavefunctions in Real Space](https://mlsanity.com/paper/2511.01671/) (2025)
- [Sorting Out Quantum Monte Carlo](https://mlsanity.com/paper/2311.05598/) (2023)
- [Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schr"odinger Equation](https://mlsanity.com/paper/2502.05318/) (2025)
- [Understanding the Capabilities of Molecular Graph Neural Networks in Materials Science Through Multimodal Learning and Physical Context Encoding](https://mlsanity.com/paper/2505.12137/) (2025)

## Abstract (arXiv)

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.
