Learning Equivariant Non-Local Electron Density Functionals

TL;DR

The paper tackles the challenge of learning accurate non-local exchange-correlation functionals for KS-DFT by proposing the Equivariant Graph Exchange Correlation (EG-XC). EG-XC combines nuclei-centered, $SO(3)$-equivariant embeddings with an equivariant GNN to capture molecular-range density interactions, and a non-local meta-GGA correction whose weights are learned via a differentiable SCF training pipeline using only energy targets. Empirically, EG-XC yields substantial improvements over semi-local ML functionals across MD17, 3BPA, and QM9, including strong extrapolation and data-efficiency properties, while maintaining favorable runtime scaling compared to hybrid functionals. These results suggest a promising path toward accurate, scalable, data-efficient DFT functionals that integrate physics-inspired biases with flexible non-local modeling, with potential extensions to periodic systems and orbital-free regimes.

Abstract

The accuracy of density functional theory hinges on the approximation of non-local contributions to the exchange-correlation (XC) functional. To date, machine-learned and human-designed approximations suffer from insufficient accuracy, limited scalability, or dependence on costly reference data. To address these issues, we introduce Equivariant Graph Exchange Correlation (EG-XC), a novel non-local XC functional based on equivariant graph neural networks (GNNs). Where previous works relied on semi-local functionals or fixed-size descriptors of the density, we compress the electron density into an SO(3)-equivariant nuclei-centered point cloud for efficient non-local atomic-range interactions. By applying an equivariant GNN on this point cloud, we capture molecular-range interactions in a scalable and accurate manner. To train EG-XC, we differentiate through a self-consistent field solver requiring only energy targets. In our empirical evaluation, we find EG-XC to accurately reconstruct `gold-standard' CCSD(T) energies on MD17. On out-of-distribution conformations of 3BPA, EG-XC reduces the relative MAE by 35% to 50%. Remarkably, EG-XC excels in data efficiency and molecular size extrapolation on QM9, matching force fields trained on 5 times more and larger molecules. On identical training sets, EG-XC yields on average 51% lower MAEs.

Figures (7)

• Figure 1: Illustration of Equivariant Graph Exchange Correlation (EG-XC)'s four components: (1) We obtain a finite atomic-range point cloud representation ${\bm{\mathsfit{H}}}^{(0)}$ by convolving the electron density $\rho$ with radial filters $\Gamma_k:\mathbb{R}_+\to\mathbb{R}$ and spherical harmonics $Y^l_m:\mathbb{R}^3\to\mathbb{R}$ at the nuclear position. (2) The embeddings ${\bm{\mathsfit{H}}}^{(0)}$ are updated using equivariant message passing to obtain molecular-range effects in ${\bm{\mathsfit{H}}}^{(T)}$. (3) We define a non-local feature density ${{\bm{g}}_\text{NL}}:\mathbb{R}^3\to\mathbb{R}^d$ from which we derive the exchange correlation density ${\epsilon_\text{XC}}:\mathbb{R}^3\to\mathbb{R}$. (4) We add a graph readout of ${\bm{\mathsfit{H}}}^{(T)}$ to learn additional corrections. To obtain $E_{\text{XC}}[\rho]$, we integrate ${\epsilon_\text{XC}}$ and add it to the global graph readout.
• Figure 2: Two-dimensional slice of the potential energy surface of 3BPA with the dihedral $\beta=120°$. Pure force fields like NequIP struggle to recover the shape of this out-of-distribution energy surface. When paired with DFT calculations, one can see that the energy surface moves closer to the target shape but introduces additional extrema. Learnable XC functionals like dickHighlyAccurateConstrained2021 and EG-XC demonstrate significantly better reproduction of the target energy surface.
• Figure 3: MAE in m on QM9 size extrapolation. Each row represents a different method. The four groups indicate the maximum number of heavy atoms in the training set and its size. Each column represents the test subset with the number of heavy atoms listed above.
• Figure 4: Energy surfaces of the 3BPA dataset at $\beta=120°$.
• Figure 5: Energy surfaces of the 3BPA dataset at $\beta=150°$.