Density Functional Theory (DFT), specifically the Kohn-Sham formulation (KS-DFT), is a cornerstone of computational materials science for predicting electronic structures. The ground state electron density, derived from KS-DFT, is rich in information about material properties, ranging from structural parameters and elastic constants to magnetic properties and phonon spectra. It also serves as the foundation for calculations of excited-state phenomena. However, the cubic scaling of conventional KS-DFT calculations with the number of atoms presents a significant computational barrier, particularly for large and complex systems. To overcome this, several approaches have been developed, including techniques that avoid explicit Hamiltonian diagonalization and methods that improve parallel scalability on high-performance computing platforms. Despite advances, routine application to complex systems using modest resources remains challenging. Machine learning (ML) offers an attractive alternative by creating surrogate models to replace computationally expensive KS-DFT calculations. Significant research has focused on ML models predicting atomic energies and forces, leading to ML-based interatomic potentials. Direct ML prediction of the ground state electron density, using KS-DFT simulations as training data, is particularly appealing because it directly captures a wealth of material information. Previous approaches to ML electron density prediction have employed different density representations, either atom-centered basis functions or grid point predictions. This work utilizes the latter, which, while challenging for inference, has shown great promise for bulk materials. A further challenge is ensuring the predicted electron density respects physical transformations (rotation, translation, permutation invariance). This is often addressed using equivariant neural networks or invariant descriptors; this study uses the latter approach. The primary challenge in creating accurate ML surrogate models for the ground state electron density remains the significant offline cost of generating training data from KS-DFT calculations. Previous work has relied exclusively on data from large systems, limiting its scalability to complex materials like multi-principal element alloys. This paper introduces a machine learning model that accurately predicts the ground state electron density of bulk materials at any scale while providing uncertainty quantification. The high cost of training data for larger systems is addressed through a transfer learning approach, combining a large quantity of inexpensive data from small systems with a smaller amount of data from larger systems, significantly reducing the overall cost. Bayesian Neural Networks (BNNs) are employed to achieve uncertainty quantification, providing a measure of confidence in predictions. Simple, scalar product-based descriptors are used to represent the local atomic neighborhood, avoiding handcrafted descriptors used in previous works. Thermalization, using ab initio molecular dynamics (AIMD) simulations at various temperatures, provides a comprehensive sampling of the descriptor space.

Numerous studies have explored the use of machine learning to accelerate electronic structure calculations. Early work focused on developing ML-based interatomic potentials that predict energies and forces, matching KS-DFT calculations with ab initio accuracy. These potentials have been successfully employed in molecular dynamics simulations. A separate line of research has focused on directly predicting the ground state electron density using ML models. This approach is attractive because the electron density contains extensive information about material properties that goes beyond energies and forces. Different strategies have been employed to represent the electron density, such as atom-centered basis function expansions and grid point predictions. While atom-centered expansions offer compact representation, they require optimized basis sets tailored to specific chemical species. In contrast, grid-point prediction offers flexibility but necessitates inference on a large number of grid points. Equivariant neural networks have been proposed to address the need for preserving physical transformations in the predicted electron density. Alternatively, invariant descriptors can be used, as employed in this research. The generation of sufficient and representative training data remains a bottleneck. Previous studies on electron density prediction have often relied on data exclusively from large systems, limiting scalability to complex systems. Furthermore, previous works have often relied on heuristic choices of descriptors, lacking a systematic method of descriptor selection. This research addresses the limitations of previous works by using a novel combination of methods to systematically address the challenges in machine learning for electron density prediction.

This research utilizes a machine learning model to predict the ground state electron density of bulk materials across various scales, including uncertainty quantification. The model employs a transfer learning approach to mitigate the high computational cost associated with generating training data using KS-DFT for large systems. The process begins by training the model on a substantial quantity of data generated from simulations of small systems. This inexpensive data provides a foundation for the model. Subsequently, a portion of the model is retrained using a smaller dataset from simulations of larger systems. This two-stage training process significantly reduces the overall computational cost. Bayesian Neural Networks (BNNs) are used to quantify the uncertainty in predictions, enabling assessment of confidence levels, especially for large systems where direct validation against KS-DFT is impractical. The model's input consists of atomic neighborhood descriptors derived from the local atomic environment around each grid point. These descriptors are simple scalar product-based features encompassing distance and angle information, ensuring invariance under translations, rotations, and permutations of atomic indices. A key aspect of this methodology is the systematic selection of an optimal set of descriptors for a given dataset. This selection process, detailed in the paper, avoids the trial-and-error approaches used in earlier studies. The selection is guided by the nearsightedness principle and the observation that electron density at a grid point is largely influenced by nearby atoms. To adequately sample the descriptor space, the research employed thermalization through ab initio molecular dynamics (AIMD) simulations at various temperatures. AIMD simulations were performed using the SPARC code with the GGA PBE exchange-correlation functional and ONCV pseudopotentials. Appropriate mesh spacing and convergence criteria were employed. For the AIMD runs, a standard NVT-Nosé Hoover thermostat was used along with Fermi-Dirac smearing. Snapshots of atomic configurations and electron densities were collected at intervals to avoid correlation. AIMD simulations were performed at temperatures spanning a wide range, enabling sampling of diverse system configurations. Data for systems with defects (monovacancies, divacancies, dislocations, grain boundaries) and under strain were also generated. The ML model is a Bayesian neural network, trained to map the atomic neighborhood descriptors to the electron density at each grid point. The use of BNNs provides a systematic route to uncertainty quantification through the stochastic nature of its parameters. The uncertainty is decomposed into aleatoric (inherent data variability) and epistemic (model parameter uncertainty) components. Aleatoric uncertainty is modeled using a heteroscedastic noise model, capturing spatial variations in noise. The epistemic uncertainty is assessed using samples from the posterior distribution of the model parameters. The post-processing of the ML-predicted electron density involves scaling the density to match the total number of electrons before calculating quantities of interest like the total ground state energy. The Harris-Foulkes formula is used for energy calculations. The transfer learning approach is implemented by first training on smaller systems, and then retraining a portion of the network using larger system data, optimizing data use and reducing cost.

The study demonstrates the effectiveness of the proposed ML model through predictions for bulk aluminum and silicon-germanium (SiGe) alloy systems. The model accurately predicts electron densities for a wide range of test systems well beyond the training data, including systems with defects, various alloy compositions, and those significantly larger than systems included in training. The Root Mean Squared Error (RMSE) and L¹ norm per electron are reported, showing high accuracy for both aluminum and SiGe. The model shows remarkable accuracy even for systems containing thousands of atoms (e.g., 1372 atoms for aluminum and 512 atoms for SiGe). The model's accuracy extends to systems containing defects (mono-vacancies, di-vacancies, grain boundaries, dislocations) and alloy compositions beyond those used for training. The error magnitudes remain low, even for significant variations in composition. The model accurately predicts not only the electron density but also ground state energies (within chemical accuracy), lattice parameters (to a fraction of a percent), and bulk moduli (within 1% of DFT values). Uncertainty quantification results show comparable magnitudes of aleatoric (inherent data variability) and epistemic (model uncertainty) uncertainties. The aleatoric uncertainty is higher near nuclei and defect sites, reflecting higher variability in the data at these locations. Epistemic uncertainty is higher near nuclei due to a limited amount of training data from such regions. Adding data from systems with defects significantly reduces error and uncertainty at defect sites. Transfer learning is shown to reduce training data generation time by more than 50%, while maintaining model accuracy. The ML model scales linearly with system size, offering a significant computational advantage (more than two orders of magnitude faster than DFT) for systems with a few hundred atoms. This advantage extends to multi-million atom systems where DFT calculations become impractical. Electron densities for four-million atom aluminum and one-million atom SiGe systems were successfully predicted using the model, with comparable uncertainty levels to smaller systems. This demonstrates the ability to predict electronic structure at scales inaccessible to DFT.

The findings demonstrate the potential of this uncertainty-quantification-enabled machine learning approach to revolutionize electronic structure prediction in materials science. The ability to make accurate and confident predictions for systems with millions of atoms addresses a major bottleneck in computational materials science, opening doors to simulations that were previously computationally infeasible. The transfer learning strategy significantly reduces the computational cost of training data generation, making this approach more accessible to researchers. The uncertainty quantification feature allows for a more rigorous assessment of the reliability of predictions, improving the trust in the results. The method's accuracy in predicting not only electron densities but also derived physical parameters (e.g., energies, lattice constants, bulk moduli) further validates its reliability. The success with diverse material systems (metals and semiconductors) and system configurations (defects, various compositions) highlights the model's transferability and robustness. The linear scaling of the model with system size makes it suitable for high-throughput computations and the exploration of large compositional spaces.

This research presents a highly efficient and accurate machine learning model for predicting ground state electron densities in bulk materials. The combination of transfer learning, Bayesian neural networks, and systematically selected descriptors results in a model that is both computationally efficient and robust, capable of handling multi-million atom systems. The uncertainty quantification capabilities add a layer of reliability to the predictions. Future work could focus on extending this approach to other material classes, exploring active learning strategies for further optimization, and applying this model to investigate complex material phenomena at unprecedented scales.

While the model shows remarkable accuracy and efficiency, some limitations should be noted. The model's accuracy might decrease when predicting systems significantly far from the training data composition. The transfer learning approach is inherently limited by the largest feasible system size that can be simulated using DFT. The accuracy of the predictions depends on the quality of the DFT data used for training. Furthermore, although the model provides uncertainty quantification, it does not explicitly model all potential sources of error in the DFT calculations themselves.