ECD: A Machine Learning Benchmark for Predicting Enhanced-Precision Electronic Charge Density in Crystalline Inorganic Materials

Authors: Pin Chen, Zexin Xu, Qing Mo, Hongjin Zhong, Fengyang Xu, Yutong Lu
Primary area: datasets and benchmarks
Keywords: Electronic Charge Density, Crystalline Inorganic Materials, Graph Neural Network, Dataset

Abstract:
Supervised machine learning techniques are increasingly being adopted to speed up electronic structure predictions, serving as alternatives to first-principles methods like Density Functional Theory (DFT). Although current DFT datasets mainly emphasize chemical properties and atomic forces, the precise prediction of electronic charge density is essential for accurately determining a system's total energy and ground state properties. In this study, we introduce a novel electronic charge density dataset named ECD, which encompasses 140,646 stable crystal geometries with medium-precision Perdew–Burke–Ernzerhof (PBE) functional data. Within this dataset, a subset of 7,147 geometries includes high-precision electronic charge density data calculated using the Heyd–Scuseria–Ernzerhof (HSE) functional in DFT. By designing various benchmark tasks for crystalline materials and emphasizing training with large-scale PBE data while fine-tuning with a smaller subset of high-precision HSE data, we demonstrate the efficacy of current machine learning models in predicting electronic charge densities.
The ECD dataset and baseline models are open-sourced to support community efforts in developing new methodologies and accelerating materials design and applications.

Decision: accept (oral)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730250455393

thread_type:
official_review

<thread_object>
creator: reviewer_tq9x

<content>
summary:
The author introduces a new benchmark for predicting electronic charge density for crystalline inorganic materials. The dataset includes 140,646 stable crystal geometries with medium-precision PBE functional, and high accuracy HSE from subsets of the meticulously curated geometries.

strengths:
1. The computational efforts used in this study is immense, and the resulting benchmark could be a valuable asset to the ML community.
2. Interesting task on predicting the gaps across different functionals.

weaknesses:
1. The work provides minimal justification for predicting electrons instead of the Kohn-Sham Hamiltonian.
2. There is no comparison with other quantum chemistry benchmarks, like QH9. This will make the paper's calculation comparable in the literature.
3. The EXP-id lacks detail on its 47 entries, which is essential for practical use. Additionally, the assessment of practical application is sparse.
4. The writing lacks cohesion, with many GPT-like passages. Figure 2 is unclear: does moving up or down the ladder imply increased simplicity?
5. The metrics are limited to MAE for charge density and acceleration ratio. Other measures, such as total energy, HOMO-LUMO gap, electronic spatial extent, or dipole moment, should be included, as this is particularly important for practical applications.

questions:
1. Is there a fundamental difference between predicting the Kohn-Sham Hamiltonian and electron density? The author should discuss the computational challenges and benefits of predicting electron density, as several prior works have focused on predicting the Kohn-Sham Hamiltonian.

2. Can the MAE of electron density be compared to that of a Kohn-Sham Hamiltonian prediction model?

3. In line 528, "due to the high dimensionality of charge density data, often involving over 1,000 nodes, the complexity of handling node and edge interactions presents significant computational challenges"—what do "node" and "edge" refer to here?

4. What does the regression output look like for a given structure? A more detailed example or description of the matrix's size and shape would be helpful.

5. How does the model perform on predicting HOMO and LUMO?

6. Could the author provide a more detailed discussion on error? With an HSE error of 1 eV, what is the impact on practical applications like energy band predictions? For a usable model, is 1 eV sufficient for downstream applications?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730539512523

thread_type:
official_review

<thread_object>
creator: reviewer_nk1z

<content>
summary:
This paper introduces the Electronic Charge Density (ECD) dataset, a benchmark designed for predicting electronic charge densities in crystalline inorganic materials. The dataset comprises 140,646 stable crystal structures computed with PBE functional, alongside a subset of 7,147 structures calculated with HSE functional. Using the ChargE3Net model as a baseline, the authors evaluate its performance across multiple tasks, including prediction accuracy on the ECD dataset and the acceleration of DFT computations. Overall, this paper makes a valuable contribution to machine learning in quantum chemistry and computational materials science, and I recommend acceptance conditional on addressing the concerns outlined below.

strengths:
* Originality: Good. The data generated in this work is indeed original.
* Quality: Good. The functionals PBE and HSE used in this work are indeed the widely accepted ones to compute crystalline inorganic materials. Thus, the data generated in this work are of good qualities.
* Clarity: The paper is well-structured, with a clear explanation of the dataset’s construction, tasks, and evaluation metrics. 
* Significance: This is a significant contribution to this area. Due to the computational cost to carefully converge the charge densities, it is not easy to generate or to find publicly available data sets. Thus, this work will clearly benefit the researchers in this area.

weaknesses:
* Limited High-Precision Data: I totally understand that the HSE calculations are much more expensive, but it would be much better if the authors can keep increasing the number of calculations done with HSE functional.

questions:
Although I generally believe this paper is a good contribution to this area, I still have the following concerns and I hope the authors may address:
1. **Codes and Data Availability**: The links to the workflows and datasets are no longer valid, thus I could not review them. It would be nice if the authors provide new links to them.
2. **Computational Setup**:

    a) Why did the authors decide to mix charge densities with PBE and PBE+U functionals? Are you sure that the charge densities at these two levels are compatible to be mixed up? More information for you to consider: In public datasets like MP, they mixed PBE and PBE+U functionals mostly because they justified that they can align the formation energies of materials with further compatibility corrections. However, here you are computing charge densities, so I am skeptical if these two levels of theories should be mixed up.

    b) The choice of hybrid functional is not justified. There are multiple hybrid functionals available on the market, including PBE0, B3LYP, the HSE you are using, etc. Why did you pick up HSE as your functional? In addition, you should specify which version of HSE you used, HSE03 or HSE06?

    c) Did you apply any convergence test on the charge densities with respect to the energy cutoff and Brillouin zone sampling? My impression is that the 520eV cutoff you are using might not have converged the charge densities.

3. **Data Storage Schema**: In large-scale calculations used for machine learning, it would be nice if you include the information of the data storage schema in your supplementary information. For example,  what information you stored and what units for each of the information. 

4. **Experimental Setups**

    a) **Question on OOD evaluation using GNoME dataset**: The authors *somehow* selected 2k materials from GNoME dataset, which does not make sense to me that these are naturally *out of distribution*. I suggest the authors conduct structure matching test to see how many in the 2k materials are truly out of distribution. Otherwise, this is not good enough to claim that this is an evalutation on OOD materials.

    b) There is a property that the charge densities have to satify:
    $$\int \rho(\mathbf{r}) \mathrm{d}\mathbf{r} = \text{No. of electrons}$$
    How did you enforce this in your model?

    c) The evaluation of the prediction of charge densities is based on the MAE with respect to the ground truth, which makes sense in terms of developing ML models. However, I am interested to see how much error in total energy or band gap if the predicted charge densities are put back to DFT iterations. These are the realistic application of the charge densities, otherwise one can not really infer any useful information from predicted charge densities. The authors may consider commenting on this or conducting one or a few examples to show the performance.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671134203

thread_type:
official_review

<thread_object>
creator: reviewer_4vph

<content>
summary:
The paper introduces a benchmark dataset ECD to enhance machine learning models for predicting electronic charge densities for crystalline materials. ECD includes over 140K crystal structures with medium-precision PBE data and 7147 high-precision HSE data entries. The authors propose tasks for training, fine-tuning, and out-of-distribution testing to improve predictive accuracy and speed up DFT calculations, using ChargE3Net to demonstrate significant improvement. This dataset offers a robust foundation for advancing scalable, high-accuracy electronic structure predictions in materials science.

strengths:
1. The ECD dataset is a substantial and valuable contribution, offering over 140K PBE-calculated crystal structures and 7147 HSE-calculated structures, providing a robust foundation for electronic charge density prediction.

1. The paper conducts a comprehensive benchmarking of multiple models, assessing their performance across various tasks and metrics, including both in-distribution and out-of-distribution scenarios, which enhances the reliability and applicability of the findings.

1. The paper demonstrates that fine-tuning models trained on PBE data with high-precision HSE data improves predictive accuracy. This approach can guide future research on leveraging high-cost data sparingly to optimize model performance.

1. The paper highlights the model’s performance on wet-lab experimental data and out-of-distribution data, validating its potential for applications in real-world materials discovery and screening.

weaknesses:
1. The description of ECD-HSE-id seems missing in Section 2.5. I assume that the data split is consistent with ECD-PBE HSE-id and ECD-PBE HSE tuning-id, but it would be great if the authors could clarify that. 

1. ChargeE3Net surpasses its invariant counterparts and other non-equivariant models marginally but significantly increases the computation complexity, thus increasing the time of cost for computation. For example, comparing CrystalNet-TL to ECD-PBE HSE, the MAE is 0.70 vs 0.65 while the time per structure is 1.58s vs 14.4min. 

1. In the OOD performance section, while the authors demonstrate the model's generalization, there is minimal analysis of failure cases or instances where the model struggled. Understanding these limitations in detail would be valuable.

1. While the ECD dataset is a major contribution, data augmentation techniques or other methods to improve model generalization are not discussed. Addressing potential improvements for model robustness could make the dataset more valuable for future research.

questions:
1. ChargeE3Net is compared to other models in Table 2, whereas the detailed description for invDeepDFT, DeepDFT, and invChargeE3Net is missing. Can the authors justify the choice of using DeepDFT for comparison and elaborate on why it is an appropriate baseline? Besides, I guess invChargeE3Net is the invariant version of ChargeE3Net, but it would be better to clarify that in the manuscript. 

1. Given the experimental results, it seems that equivariant models do not significantly decrease the errors in ECD prediction. Considering the lower computation complexity of invariant models, equivariant models may not have an edge over invariant models. Can the authors provide the DFT calculation acceleration of invChargeE3Net and make a comprehensive comparison between invariant and equivariant models taking account of both accuracy and speed?

1. Given the results in Table 3, the acceleration from ML-predicted densities is moderate as the achieved ratio is around 0.7. Furthermore, the achieved ratio is around one order of magnitude larger than the optimal ratio. Can the authors provide insight of how to further improve the acceleration or discuss what factors prevent the model from achieving the lower bound on the number of steps?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730679201616

thread_type:
official_review

<thread_object>
creator: reviewer_dhoo

<content>
summary:
Paper proposes an electronic charge density dataset of 140,646 crystals using DFT PBE functional and 7,147 crystals with HSE functional. Paper trains the ChargE3Net on PBE dataset and with and without fine-tuning on HSE dataset, and compares with original ChargE3Net on MP data. Paper also shows that using the charge density predicted by ChargE3Net (as initialization) can accelerate DFT calculations. Paper compares band gaps predicted by their model with other published results.

strengths:
Paper is clearly written. Paper presents dataset which was generated by significant computation expense. The problem of materials design and prediction is significant.

weaknesses:
Charge density data is high memory, and may be computationally expensive to train ML models. (authors admit this in limitations section). The optimal vs achieved ratio of Table 3 seem quite far from each other- it seems to indicate that initializing dft calc with ML charge density model prediction does not improve efficiency as much as one would like.

questions:
what is invChargE3Net in Table 2? could authors provide a clear definition or explanation of invChargE3Net in the paper, as this term is not defined.

What is ECD-HSE-id in Table 2? could authors clearly define or explain ECD-HSE-id in the paper, as this term is not explained.

Reviewer would suggest to authors to include a detailed description of structure of each dataset entry, possibly in appendix, if not already included.

Could authors provide more context about the Choudhary 2018 HSE-based dataset model?, including reasons for its performance characteristics, such as its error and time performance.

For Table 5, are there error bars? I am assuming the eV is band gap error, what is the acceptable error for band gap prediction? 

What is reason for discrepancy between optimal and achieved ratio in Table 3, and how was optimal ratio determined?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a repository for predicting Electronic Charge density prediction of crystalline materials, comprising of 140K crystal structures with medium precision PBE and high precision HSE data.

Strengths:
1.  This dataset offers a robust foundation for advancing scalable, high-accuracy electronic structure predictions in materials science. The paper conducts a comprehensive benchmarking of multiple models, assessing their performance across various tasks and metrics, including both in-distribution and out-of-distribution scenarios, which enhances the reliability and applicability of the findings [Reviewer 4VpH]

2.  Using the ChargE3Net model as a baseline, the authors evaluate its performance across multiple tasks, including prediction accuracy on the ECD dataset and the acceleration of DFT computations. Overall, this paper makes a valuable contribution to machine learning in quantum chemistry and computational materials science. This is a significant contribution to this area. Due to the computational cost to carefully converge the charge densities, it is not easy to generate or to find publicly available data sets. Thus, this work will clearly benefit the researchers in this area. [Reviewer NK1Z]

3. The dataset includes 140,646 stable crystal geometries with medium-precision PBE functional, and high accuracy HSE from subsets of the meticulously curated geometries. The computational efforts used in this study is immense, and the resulting benchmark could be a valuable asset to the ML community [Reviewer Tq9X]

Authors have patiently and objectively addressed questions and concerns raised by the reviewers during the rebuttal phase.

Weakness: 
Most of the weaknesses highlighted by the reviewers were addressed by the authors during the rebuttal phase. However, from the results, the MAE of CrystalNet-TL  is close to ECD-PBE HSE (0.70 vs 0.65) but the required time is much lower than ECD-PBE HSE [Reviewer 4VpH]. Although this is a minor weakness.