1 Introduction↩︎

In the past decade, machine learning (ML) methods have gained popularity in materials modeling due to their favorable \(O(N)\) scaling in supervised tasks[1], thanks to the so-called nearsightedness principle of electronic matter[2]. A major source of training data for ML methods comes from density functional theory(DFT) calculations. DFT is widely used in quantum atomistic simulations owing to its good balance between accuracy and efficiency[3]. However, the \(O(N^3)\) scaling for solving the Kohn-Sham equations[4] self-consistently hindered its application to large scale simulations. To address this, various ML approaches have been developed to learn the mapping from atomic structures to DFT-computed properties directly. Among these, neural network potentials (NNPs), trained to predict formation energies and atomic forces, are the most extensively studied and form the foundation of other ML methods in materials modeling[5], [6]. The success of NNPs has inspired researchers to adapt their architectures to learn other properties, such as the ground state charge density[7].

Accurate ML charge density is desired for several reasons. Firstly, the first Hohenberg-Kohn theorem states that the ground-state charge density uniquely determines all ground-state properties of an interacting electron system[8], [9]. In principle ML charge density could allow us to bypass the self-consistent procedure in Kohn–Sham formalism to obtain ground-state properties by one diagonalization process. However, DFT calculations are extremely sensitive to the self-consistent charge density. In practice, the more accurate the ML-predicted charge density, the more likely it is to achieve self-consistency, and the more reliable the properties are derived from non-self-consistent calculations. Secondly, the ground-state charge density can be used further as descriptors for downstream machine learning tasks. In NNPs this indirect approach has been reported to outperform significantly[7], [10], [11]. Lastly, encoding rich electronic structure and chemical bonding information, charge density–along with its derived quantities such as Bader charge[12], electron localization function[13] and Hartree potential–can be used directly to characterize a wide range of phenomena including polarization[14], defect states[15], charge transfer in photovoltaics[16] and lithium batteries[17], to name a few.

Chronologically, the architecture of ML charge density models has evolved from kernel regression to message passing neural networks, and from obeying invariance to incorporating equivariant hidden layers, for better transferability. Meanwhile, the choice of fitting targets ranges from basis expansion coefficients to grid representations, depending on the trade-off between efficiency and accuracy. Early triumphs in modeling charge density use kernel ridge regression to fit basis expansion coefficients of charge density[7], [18]–[20]. By comparing unknown structures to those in the training set through a kernel function, kernel-based ML charge density models have achieved transferability on molecular dynamics trajectories and diverse molecular datasets. However, kernel methods scale cubically with the size of dataset, making it unfeasible to tackle with the full materials space. In recent years message-passing neural networks–which update iteratively node and edge features of a graph by aggregate neighbor messages–have gained prominence in materials science due to their flexibility and transferability on large dataset, albeit at the cost of longer training time compared to kernel methods[21], [22]. MPNN charge density models can represent charge density with either basis expansion coefficients[23] or grid quantities[24]–[26]. Grid-based methods place fictitious probe nodes at grid points, which receive messages only from nearby atoms. Although much slower than basis-based methods, grid-based methods have higher accuracy and are not affected by the choice of basis. Finally, although charge density is itself an invariant quantity which transforms like scalar field, researchers have shown that it is necessary to incorporate high-order equivariant hidden layers to accurately address the complex chemical environments in molecules and solids[27].

Recent developments such as \(\Delta\)-learning and physics-informed neural networks suggest that incorporating physical priors can improve data efficiency and model accuracy[28]–[30]. However, to the best of our knowledge, all ML charge density models reported so far learn charge density from scratch, using the total charge density as training target. Nevertheless, there are groups using atomic charge densities (SAED) as input of charge density models, demonstrating the potential of SAED to provide physical information for data driven tasks. In [31] the authors emulated the Kohn-Sham self-consistent iterations with SAED as the initial guess. They used a basis-based machine learning charge density model to predict the Kohn-Sham density map. Another group used SAED as input charge density to recover DFT charge density using image super-resolution techniques[32]. However, since image super-resolution methods operate convolutions over regular uniform grids instead of atomic graphs, they do not inherently preserve rotational invariance of charge density.

In this work, we propose using the difference charge density as the training target, referred to as \(\Delta\)-SAED, which robustly improves accuracy at no additional cost. Using a state-of-the-art grid-based MPNN model, we demonstrated that \(\Delta\)-SAED reduces the MAE of charge density for more than 99% of structures in the QM9 and NMC datasets, and for 90% in the MP dataset. On a Si allotropy dataset, we further showed that \(\Delta\)-SAED almost always enhances the transferability to chemical accuracy for physical properties derived from non-self consistent DFT calculations, even when the accuracy of charge density itself is not improved. Therefore, \(\Delta\)-SAED is more reliable when applied to unseen structures. We examined the effects of \(\Delta\)-SAED on the radial distributions of the training targets and attributed its advantage to the physical priors from SAED, which reduce the complexity of the radial and angular dependencies perceived by the MPNN, alleviating its limited expressive power.

2 Results↩︎

2.1 Work Flow↩︎

In DFT, we are concerned with the ground-state electronic charge density \(\rho_t(\mathbf{r})\). For a single isolated atom \(i\), the charge density of its electrons \(\rho_a^{i}(\mathbf{r} - \mathbf{r}_i)\) is localized around the nucleus \(\mathbf{r}_i\). For a system of electrons and nuclei, we can define an auxiliary atomic charge density \(\rho_a(\mathbf{r}) = \sum_i\rho_a^i(\mathbf{r} - \mathbf{r}_i)\) which is the superposition of isolated atomic electron densities (SAED) centered at corresponding nuclei. The difference charge density (DCD) is introduced naturally as the difference between total charge density (TCD) and SAED: \[\label{eq:difference32charge32density} \rho_d(\mathbf{r}) = \rho_t(\mathbf{r}) - \rho_a(\mathbf{r})\tag{1}\]

The \(\Delta\)-SAED workflow is illustrated in Fig. 1. When preparing charge density datasets using established DFT codes, we collect not only atomic structures (model inputs) and \(\rho_t\) (ultimate goal), but also \(\rho_a\), which can be easily extracted. During training the model parameters are optimized by minimizing \(||\hat{\rho}_d - \rho_d||\), where \(\hat{\rho}_d\) is the predicted DCD. In the inference stage, the predicted TCD \(\hat{\rho}_t\) is obtained by adding \(\rho_a\) and the model-predicted \(\hat{\rho}_d\).

Aiming to learn DCD defined in Eq. 1 , \(\Delta\)-SAED is not a trivial extension of the common trick in NNPs of training on atomization energy instead of total energy. Rather, it embodies an application of \(\Delta\)-learning[28], [30], which follows the general principle that, given a physical Ansatz, the neural network only needs to learn the residual part \(\Delta\), potentially leading to improved accuracy. NNPs primarily focus on predicting forces. Since the superposition of atomic energies is independent of atomic coordinates, it does not contribute to the forces. Hence we do not consider the superposition of atomic energies as a form of physical prior in NNPs. In contrast, \(\rho_t\) has clear physical meanings, allowing us to calculate exchange-correlation potential and other related quantities. By training \(\rho_d\) as target, we effectively introduce physical prior knowledge or bias into the neural network. The analogy of SAED in NNPs is \(\Delta\)-DFT[30] which uses DFT energies as physical prior of CCSD(T) energies to achieve chemical accuracies.

2.2 Performance on Benchmark Datasets↩︎

The performance of \(\Delta\)-SAED is evaluated on three datasets used for benchmarking charge density models in prior studies: organic molecules (QM9)[25], [33], nickel manganese cobalt battery cathode materials (NMC)[26], and all structures with available charge density data from the Materials Project (MP) inorganic crystal database[34]. We choose Charge3Net[27], an E(3)-equivariant charge density model incorporating high-order messages, to learn the total and difference charge density. Charge3Net is grid-based, offering a more expressive representation of density compared to basis-based methods.

Following [7], [26], [27], \(\varepsilon_{mae}\) is used consistently as the evaluation metric for both TCD and DCD models, and is defined as: \[\label{eq:density32mae} \varepsilon_{mae} = \frac{\int_{\Omega}d^3\mathbf{r}|\rho(\mathbf{r}) - \hat{\rho}(\mathbf{r})|}{\int_{\Omega}d^3\mathbf{r}\rho_t(\mathbf{r})} \times 100\%,\tag{2}\] where \(\rho\) and \(\hat{\rho}\) denote the target and predicted total (or difference) charge densities, respectively, \(\rho_t\) denotes the target total charge density used for normalization. \(\varepsilon_{mae}\) is reported in percentage. Since the goal of ML-predicted charge density is to accelerate DFT calculations and enable further charge analyses—which require charge density information on all grid points—throughout our study we perform integrations over the full set of grid points used in the DFT calculations unless otherwise noted. For a given structure, \(\Delta\)-SAED is assessed by the relative reduction of \(\varepsilon_{mae}\): \(\Delta_{rel}\varepsilon_{mae} = \frac{\varepsilon_{mae}^{TCD} - \varepsilon_{mae}^{DCD}}{\varepsilon_{mae}^{TCD}}\), where \(\varepsilon_{mae}^{TCD}\) (\(\varepsilon_{mae}^{DCD}\)) is the density error of total (difference) charge density.

The result of DCD and TCD models in Table 1 together with the \(\Delta_{rel}\varepsilon_{mae}\) distribution via violin plot in Fig 2 (a) demonstrate the effectiveness of \(\Delta\)-SAED across three benchmark datasets. The validity of the results are ensured by careful contrast settings which will be elaborated in Sec 4.2. Predicting \(\rho_d\) yields lower average \(\varepsilon_{mae}\) consistently, with an average of 17.0% \(\Delta_{rel}\varepsilon_{mae}\) for the NMC dataset, 33.6% for the QM9 dataset, and 14.4% for the MP dataset. Moreover, the percentage of test structures showing improved prediction accuracy exceeds 99% for both the QM9 and NMC datasets, and remains as high as 89.1% for the more diverse MP dataset, further demonstrating the capability of \(\Delta\)-SAED as a simple yet robust method to enhance accuracy over baseline approaches. Notably, there are a few outliers for the NMC and the MP dataset, which show much less accuracy of \(\varepsilon_{mae}\) after imposing \(\Delta\)-SAED technique. We hypothesize that this is due to the rarity of the corresponding chemical environments in the training data, which may be mitigated by employing training strategies such as active learning[35].

2.3 Improved Transferability for Properties from Non-Self Consistent Calculations↩︎

Using a Si allotropy dataset composed of diverse crystal structures from Materials Project[36], we report the improved transferability of \(\Delta\)-SAED by comparing the MAEs of per-atom energies, forces, band gaps and band energies derived from non-self-consistent calculations with predicted charge densities. We first collected 42 Si allotropes from Materials Project, 24 of which were used to construct the Si dataset, while the remaining 18 allotropes were reserved for transferability evaluation. The 24 Si allotropes for the dataset, among which one-fourth are insulators, contain no more than 60 atoms and exhibit good stability. Twelve of these were further expanded into larger supercells. We then introduced random displacements of 0.1Å to the original structures 30 times, generating a total of 1080 structures. The Si dataset was then split into training, validation, and test sets with a ratio of 8:1:1.

Using identical hyperparameters, the DCD (TCD) models yield MAEs of 0.02 meV/atom (0.2 meV/atom), 0.5 meV/Å (1.5 meV/Å), 0.4 meV (0.9 meV), 0.06 meV (0.1 meV), for per-atom energy, force, band energy and band gap respectively, while the commonly adopted chemical accuracy thresholds are 1 meV/atom, 30 meV/Å, 42 meV and 43 meV. Not surprisingly, both the DCD and TCD models achieved chemical accuracy predicting random displaced structures in the training set, with the DCD model yielding slightly lower MAEs, although the average density error \(\varepsilon_{mae}\) is similar (0.3%) for both models. This suggests the comparable capability and even potential advantages of \(\Delta\)-SAED in related applications such as phonon calculations, where for ML models good transferability is required for similar chemical environment with larger system size.

We further evaluate the models’ transferability with 18 Si allotropes in MP outside of the training set. The average \(\varepsilon_{mae}\) of DCD model (1.42%) is slightly smaller than that of TCD model (1.62%), which is consistent with the trend in Table 1. The MAEs of non-self-consistent properties are presented in Fig 3. For mechanical properties, the average per-atom energy MAEs are 6.4 meV/atom (DCD) and 120 meV/atom (TCD). The average force MAEs are 46 meV/Å (DCD) and 260 meV/Å (TCD). The force MAEs of the DCD model is close to the common threshold for structure relaxations (30 meV/Å). For electronic properties, the average band energy MAES are 48 meV (DCD) and 92 meV (TCD), while the band gap MAEs are 29 meV (DCD) and 87 meV (TCD). The band gap MAEs for some structures are exactly zero, indicating that their metallic states are predicted correctly for both DCD and TCD models. Our results indicate that models trained on DCD exhibit superior transferability over which trained on TCD, suggesting DCD models are more reliable when confronting with unseen structures.

2.4 DCD Exhibits Simpler Radial and Angular Dependencies in the Core Region↩︎

In this section, we investigate why and to what extent \(\Delta\)-SAED is beneficial. We attribute its advantage to the simpler radial and angular dependencies of DCD in the core region. In Fig. 2 (b) the distribution of the \(\Delta_{rel}\varepsilon_{mae}\) for metals and insulators in the MP dataset, which occur in approximately equal proportions, are present separately. \(\Delta\)-SAED achieves better performance on insulators, with an average \(\Delta_{rel}\varepsilon_{mae}\) of 17.2%, compared to 11.3% for metals. The performance disparity of \(\Delta\)-SAED among insulators and metals suggests that \(\Delta\)-SAED works in line with the \(\Delta\)-learning principle, i.e., incorporating physical priors would smooth the regression task. This is evidenced by the empirical observation that insulators’ \(\rho_d\) is closer to zero than metals’ due to less complex charge transfer process, indicating that for insulators the smoothness advantage of \(\rho_d\) over \(\rho_t\) is more prominent than that of metals.

In Fig. 4 we displays 4 Si allotropes’ radial distributions, defined by grid points within 2Å of each atom, for \(\rho_t\), \(\rho_d\) and their corresponding absolute error \(|\rho_t - \hat{\rho}_t|\), \(|\rho_d - \hat{\rho}_d|\). We choose mp-1202745 and mp-988210 as representatives of the 18 test structures in Fig. 3 becuase their MAEs are close to chemical accuracy and their radial distributions resemble those of other test structures, while mp-988210 and mp-1403870 are are chosen as examples with outlier errors. Both TCD and DCD models struggle to predict the charge density and properties for the two allotropes, because the MPNN charge density model has less transferability for complex geometry chemical environment than complex compound combinations, as also reported in [37], though \(\Delta\)-SAED still outperforms. Overall, \(\rho_d\) is smoother as function of \(r\) and shows a narrower spread, indicating simpler angular and radial dependencies. This confirms that SAED provides useful physical prior information, making \(\rho_d\) easier to learn. The radial error distributions in Fig. 4 (a), 4 (b) show that, near atomic centers, the error increases for the TCD model but decreases for the DCD model. As expected, \(\Delta\)-SAED performs better around atomic centers, although the overall average error of DCD is not always lower than that of TCD (mp-1244933 for example). Remarkably, even in such cases, the DCD model achieves better performance on structural properties. This suggests \(\Delta\)-SAED’s reliability for non-self-consistent calculations, and the limitation of relying solely on the average charge density error as an evaluation metric for charge density models.

In Fig. 5 we show the per-atom \(\varepsilon_{mae}\) of the MP test set, by integrating within a sphere of radius 2Å centered at each atom using Eq. 1 . We find \(\Delta\)-SAED improves across the periodic table consistently, with an average 0.08% of absolute improvement \(\Delta\varepsilon_{mae} = \varepsilon_{mae} ^{TCD}- \varepsilon_{mae}^{DCD}\), despite as element gets heavier the sampling probability declines. The improvement is not because SAED compensates for inefficient uniform sampling of the TCD model around the core, as the radial error distributions for Si allotrope in Fig. 4 shows that TCD model can predict core charge densities reasonably well, though not as accurately as DCD model does. Rather, the improvement reflects the limited expressive power of the model itself in capturing the complex radial and angular dependence of \(\rho_t\), and SAED provides useful physical priors across the periodic table to compensate for the model architecture’s expressive limitations.

3 Discussion↩︎

This work introduces \(\Delta\)-SAED, a robust and cost-free method to improve the accuracy and transferability of charge density and derived properties beyond baseline models. Inspired by the spirit of \(\Delta\)-learning, \(\Delta\)-SAED applies superposition of atomic electron density as physical priors and trains on the difference charge density instead of the total charge density. This reduces the complexity of training target’s radial and angular dependencies at core regions, mitigating machine learning models’ limited expressiveness in complex fitting task. Although advantages have so far been demonstrated only on grid-based architectures, we expect \(\Delta\)-SAED to also work for basis-based architectures, as it provides an exact component of the charge density, bypassing the less accurate basis-expansion formulation.

Looking forward, \(\Delta\)-SAED paves the way for building charge density foundation models that integrates calculations from different DFT codes. Given a structure and functional, TCD is not unique in different DFT implementations due to approximations like pseudopotentials and basis sets. In contrast, DCD is, in principle, free from this issue once calculations converge, since it captures the charge transfer from SAED–a well defined state by solving the radial Schrödinger equations for atoms individually.

4 Methods↩︎

4.1 Datasets↩︎

The total charge densities of the NMC and QM9 datasets are provided by Jørgensen and Bhowmik[26], [30]. The total charge densities of the MP dataset are calculated by [34]. The NMC dataset consists of 2000 randomly sampled nickel manganese cobalt oxides containing varying levels of lithium content, with training, validation, and test splits of size 1450, 50, and 500 respectively. The QM9 dataset contains 133845 small organic molecules, with training, validation, and test splits of size 123835, 50, and 10,000 respectively. The MP dataset contains 117173 structures, splited randomly into training, validation, and test sets with sizes 114661, 512, and 2000 respectively. Our MP dataset is different from the one used in [27] because structures with available charge density information are updated frequently.

4.2 Training↩︎

The training settings, including model hyper parameters and optimizing schedules, follow those reported by Koker[27]. For every dataset, the settings used for training \(\rho_d\) are the corresponding same settings for \(\rho_t\). The training setting for the Si dataset is identical to that of the NMC dataset. The model parameters for the NMC TCD and QM9 TCD models are adopted from [27]. Notably, the average \(\varepsilon_{\text{mae}}\) for TCD models in Tab. 1 differs from those reported by Koker, as our evaluation uses all probing points in line with DFT, whereas Koker employed a random sampling for the integral in Eq. 2 . The structures for the Si allotropy dataset and all trained model parameters are available on Zenodo (10.5281/zenodo.17180513).

4.3 DFT Calculations↩︎

All DFT calculations were performed using plane wave based VASP code[38] with PBE functional[39] and projector-augmented wave (PAW) pseudopotentials[40], [41]. All Si allotropes are calculated with a \(6\times6\times6\) Monkhorst-Pack k-point mesh[42]. After performing self-consistent-field calculations, \(\rho_a\) is computed using the same pseudopotential, with the electronic relaxation disabled (NELM = 0), the charge density initialization scheme set to SAED(ICHARG = 12), the DFT grid forced consistent. Notably, each element in the MP dataset uses a unique, consistent pseudopotential across all structures, avoiding any physical discrepancies due to pseudopotential mixing.

Acknowledgements↩︎

We acknowledge financial support from the National Key R&D Program of China (No. 2022YFA1402901), NSFC (grants No. 12188101), Shanghai Science and Technology Program (No. 23JC1400900), the Guangdong Major Project of the Basic and Applied Basic Research (Future functional materials under extreme conditions–2021B0301030005), Shanghai Pilot Program for Basic Research—Fudan University 21TQ1400100 (23TQ017), the robotic AI-Scientist platform of Chinese Academy of Science, and New Cornerstone Science Foundation.

References↩︎