Equivariant Flow Matching for Electron Density Prediction

2026-10-01

AI Summary

This study addresses the tension between the high computational cost of grid-based methods and the neglect of structural correlations in basis-set approaches for density functional theory (DFT) electron density prediction. To this end, we propose OrbFlow, a model built upon an SE(3)-equivariant flow matching framework that directly predicts Gaussian basis set coefficients. Furthermore, a two-stage trajectory curriculum learning strategy is designed to mitigate discretization drift arising from numerical integration. Experimental results on the QM9 dataset demonstrate that OrbFlow reduces prediction error by 13.6% and decreases self-consistent field (SCF) iteration counts by 68%. The model also exhibits zero-shot transferability and accurately recovers high-order multipole moments, thereby achieving an effective balance between computational efficiency and predictive accuracy.

Abstract

Machine learning surrogates for density functional theory (DFT) have been increasingly used to reduce the cost of first-principles calculations. In this arena, predicting real-space electron densities offers a scalable and transferable initialization for self-consistent field (SCF) procedures. However, current methods face a clear dilemma. That is, grid-based architectures incur a high computational cost, while basis-set methods fail to capture the structural correlations inherent in the coefficient space. Here, we develop OrbFlow, an $\mathrm{SE}(3)$-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching. OrbFlow retains the efficiency of a compact atom-centered basis while replacing pointwise regression with a learned probability path over the full coefficient space. It is trained through a two-phase trajectory curriculum that mitigates discretization drift during numerical integration. OrbFlow achieves state-of-the-art accuracy on QM9, reducing density error by 13.6% relative to the previous best model, and reduces error by 51% to 63% on every molecule of the MD benchmark relative to the strongest prior method sharing its basis. The predicted density also cuts SCF iterations by up to 68% with zero-shot transfer to unseen exchange-correlation functionals and recovers dipole and quadrupole moments to within a few percent of DFT references without any SCF calculation.

Problem

• Electron Density Prediction
• Density Functional Theory
• Basis-set Methods
• Grid-based Architectures
• Machine Learning Surrogates

Innovation

• Equivariant Flow Matching
• Electron Density Prediction
• SE(3)-equivariant Generative Model
• Gaussian-type Orbital Coefficients
• Two-phase Trajectory Curriculum