Higher-Order Equivariant Neural Networks for Charge Density Prediction in Materials

TL;DR

The paper introduces ChargE3Net, an E(3)-equivariant grid-based graph neural network that leverages higher-order ($L=4$) tensor representations to predict electron charge densities with high accuracy across diverse materials. By predicting converged densities, the model significantly reduces KS-DFT self-consistent field iterations (median reductions around 27%), while non-self-consistent properties derived from the predicted densities approach chemical accuracy for many materials. The approach demonstrates strong performance on large-scale datasets (Materials Project) and shows linear-time scalability with system size, enabling density predictions for systems orders of magnitude larger than typical ab initio calculations. The results highlight the value of higher-order angular information for materials with covalent bonding and angular variation, and point to future directions toward spin densities, foundation-model DFT, and broader generalization to unseen structures.

Abstract

The calculation of electron density distribution using density functional theory (DFT) in materials and molecules is central to the study of their quantum and macro-scale properties, yet accurate and efficient calculation remains a long-standing challenge. We introduce ChargE3Net, an E(3)-equivariant graph neural network for predicting electron density in atomic systems. ChargE3Net enables the learning of higher-order equivariant feature to achieve high predictive accuracy and model expressivity. We show that ChargE3Net exceeds the performance of prior work on diverse sets of molecules and materials. When trained on the massive dataset of over 100K materials in the Materials Project database, our model is able to capture the complexity and variability in the data, leading to a significant 26.7% reduction in self-consistent iterations when used to initialize DFT calculations on unseen materials. Furthermore, we show that non-self-consistent DFT calculations using our predicted charge densities yield near-DFT performance on electronic and thermodynamic property prediction at a fraction of the computational cost. Further analysis attributes the greater predictive accuracy to improved modeling of systems with high angular variations. These results illuminate a pathway towards a machine learning-accelerated ab initio calculations for materials discovery.

Figures (33)

• Figure 1: Overall architecture of the ChargE3Net model.a Acceleration of DFT property prediction with ChargE3Net charge density initialization. b Illustration of single charge density probe point (center, gray) and its local environment within a periodic atomic system. c Sub-graph of local environment graph with atom nodes $z_i$, $z_j$ and probe node $g_k$. d Neural network architecture. Predicted charge density $\hat{\rho}(\vec{g})$ at probe point $\vec{g}_k$ is computed from the atomic species $z_i$, interatomic displacement vectors $\vec{r}_{ij}$, and probe-atom displacement vectors $\vec{r}_{ik}$ through a series of convolution and gate operations. e$\text{Conv}_\text{atom}$ sends higher-order representations bidirectionally to atom nodes. $\text{Conv}_\text{probe}$ sends higher-order representations from atoms to probes.
• Figure 2: DFT calculations with ChargE3Net. Comparison of DFT calculations initialized from superposition of atomic densities (SAD), ChargE3Net predictions, and ground-truth densities from existing self-consistent calculations (SC). a Average number of SCF steps required to converge the total energy to $5\times 10^{-5}$ eV/atom for non-magnetic materials in the MP test set, with respect to the number of unique species within each material. Error bar shows $\pm$ one standard error. b Number of steps required for self-consistent (SC) and non-self-consistent (NSC) calculations. Box-plots show median as center line, upper and lower quartiles as box limits, 1.5x interquartile range as whiskers, and outliers as points. c Non-self-consistent physical property calculations for total energy, atomic forces, band energies, and band gaps. Dashed lines indicate chemical accuracy (1 meV/atom for per-atom energies, 0.03 eV/Å for forces, and 43 meV for band energies and band gap). Errors of exactly zero are replaced with $10^{-8}$ for visualization purposes. Interior of violin-plots shows median as center line, upper and lower quartiles as box limits, and 1.5x interquartile range as whiskers. d Percentage of materials achieving chemical accuracy for non-self-consistent calculations.
• Figure 3: Effect of angular variance. Comparison of materials with high and low angular variance. a Material Cs(H2PO4) with high angular variance. b Material Rb2Sn6 with low angular variance. Left to right: visualization of charge density isosurfaces (gray); plot of DFT-computed charge density with respect to radial distance from nearest atom; $\epsilon_\text{mae}$ for ChargE3Net model predictions on these materials.
• Figure 4: Effect of rotation order on element-wise error. The lower-left triangles indicate the average $\epsilon_\text{mae}$ per element at $L=0$, and the upper right triangles represent $\epsilon_\text{mae}$ per element at $L=4$. Performance improvements are shown across the periodic table. Values are computed over the set of grid points $G_\text{local}$, and averaged across all materials containing those elements in the test set. Only elements in at least 10 test set materials are shown.
• Figure A.1: t-SNE visualization of Smooth Overlap of Atomic Positions (SOAP) descriptors of the MP train and test subsets, along the GNoME subset used for evaluation. Descriptors are computed using DScribe dscribe with a cutoff radius of 4 Å and compressed darby2022compressing with $\mu=1,\nu=1$.