Equivariant Flow Matching for Electron Density Prediction
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 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.
Related
- Multi-electron chemistry for high-energy-density aqueous organic redox flow batteries.
- Steerable Single-step Retrosynthesis with Electron-pushing Discrete Flow Matching and Feynman-Kac Steering
- FLOWR.ROOT – A flow matching-based foundation model for joint multi-purpose structure-aware 3D ligand generation and affinity prediction
- Electron-Density-Informed Molecular Property Prediction Using Surrogate Functionals
- FlowBatt: Flow Matching for Probabilistic Battery Degradation Prediction
BibTeX
@misc{liang2026equivariant,
title = {{Equivariant Flow Matching for Electron Density Prediction}},
author = {Chenxing Liang and Chengdong Wang and Yuchao Lin and Xiaofeng Qian and Shuiwang Ji},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.02651},
eprint = {2610.02651},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.02651},
}
RIS
TY - UNPB
TI - Equivariant Flow Matching for Electron Density Prediction
AU - Chenxing Liang
AU - Chengdong Wang
AU - Yuchao Lin
AU - Xiaofeng Qian
AU - Shuiwang Ji
PY - 2026
DA - 2026/10/02
PB - arXiv
DO - 10.48550/arxiv.2610.02651
UR - https://arxiv.org/abs/2610.02651
AB - 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 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.
ER -