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

Teddy Koker,Keegan Quigley,Eric Taw,Kevin Tibbetts,Lin Li

2024-05-14

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.

Computational Physics,Materials Science,Machine Learning

What problem does this paper attempt to address?

This paper mainly addresses the accuracy and efficiency issues of computing electronic density distribution using Density Functional Theory (DFT) in materials science. Due to the complexity of self-consistent field iteration, DFT calculations are computationally expensive for large-scale systems. The paper proposes a neural network model called ChargE3Net, which is an E(3) group equivariant graph neural network model, for predicting the electronic density of atomic systems. ChargE3Net is able to learn higher-order equivariant features, improving prediction accuracy and model expressive power.

Research shows that ChargE3Net outperforms previous works on various molecular and material datasets, and after training on the Materials Project database containing over 100,000 materials, the model captures the complexity and variability in the data. Using ChargE3Net as initialization for DFT calculations significantly reduces the number of self-consistent field iterations and lowers computational costs. Additionally, non-self-consistent DFT calculations can achieve near-DFT performance at a lower cost.

Furthermore, the paper points out that ChargE3Net improves prediction accuracy when dealing with systems with high-angle variations. By introducing higher-order equivariant features, the model is able to better simulate electronic density, especially in materials involving covalent bonds. ChargE3Net's time complexity scales linearly with system size, making it capable of handling large-scale systems with tens of thousands of atoms, surpassing the feasibility of first-principle calculations.

In conclusion, this paper addresses the efficiency issues in DFT calculations and proposes a new machine learning approach to accelerate quantum mechanical computations in the process of materials discovery.