Transformer models predict electron densities accurately and efficiently
Transformer-inspired architectures make machine-learned electron densities accurate and transferable enough to bypass self-consistency
Machine Learning
Summary
Predicting how electrons arrange themselves in molecules is important for understanding chemistry but usually requires slow, repeated calculations. This paper presents machine learning models inspired by Transformer architectures that can quickly and accurately predict electron densities without needing these repeated steps. The authors show their models work well on standard molecule datasets and can handle molecules with new elements or larger sizes. This makes simulations of molecules faster while keeping high accuracy.
What this means in practice
- •For computational chemists: Use the models to accelerate electronic structure calculations by bypassing iterative self-consistent steps while maintaining accuracy in predicted energies and forces.
- •For materials simulation teams: Drive molecular dynamics simulations faster by employing predicted forces directly from machine-learned electron densities without costly quantum iterations.
Authors
Jeheon Woo
Abstract
We introduce machine learning models that accurately predict the electron density, the central variable of density functional theory. They are built on Transformer-inspired architectures, which require little computation per parameter and can therefore be enlarged at low cost. The proposed models, SwinUNETR-V2 and its attention-free counterpart NeXtUNETR, achieve state-of-the-art accuracy on QM9, with normalized mean absolute errors in the density averaging 0.075% and 0.046%, respectively. The predicted densities are accurate enough to bypass the self-consistent iteration: a single Kohn-Sham step reproduces the self-consistent electronic structure. For NeXtUNETR, the median errors are 0.036 meV in the total energy, 0.65 meV/angstrom in the forces and 1.8 meV in the HOMO-LUMO gap. The forces not only recover the self-consistent minima in geometry optimization but also successfully drive ab initio molecular dynamics. Fine-tuned on small datasets, both networks also transfer accurately to systems containing elements absent from QM9 or larger than its molecules.