A Recipe for Charge Density Prediction

TL;DR

This paper tackles accelerating charge-density prediction to speed up DFT workflows by bypassing expensive Kohn–Sham SCF iterations. It introduces a three-ingredient recipe: (i) an atomic plus virtual orbital charge-density representation using an even-tempered Gaussian basis with trainable exponents, (ii) a high-capacity SE(3)-equivariant backbone (eSCN) that predicts basis coefficients and per-basis scaling, and (iii) a carefully staged training procedure including fine-tuning of exponents. On the QM9 charge-density benchmark, the method achieves state-of-the-art accuracy (best NMAE around $0.178$) and substantial efficiency gains (up to $125.29$ mol/min and up to $171\times$ speedups relative to prior SOTA), illustrating a favorable accuracy–throughput Pareto. The results demonstrate strong potential to accelerate DFT-based materials and molecular discovery by providing accurate, scalable charge-density predictions and useful descriptors for downstream properties.

Abstract

In density functional theory, charge density is the core attribute of atomic systems from which all chemical properties can be derived. Machine learning methods are promising in significantly accelerating charge density prediction, yet existing approaches either lack accuracy or scalability. We propose a recipe that can achieve both. In particular, we identify three key ingredients: (1) representing the charge density with atomic and virtual orbitals (spherical fields centered at atom/virtual coordinates); (2) using expressive and learnable orbital basis sets (basis function for the spherical fields); and (3) using high-capacity equivariant neural network architecture. Our method achieves state-of-the-art accuracy while being more than an order of magnitude faster than existing methods. Furthermore, our method enables flexible efficiency-accuracy trade-offs by adjusting the model/basis sizes.

Figures (5)

• Figure 1: (a) Illustration of the orbital-based method for charge density representation for an example molecule (indole, C8H7N). The overall charge density is represented as a sum over spherical-harmonics-based atomic orbital basis functions (spherical fields) centered at each atom. (b) Left: Illustration of the probe-based method for charge density representation. The charge density is represented as a voxel where each grid point (probe node) represents a scalar density at that coordinate. The voxel for the example molecule is of size $108 \times 96 \times 40$. Grid points with very small charge densities ($<0.05$) are not visualized. Right: For a probe-based machine learning prediction model, the voxel contains too many grid points to be processed simultaneously. Sampling of the voxel points is needed during training and inference. All charge densities use the same colormap scale at the right-most side of the figure. Atom color code: H (white), C (gray), N (blue). The charge density is from the QM9 charge density dataset jorgensen2022equivariant.
• Figure 2: (a) Two example molecules (left: indole, C8H7N; right: methanol CH3OH), before and after the bond-midpoint-based virtual coordinates (small black points) are inserted. Atom color code: H (white), C (gray), N (blue), O (red), virtual nodes (small, black). (b) The number of Gaussian-type orbital basis functions for selected elements in the def2-QZVPPD basis set and even-tempered Gaussian basis sets derived from it under different $\beta$, which controls the number of basis functions as described in \ref{['eqn:even_tempered']}.
• Figure 3: Efficiency--accuracy trade-off for SCDP models. The models with scaling factor fine-tuning form the Pareto front.
• Figure 4: Visualization of the reference charge density and prediction errors for select SCDP models with two representative test molecules (top: C2H3NO2 and bottom: C8H18O). The first column is the ground truth charge density with the corresponding color scale. The next five columns are prediction errors from various models which all use the same color scale in the rightmost for error magnitude. The prediction errors significantly reduce with larger model size, virtual orbitals, orbital exponent scaling, and a larger basis set. VO stands for virtual orbitals. Scaling stands for scaling factor fine-tuning. The virtual orbitals significantly reduce errors around chemical bonds. Atom color code: H (white), C (gray), N (blue), O (red), virtual nodes (small, black).
• Figure 5: Convergence of validation NMAE during pretraining and finetuning.