A Recipe for Charge Density Prediction

Authors: xiang  fu
andrew  rosen
kyle  bystrom
rui  wang
albert  musaelian
boris  kozinsky
tess  smidt
tommi  jaakkola
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 22:17:06

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.

Main paper: 
 introduction: 
 Density functional theory (DFT) is a computational quantum chemistry method that has enabled countless advancements in the chemical sciences by providing a tractable means to calculate the electronic structure of molecules and materials (Jain et al., 2016) . The central concept in DFT is the charge density, a fundamental quantity from which all derivable ground-state physicochemical properties of a system, such as energy and forces, can, in principle, be derived. The most widely used Kohn-Sham formalism (Kohn & Sham, 1965) of DFT offers a reasonable balance between accuracy and computational efficiency among conventional DFT workflows. However, it still scales with a complexity of roughly O(N 3 e ) where N e is the number of electrons, rendering it computationally expensive and limiting its viability for both large-scale systems and long-timescale ab initio molecular dynamics simulations.In DFT, the solution to the Kohn-Sham equations are reliant on an iterative calculation to identify the charge density that minimizes the potential energy functional for a given atomic configuration. This process, known as converging the selfconsistent field, is the main computational expense within DFT. With a machine learning (ML) model that can effectively bypass the Kohn-Sham equations by accurately and efficiently predicting the charge density, the number of steps required to converge the ground-state electron density can be drastically reduced or potentially eliminated altogether by using the predicted charge density as the initial guess. If accurate enough, a machine-learned charge density could also be used to directly predict electronic structure properties, such as the band gap, band structure, and electronic density of states of a material. Furthermore, the charge density itself can provide an enormous amount of insight into a molecule or material. From the charge density, partial atomic charges, dipole moments, atomic spin densities, and effective bond orders can all be directly computed through one of several population analysis methods (Fonseca Guerra et al., 2004; Limas & Manz, 2018) . For some materials discovery tasks, the charge density can also be a crucial descriptor depending on the application area (Shen et al., 2023; Yao et al., 2024) . Therefore, efficient and accurate representations and ML models for charge density prediction are highly desirable as a means of accelerating the discovery of promising molecules and materials.In machine learning workflows, the charge density is a volumetric, data-rich object, usually represented as voxels with a grid resolution of around 0.1 Å (Jørgensen & Bhowmik, 2022; Shen et al., 2022) . This poses a challenge, as even relatively small molecules and materials can require hundreds of thousands to millions of grid points to represent the charge density at this (relatively coarse) resolution. At the same time, small deviations in the charge density that result from too coarse of a representation can have a substantial impact on the energy and other derivable properties. This need for both efficiency and accuracy creates a significant challenge for ML methods. Existing literature has mainly focused on two approaches for learning to predict the charge density. The first approach (orbital-based), illustrated in Figure 1 (a), is to predict atomic orbital basis set coefficients by regressing over coefficients extracted from DFT data (Fabrizio et al., 2019; Qiao et al., 2022; Rackers et al., 2023; Cheng & Peng, 2023; del Rio et al., 2023) . The atomic orbital basis functions are based on the composition of radial functions and spherical harmonics. Under this scheme, the charge density is represented as a set of spherical fields centered around each atom. The real space charge density voxel can be constructed by overlaying the spherical fields and evaluating at each grid point. For orbital-based ML models, both the prediction of the basis set coefficients and the evaluation of the spherical fields are relatively scalable, making this approach efficient at inference time. However, this approach can suffer from sub-optimal accuracy due to the limited representation power of the chosen basis set. In particular, it is challenging for the atom-centered atomic orbitals to complex electronic structures between atoms.The second approach (probe-based) (Gong et al., 2019; Jørgensen & Bhowmik, 2022; Koker et al., 2023; Pope & Jacobs, 2023) , illustrated in Figure 1 (b), is to predict the charge density by inserting "probe nodes" at all grid coordinates of the charge density voxel and applying graph message passing between the atoms and these probe nodes. Finally, the scalar charge density at each grid coordinate is predicted through node-wise readout over the probe nodes. This approach, while expressive and accurate, is computationally expensive. To see why, recall that the number of grid points in the charge density voxel is usually very large for even a small atomic system. Conducting neural message passing over millions of nodes is both computationally and memory intensive. The large number of nodes usually requires sampling a subset of grid points from the charge density voxel (Figure 1 (b), right) in each training or inference step (Jørgensen & Bhowmik, 2022) .This paper aims to address this accuracy-efficiency dilemma with a new recipe for building representations and ML models for charge density prediction. We identify three key ingredients:1. We represent the charge density using an atomic orbital basis set (spherical fields centered at each atom) to leverage its efficiency and equivariant properties. Beyond orbitals placed at the atomic coordinates, we further introduce virtual orbitals to improve expressivity.In other words, we also place spherical fields centered at coordinates other than the atomic centers, while ensuring the placement algorithm is SE(3)-equivariant.2. We use domain-informed and expressive basis sets. In particular, we construct an even-tempered Gaussian basis from an atomic orbital basis set. This allows us to smoothly control the expressivity of the atomic orbitals and enable flexible accuracy-efficiency tradeoffs. We make the basis set exponents learnable to further improve expressivity.3. We use a high-capacity equivariant neural network architecture (eSCN (Passaro & Zitnick, 2023) ), which enables efficient training and inference with features of high tensor order for a large dataset.We apply our recipe to the widely used QM9 charge density benchmark (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014; Jørgensen & Bhowmik, 2022) . Our method outperforms existing state-of-the-art methods while being around 30× faster. Furthermore, we can flexibly trade off accuracy and efficiency by adjusting the model/basis size; in doing so, we achieve up to 171× efficiency compared to state-of-the-art methods with only a slight degradation in accuracy. This tunability is valuable, as different applications, material classes, and available computing resources may require drastically different levels of accuracy in the charge density prediction. We conduct an ablation study to justify the significance of each proposed ingredient. 
results: 
 We highlight the experimental results of our proposed method in the main paper and defer the related works and method sections to the appendix. Our experiments on the QM9 charge density benchmark aim to validate the effectiveness of our proposed recipe in both accuracy and efficiency. We refer to our method as SCDP models, which stands for Scalable Charge Density Prediction models.The QM9 charge density dataset (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014; Jørgensen & Bhowmik, 2022) contains charge density calculations for 133,845 small organic molecules using the Vienna Ab initio Simulation Package (VASP). We adopt the original split, where 123,835, 50, and 10,000 data points are used for training, validation, and testing, respectively. There are on average 18 atoms in each molecule and 666,462 grid points in each charge density voxel. The entire dataset takes 1.1 TB of disk space. Following previous works (Jørgensen & Bhowmik, 2022; Koker et al., 2023) , we benchmark the prediction accuracy with the normalized mean absolute error, defined as:EQUATIONwhere the integration is approximated by summing over the full charge density voxel. We benchmark the efficiency of different methods by the number of molecules predicted per minute (Mol. per min.) on a single NVIDIA A100-80GB-PCIe GPU for the QM9 test split.We compare SCDP to several previous works on the QM9 charge density prediction benchmark (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014; Jørgensen & Bhowmik, 2022) . i-DeepDFT, e-DeepDFT (Jørgensen & Bhowmik, 2022) , and ChargE3Net (Koker et al., 2023) are probe-based methods with different backbone architectures: i-DeepDFT Accuracy (1-NMAE [%]) uses SchNet (Schütt et al., 2021 ), e-DeepDFT uses PaiNN (Schütt et al., 2021) , while ChargE3Net uses higherorder equivariant features under the tensor field network framework (Thomas et al., 2018; Geiger & Smidt, 2022) .K = 4, L = 3, = 2.0 K = 8, L = 6, = 2.0 K = 8, L = 6, = 1.5 VO, K = 8, L = 6, = 2.0 VO, K = 8, L = 6, = 1.5 VO, K = 8, L = 6, =InfGCN (Cheng & Peng, 2023) combines GTOs and a shallow network for probe-based inference. It also has a more efficient but less accurate GTO-only variant. GPWNO (Kim & Ahn, 2024) combines GTOs and plane-wave basis sets but still requires a large number (64,000) of probe nodes for constructing the plane wave prediction. NMAE results for i-DeepDFT, e-DeepDFT, and ChargE3Net are from (Koker et al., 2023) . NMAE results for InfGC and GPWNO are also from the original papers, which uses a different test split from the default QM9 test split (last 1,600 molecules from the QM9 test split). We benchmark the efficiency of baseline models on our hardware when the source code and pretrained model are publicly available (ChargE3Net and InfGCN). We do not apply any modification to the original code but use optimized configurations for inference to better utilize our GPU: for ChargE3Net, we process 20,000 probes in each batch instead of the default setting of 2,500 probes per batch, and for InfGCN, we process 40,000 probes in each batch with a batch size of 4.A significant advance in both accuracy and efficiency.The metrics for all methods are presented in Table 1 . We 3 .Ablation Analysis. We discuss the effectiveness of all in-gredients through an ablation analysis of the performance of different SCDP models. Starting from the most lightweight model with K = 4, L = 3, β = 2.0 and no virtual orbitals, we first observe that increasing the model size to K = 8, L = 6 significantly improves the performance, reducing the NMAE from 0.504% to 0.434%. Next, we increase the basis set size by adjusting β from 2.0 to 1.5, which further reduces the NMAE to 0.381%. The introduction of the virtual orbitals renders a significant gain in accuracy by reducing the error from 0.434% to 0.237% for β = 2.0 and from 0.381% to 0.196% for β = 1.5. In particular, the charge density near chemical bonds is significantly more accurate after introducing the virtual orbitals, as visualized in Figure 2 . On the other hand, for models with higher capacity, the improved accuracy comes at the cost of efficiency. As shown in Table 1 and Figure 3 , higher capacity consistently improves performance while sacrificing efficiency. At the same time, all SCDP models remain highly efficient compared to baseline models. When the scaling factors are introduced, accuracy further improves at a slight cost on efficiency for all models. As shown in Figure 3 , models with scaling factors from the Pareto front of all SCDP models benchmarked. 
discussion: 
 Charge density is a fundamental quantity for atomic systems and is central to DFT. ML methods for charge density prediction are promising as a means of greatly accelerating DFT by circumventing the iterative procedure used to find the ground-state charge density given a set of atomic coordinates. In this paper, we propose a recipe that combines three ingredients: (1) virtual nodes;(2) expressive basis sets; and (3) high-capacity equivariant networks that collectively outperform state-of-the-art methods in accuracy while being more than an order of magnitude faster. Nevertheless, there are still many directions for further improving the performance of our proposed model. First, the simple heuristic of assigning virtual node coordinates to bond centers may not be optimal. With recent advances in auto-regressive (Daigavane et al., 2023) and diffusion-based (Hoogeboom et al., 2022; Xu et al., 2022) equivariant generative models for 3D atomic structures, learning to insert the virtual orbitals may be a promising avenue for optimizing the placement of virtual nodes, thus improving charge density prediction. Due to the "nearsighted" nature of electronic matter (Prodan & Kohn, 2005) , an automated method for placing a higher density of virtual nodes near sites of chemical relevance may also be worthwhile to pursue. Second, we can use basis functions beyond Gaussian-type orbitals (e.g., Slater-type orbitals (Slater, 1930; Chong et al., 2004) or non-decay radial basis functions (Jing et al., 2024) ) that may require fewer functions to achieve the same level of accuracy.There are several limitations of the current paper that we aim to address in future work: (1) Despite substantial improvements in efficiency, the computational cost for training the current model is still significant: our best-performing model was pretrained for six days and fine-tuned for six days over four NVIDIA A100 GPUs for the QM9 charge density prediction task. The scaling factor fine-tuning stage requires a small learning rate, which prolongs training. The prediction model can benefit from resolving the training instability issues with the scaling factors as well as further improvement on the model architecture (Liao et al., 2023) .(2) While our approach achieves state-of-the-art performance on the QM9 charge density prediction benchmark, its effectiveness in crystalline materials (Jain et al., 2013; Shen et al., 2022 ) has yet to be validated. The GTOs and the equivariant network can be applied to materials without modification. The bondmidpoint-based virtual node assignment for molecules can be generalized to crystals through a crystal graph construction algorithm, such as CrystalNN (Zimmermann & Jain, 2020) . Alternatively, virtual nodes can be iteratively added to occupy void space inside the unit cell of the material using an algorithm based on the Voronoi diagram (Alexa et al., 2003) . The virtual nodes are expected to play an even more important role in prediction accuracy -this is because the diverse atomic species in materials and their complex interactions induce even more complex charge density patterns.(3) To better validate the practical utility of the predicted charge density, evaluation on the reduction of self-consistent field calculations, or on recovering physical observables, such as energy and forces (Jørgensen & Bhowmik, 2022; Sunshine et al., 2023; Koker et al., 2023) , will be highly valuable.ML methods for charge density prediction. Orbital-based methods predict coefficients for the orbital basis set functions to recover the target charge density. Past works have explored Gaussian processes (Fabrizio et al., 2019) and graph neural networks (Qiao et al., 2022; Rackers et al., 2023; Cheng & Peng, 2023; del Rio et al., 2023) in small molecules, water, and materials systems. (Focassio et al., 2023) used Jacobi-Legendre expansion -a many-body extension of atomic orbitals -for representing and predicting the charge density. These approaches, while efficient, suffer from lower accuracy in benchmarks such as QM9 (Cheng & Peng, 2023) or plane-wave basis sets (Kim & Ahn, 2024) . However, both methods still require neural message passing with a large number of probe nodes, which limits their scalability. In the present work, we combine virtual nodes, even-tempered Gaussian basis, and trainable basis functions to greatly improve the expressivity of orbital basis functions.Equivariant neural networks. Equivariant neural networks (Thomas et al., 2018; Anderson et al., 2019; Weiler et al., 2018; Gasteiger et al., 2021; Schütt et al., 2021; Geiger & Smidt, 2022; Liao & Smidt, 2022; Passaro & Zitnick, 2023) use equivariant representations and processing layers that can preserve roto-translational symmetries that are critical to atomistic modeling tasks. Equivariant models have shown advantages in ML potentials with respect to the accuracy, sample complexity, and molecular dynamics simulation capabilities (Unke et al., 2021; Batzner et al., 2022; Fu et al., 2023; Bihani et al., 2024) in addition to charge density prediction tasks (Jørgensen & Bhowmik, 2022; Koker et al., 2023) . This is because atomic forces and charge densities are indeed SE(3)-equivariant with regard to the input atomic coordinates. In this work, we leverage recent advances in methods for building more expressive and scalable equivariant architectures (Passaro & Zitnick, 2023) to improve the accuracy and scalability of charge density prediction.Our recipe for building ML charge density prediction capabilities involves two complementary aspects: the charge density representation and the prediction model.Gaussian-type orbitals (GTOs) are widely used as basis sets for representing electron configurations in quantum chemistry (Eichkorn et al., 1995) . They are spherical Gaussian functions centered at atomic coordinates. For an atom i at coordinate r i , a GTO basis function with exponent α, angular momentum quantum number (also called tensor order or degree) l, and magnetic quantum number m is given by the following expression:EQUATIONwhere r = ||r -r i || is the distance from a query coordinate r to the atom coordinate r i and Y l,m are real spherical harmonics. z α,l is a normalizing constant, such thatR 3 ||Φ|| 2 2 dV = 1.For the purpose of developing a machine learning model based on GTOs, we choose to represent the charge density, ρ, of an atomic system via a linear combination of many basis functions:ρ(r) = N i N i b j li,j m=-li,j c i,j,m Φ αi,j ,li,j ,m,ri (r), (3)where N is the number of atoms (including virtual ones when applicable), N i b is the number of l values (α values) for atom i. It should be noted that the charge density in Kohn-Sham DFT is not computed in this way; rather, Equation ( 3) is an artificial representation for the sake of training a machine learning model that is inspired by the orbital-like character of GTOs. The basis functions Φ αi,j ,li,j ,m,ri are chosen first as the basis set, with a fixed set of l and α values for each element. For example, the values for l and α for hydrogen in the def2-QZVPPD basis set (Weigend & Ahlrichs, 2005) are presented in Appendix C, Table 2 . The number of basis functions for an atom i can be derived asN i b j=1(2 • l i,j + 1), because m can be an integer from -l to l. A higher l value corresponds to a more complex angular part of the basis function and allows the corresponding spherical field to be more anisotropic. The number of orbital basis functions for elements H, C, N, O, and F of the def2-QZVPPD basis set (and its even-tempered variant, detailed later in this section) is included in Figure 4 . Atoms with more complex electronic structures are often represented with more basis functions. The number of basis 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 β, which controls the number of basis functions as described in Equation ( 4).functions, ls, and α values are carefully chosen in existing basis sets such as def2-QZVPPD. We refer interested readers to the original papers (Weigend & Ahlrichs, 2005; Schuchardt et al., 2007; Pritchard et al., 2019) for more details regarding the construction of atomic orbital basis sets.In training the machine learning model, after the basis set is determined, the coefficients c i,j,m are then fit such that Equation (3) best represent the charge density. GTOs have been studied in several previous works (Fabrizio et al., 2019; Rackers et al., 2023; Cheng & Peng, 2023) as a means of representing the charge density with promising results. However, their accuracy still bears significant room for improvement. We next introduce virtual orbitals, even-tempered Gaussian basis, and scaling factors for orbital exponents that greatly improve the expressive power of GTOs for charge density representation.Virtual orbitals. The atom-centered spherical fields often struggle to capture non-local electronic structures, which induces representation errors. This limitation is effectively addressed with the introduction of virtual orbitals, which define sets of spherical fields located in a position other than the atomic centers. Due to the critical importance of chemical bonds in defining the overall electronic structure, we insert virtual nodes into the midpoint of all chemical bonds for a given molecule (illustrated in Figure 4 (a) ). With this method, the coordinates to insert the virtual nodes are SE(3)-equivariant with regard to the input atom coordinates. Therefore, as long as the prediction of the basis set coefficients is SE(3)-equivariant, the overall charge density prediction will still be equivariant after the introduction of the virtual orbitals. We discuss potential extensions to virtual orbital assignments in Section 3. After the virtual nodes are created, one must decide which basis functions to use for the virtual orbitals. In this work, we use the basis functions of element O for the virtual nodes, which offers a balance in accuracy and efficiency based on preliminary experiments.Even-tempered Gaussian basis. The number of basis func-tions in existing basis sets, such as def2-QZVPPD, may be insufficient for representing complex charge densities. At the same time, expanding the number of basis functions requires care in choosing the values of l and α that improve expressivity effectively. As an example, the def2-QZVPPD basis set for hydrogen already contains basis functions with l = 1 and α = 2.292. Extending this basis set with basis functions with l = 1 and α = 2.0 will not significantly improve its expressivity because the spherical pattern will be similar to existing basis functions. A general methodology for controlling the basis set size is to use an even-tempered Gaussian basis set (Bardo & Ruedenberg, 1974) . Based on a reference atomic orbital basis set (e.g., def2-QZVPPD), the even-tempered basis set constructs a series of GTOs with a set of angular momentum quantum numbers l determined by the atomic number and exponents α given by:EQUATIONFor each spherical harmonics degree l, α and N l are chosen such that the exponents in the reference atomic orbital basis set are well-covered1 . β controls the number of basis functions -a smaller β creates a more expressive basis set with denser exponents. The use of an even-tempered basis set allows us to smoothly control the number of basis functions N i b effectively. Figure 4 (b) shows how the number of orbital basis functions for elements H, C, N, O, and F grows with a smaller β for the even-tempered Gaussian basis derived from the def2-QZVPPD basis set.Scaling factors for orbital exponents. In existing orbitalbased models (Fabrizio et al., 2019; Qiao et al., 2022; Rackers et al., 2023; Cheng & Peng, 2023; del Rio et al., 2023) , while the coefficients for the basis functions are predicted by the ML model, the exponents are fixed for each atom type and not trainable. However, atoms in different local atomic environments can exhibit significantly different charge den-sity patterns around them, especially for the virtual orbitals that aim at capturing interatomic interactions. To further improve the expressivity of the basis set, we make the exponents trainable by learning a positive scaling factor s > 0, such that Equation (2) becomes:Φ α,l,m,ri (r, s) = z α,l,s exp(-s • αr 2 )r l Y l,mrr i r .(5) The charge density is now represented with coefficients c i,j,m and scaling factors s i,j as:ρ(r) = N i N i b j li,jm=-li,j c i,j,m Φ αi,j ,li,j ,m,ri (r, s i,j ), (6) where z α,l,s is a normalizing constant such thatR 3 ||Φ|| 2 2 dV = 1.The introduction of the learnable scaling factors for the exponents significantly improves the expressive power of our charge density representation but is also prone to instability during training. We resolve the instability issue with a fine-tuning approach detailed in Appendix B.2.Using the atomic orbital basis set representation of charge density, the prediction model aims to predict the basis set coefficients c i,j,m and the scaling factors s i,j for each real and virtual node such that the predicted charge density matches the ground truth density obtained from DFT calculations. The model F takes as input the types A = {a i |i = 1, . . . , N } and coordinates R = {r i |i = 1, . . . , N } of all real and virtual nodes: {c i,j,m , s i,j |i = 1, . . . , N ; j = 1, . . . , N i b ; m = -l i,j , . . . , l i,j } = F (A, R).(7)Backbone architecture. Our construction of the ML prediction model is motivated by the following:• Charge density is SE(3)-equivariant with regard to the input atom coordinates. An equivariant model that can preserve this symmetry is desired. Concretely, the basis set coefficients c i,j,m are SE(3)-equivariant with regard to the input atom coordinates. The scaling factors s are SE(3)-invariant with regard to the input atom coordinates.• Charge density is data-rich and very sensitive to the local atomic environment. A high-capacity and expressive model is desired.• Efficiency is key for general applications of charge density prediction. The model should be efficient while being expressive.Based on these criteria, we consider equivariant model architectures and the balance between capacity and efficiency.For equivariant models, an important aspect of model expressivity is the representation of node and edge features in the form of irreducible representations (irreps) of SO(3)2 : spherical harmonic coefficients. A higher degree of representation (L) is desired for building high-capacity models. Previous works have employed the PaiNN architecture (Schütt et al., 2021; Jørgensen & Bhowmik, 2022) that is based on Cartesian features (equivalent to L = 1) or architectures based on irreps of SO(3) and tensor products (Qiao et al., 2022; Koker et al., 2023; Cheng & Peng, 2023; Geiger & Smidt, 2022) for charge density prediction. However, these models suffer from limited expressivity or scalability. Cartesian features are limited in representing angular information (L = 1); meanwhile, the O(L 6 ) complexity of tensor products limits the degree of representation that can be used while remaining computationally feasible.In this work, we adopt the equivariant spherical channel network (eSCN) architecture (Passaro & Zitnick, 2023) as our model backbone. While using SE(3)-equivariant representations and processing layers, the convolution layers in eSCN reduce the SO(3) convolutions (Zitnick et al., 2022) or tensor products (Thomas et al., 2018; Batzner et al., 2022) to convolutions in SO(2) that are mathematically equivalent.It reduces the complexity of the convolution operation from O(L 6 ) to O(L 3 ). Further, the use of point-wise, spherical non-linearity in eSCN also distinguishes itself from e3nnbased equivariant models that only apply non-linearity to the scalar features in the irreps. In our experiments, we also find that eSCN outperforms alternative architectures, such as tensor field networks (Koker et al., 2023; Thomas et al., 2018) and MACE (Batatia et al., 2022) . Using eSCN as the backbone architecture, we get the last-layer latent features x i for all real/virtual nodes:EQUATIONPrediction layers. The features x i are encoded using multichannel spherical harmonic coefficients (irreps). Note that the prediction target, basis set coefficients c i,j,m , are also encoded as multi-channel spherical harmonic coefficients. For example, for an eSCN with L = 3 and a latent dimension of 128, the last-layer latent atom features will be 128x0e + 128x1o + 128x2e + 128x3o. For the (uncontracted) def2-QZVPPD basis set of hydrogen described in Table 2, its irreps are 7x0e + 4x1e + 2x2e + 1x3e (even parity as charge density is reflection-invariant). The scaling factors s i,j are SE(3)-invariant and can be seen as scalar features of multi-channel irreps (14x0e for the def2-QZVPPD basis set of hydrogen). Therefore, we can make equivariant predictions of the basis set coefficients and invariant predictions of the scaling factors for each atom i through a fully connected tensor product layer over the atom features and additional processing:EQUATIONThe basis set coefficients are directly obtained through the fully connected tensor product. The tensor product also produces scalar features h i (128x0e for a 128-channel eSCN), which are used for predicting the scaling factors.The parameterization of Equation ( 10) allows the prediction to range from (C 1 , C 1 + C 3 ), and the scaling factors will be C 1 /C 2 + C 3 when the linear network in Equation ( 10) is zero-initialized. By setting C 1 = 1.5, C 2 = 2, and C 3 = 0.5, we can limit the range of the scaling factors to be (0.5, 2) (at most halve or double an exponent) and let initial scaling factors be 1 with a zero initialization of the linear layer in Equation ( 10).With the predicted coefficients and scaling factors, the charge density prediction ρ can be obtained efficiently by evaluating Equation ( 6) at all grid coordinates of the charge density voxel. We train the model end-to-end with a meanabsolute error loss L over the charge density:EQUATIONFine-tuning for scaling factor prediction. The scaling factors at the exponents lead to significant training instability when the network is trained from scratch. Therefore, we use a fine-tuning approach, where we first pre-train the model with fixed basis set exponents (an even-tempered Gaussian basis derived from def2-QZVPPD) and then finetune the prediction model with a small learning rate with the learning for scaling factors enabled. To achieve this, we zero-initialize the linear layer in Equation ( 10) and freeze its weights until the fine-tuning stage.Software. Basis-set-exchange-v0.9.1 (Schuchardt et al., 2007; Pritchard et al., 2019) , PySCF-v2.5.0 (Sun et al., 2020) , are used to build the orbital basis sets. PyTorch-v1.13.1 (Paszke et al., 2019) and CUDA-v11.6 (Nickolls et al., 2008) are used to build the SCDP models. The eSCN (Passaro & Zitnick, 2023) implementation is adopted from the Open Catalyst Project (Chanussot et al., 2021) . We also acknowledge Numpy (Harris et al., 2020) , ASE (Larsen et al., 2017) , Pymatgen (Ong et al., 2013) , wandb (Biewald, 2020) , Matplotlib (Hunter, 2007) , and Plotly (Inc., 2015) .Table 3 . Hyperparameters for SCDP models. 1 The cutoff distance used for building the message passing graph. 2 The cutoff distance for computing the charge density using Equation ( 6). An orbital basis function only influences all grid coordinates within this distance. Max # grid points in a forward pass for Equation (6) eSCN, K = 4, L = 3, β = 2.0 2, 000, 000 eSCN, K = 8, L = 6, β = 2.0 1, 000, 000 eSCN, K = 8, L = 6, β = 1.5 1, 000, 000 eSCN + VO, K = 8, L = 6, β = 2.0 600, 000 eSCN + VO, K = 8, L = 6, β = 1.5 400, 000 eSCN + VO, K = 8, L = 6, β = 1.3 400, 000We adopt the implementation of PySCF(Sun et al.,2020) and refer interested readers to the original paper/code for more details on the construction of even-tempered Gaussian basis.We refer interested readers to(Geiger & Smidt, 2022) and(Duval et al., 2023) for more information on equivariant geometric neural networks. 
 
 Appendix:  


Decision: Despite the range of numerical scores on this paper, the text of the reviews is essentially unanimous, and in discussion, all reviewers converged on poster presentation for this paper.

Reviewers have suggested several areas for improvement and the rebuttal includes additional experimentation which reinforces the paper's contribution within the domain of charge density prediction.

A key question of reviewers and AC in deciding between poster and spotlight presentation was generalizability of the paper's contributions to a wider machine learning audience.  As the authors point out in the rebuttal, the model has some wider application ("to any geometric deep learning backbone models that use irreps (irreducible representations of SO(3)) for feature representations, which is common in modern equivariant architectures").  This does include applications outside charge density prediction, but remains a relatively small subset of wider ML research.

Reviews: review: paper_summary: paper_summary: This application-oriented paper uses equivariant GNNs to predict charge density, represented orbital basis sets.
strength_weakness: strength_weakness: - I feel like the architectures are a simple variant of popular backbones. But this shouldn't affect the novelty of this application paper.
- the computation cost is still considerable
questions: questions: You mentioned in the text that if you use a different backbone other than eSCN, the performance is worse. It would be great if you could share some numbers to back up this statement.
limitations: limitations: The limitations are sufficiently discussed in the paper.
paper_summary: paper_summary: Overall this is a nice technical contribution towards the goal of machine-learning based prediction of charge densities trained from DFT calculations. The main merits are that the authors combine the recent eSCN equivariant network, which reduces the computational complexity of equivariant message passing compared to SO(3) message passing, and other existing ideas such as virtual atom/sites, and learnable (through the exponent rescaling) GTOs. The paper is well written and easy to follow. I have no doubt this is a valuable contribution. But my concern is lack of original ideas (all ideas were existing or straightforward) and lack of high impact application. To me you need to have at least one for acceptance.
strength_weakness: strength_weakness: Main problem is that the ideas are mainly straight application of existing algorithms and perhaps codes. It is not trivial to put all these pieces together but still I do not see much novelty in ML. The main advantage in application is efficiency compared to e.g. chargeE3Net as the accuracy is very similar. First, the issue of efficiency matters most when one has a killer application for this method. There can be but it is not demonstrated or even explained clearly. ML AIMD might be but we already have ML potentials for that. Second, the efficiency test of ChargeE3Net needs to be explained more clearly. Are we comparing both methods with the same batch size, molecular size, hyperparameters, etc?
questions: questions: Page 5. "where z_alph... =1" should be in the preceding sentence.
Page 6, equation 9. Are the hidden scalar features h_i the direct scalar output of eSCN? Or the output of FCTensorProd.?
Page 6, lines 230. If I am not mistaken, should be "c3, c1+c3" and "c1/(1+c2)+c3"
limitations: limitations: Yes.
paper_summary: paper_summary: This paper proposes a new effective approach to estimating the charge density of molecular systems using machine learning models. Although this topic has been recently actively studied, the present approaches suffer from either a lack of accuracy or scalability. The proposed approach alleviated the problem by identifying three key ingredients: (1) utilizing atomic and virtual orbitals, (2) utilizing expressive and learnable orbital basis sets, and (3) utilizing high-capacity equivariant neural network architecture. The proposed approach achieved state-of-the-art accuracy on QM9 while reducing computational cost.
strength_weakness: strength_weakness: The following (at least major) weaknesses should be approached before acceptance: 

Major weakness (the upper, the higher priority): 

1. The paper only performed the experiments on QM9 and I wonder if the method is over-fitted to the QM9 dataset. The previous methods (InfGCN, ChargE3Net, GPWNO) performed their experiments on average 3-dataset, e.g., among NMC, Cubic, MD, MP, and QM9. Since this method's scope is now limited to molecular systems, either MD or a molecule dataset used in [1] can be utilized. Because of the time limitation of the rebuttal period, at least logical validation of the effectiveness of the other dataset would be necessary (Of course the additional experiments are highly recommended, at least before publication).

2. Although not so crucial, the code is not provided in supplemental material. So, at this stage of this review, I do not have confidence that the paper's result can be reproducible. And I wonder why the authors did not submit it even though mentioning they would provide it after acceptance. Are there any proper reasons to hide the code from the reviewers?

3. Even though the author mentions at line 214-215 that eSCN outperforms tensor field networks and MACE, I cannot find the corresponding part in Section 4. An explicit indication of the experiment would be desirable. 

[1] Jørgensen, P.B., Bhowmik, A. Equivariant graph neural networks for fast  electron density estimation of molecules, liquids, and solids. npj Comput Mater 8, 183 (2022). https://doi.org/10.1038/s41524-022-00863-y

Minor (for further improvement): 

1. The explanation in "Prediction layers" can be improved. It would be strongly recommended to provide more information in the Appendix using more space.  

2. From the description of "Prediction layers", this method is a little difficult to apply to other models than eSCN. If not, it would be helpful for the readers that the authors provide additional "recipe" for the other models. 

3. The inference time scaling in terms of the molecular size (atom number) is missing, which would be important to apply for larger molecular systems, for example, protein, on which DFT has difficulty performing the calculation.
questions: questions: 1. The idea of the scaling factors for orbital exponents in Equation (4) seems to have a relation with the D4FT [2] Equations (20) and (21). Is the idea independently developed? It would be better to cite the paper and explain the difference. 

2. At line 102, roto-translational --> rotation-translational     (the former seems not strictly accepted term by all the community). 

3. In Table 1, the "Mol. per min." shows around 10% increases in SCDP when adding scaling factors. Why? Equation (4) indicates that the scaling factor would not increase the model complexity and size so much. 

[2]Li, Tianbo, et al. "D4FT: A deep learning approach to Kohn-Sham density functional theory." arXiv preprint arXiv:2303.00399 (2023).
limitations: limitations: Yes. It is provided in the "Discussion" section.
paper_summary: paper_summary: The paper proposes a new method for calculating charge density using orbital based representations along with an equivariant machine learning architecture. The paper starts by introducing density functional theory (DFT) and how charge density plays a central role for DFT. The important role of charge density in DFT motivates the need for a fast, accurate and robust method of calculating it creating an opportunity to apply machine learning to the problem. The introduction also describes how charge density is a volumetric object, which can it difficult to apply ML methods due to data density and how prior methods have aimed to address the overall challenge. After outlining the primary contributions of the paper, Section 2 provides greater detail for related work in ML methods for charge density prediction and equivariant neural networks.

Section 3 describes the primary method, which focuses on charge density representation and the prediction model. The charge density representation includes Gaussian-type orbitals approximated as a linear combination of spherical harmonics basis functions, as well as virtual orbitals to capture non-local electronic structures and interactions. The charge density approximation also includes trainable orbital exponents for each atom type making the representation more expressive. For the prediction model, the paper motivates the need for an equivariant and expressive model architecture and describes the eSCN-based architecture used for the experiments. The last paragraph of the section outlines a pretraining based approach to enhance training stability for the charge density prediction model.

Section 4 describes the experimental results performed on the charge density dataset for the QM9 database of molecules. The results include a comparison to other methods, as well as an ablation analysis related to the different SCPD models' trade-offs for accuracy and computational efficiency. The paper subsequently ends with a discussion along with a description of limitations that can be used for future research work.
strength_weakness: strength_weakness: The paper could be improved by providing more detail along the following:
* Provide more details on how charge density prediction is different from ML potentials and why that distinction matters. Concretely, what advantages and disadvantages does charge density prediction have over ML potentials and vice versa? [Significance, Clarity]
* Provide greater clarity on how their proposed method for Charge Density Representation differs from prior approaches, such as a summary table. [Clarity, Significance]
* Address the questions below.

-----------

Assessment updated during discussion period.
questions: questions: * Were there other options you considered for the def2-QZVPPD basis set? Can you describe your considerations and/or if it's reasonable to perform an ablation related to this? You somewhat discuss this in "Even-tempered Gaussian basis." but its unclear to me if an empirical evaluation is feasible.
* Can you include traditional DFT throughput numbers in the table for reference? This is to get a sense of the speedup.
* Can you describe in more detail why an equivariant network is required? Have you tried evaluating non-equivariant networks?
limitations: limitations: The paper includes a discussion on limitations in Section 5.

Comments: title: Thank you for your response
reply
Thank you for your response
Thank you for your response
reply
comment: Thank you for your response and additional suggestions! 

> Are more datasets needed beyond QM9 and MP? If so, what should researchers focus on? What types of other open problems exist that might be interesting for the machine learning community?

These are great questions. More datasets are certainly desirable -- first, the QM9 and MP datasets are not calculated with a very high level of theory -- which is where ML is going to help the most. An interesting future direction would be learning the difference between a higher and a lower level of theory, so the lower level of theory can be used as a prior, which may make the learning task easier so it requires less data/compute to train. It would also be interesting to consider other classes of systems -- QM9 is for small organic molecules, MP is for inorganic crystals -- other systems such as catalytic surfaces and metal organic frameworks could also be considered. 

From the few directions proposed in our previous rebuttal response, we believe in mid to short term we will be able to create efficient and accurate ML-based charge density prediction capabilities that can accelerate DFT through accurate SCF initialization, and obtain a variety of properties such as partial charges and dipole moments. We are also eager to see future research in pretraining/co-training of ML potentials or property predictors with charge density data. Addressing the limitations discussed in the paper to further improve charge density prediction is also exciting to us. In long term, we believe charge density modeling is crucial for high-fidelity modeling of more complex scenarios such as excited-state DFT and chemical reaction, and potential impact in designing better DFT functionals.

We truly appreciate your time, consideration, and feedback that greatly help us improve our paper.
Thank you for the additional details. I think they make the paper stronger and have adjusted my score accordingly. I have one additional suggestions:

* For future work or related work, it could be helpful to understand what would be required to cover a greater extend of relevant charge density prediction cases. Are more datasets needed beyond QM9 and MP? If so, what should researchers focus on? What types of other open problems exist that might be interesting for the machine learning community?
I appreciate the authors' response. I fully agree with the authors that the ability to predict charge density accurately opens the door to many downstream applications, and never doubted that. My main concern was and still is whether this specific contribution can be demonstrated to be generalizable and useful in applications. I'll keep my rating. Solid technical paper nevertheless.
Thank you for your fruitful information. 
I look forward to the camera-ready version and re-evaluate my score.
Thank you for your response! A key distinction between eSCN and alternative backbones is eSCN uses point-wise spherical non-linearities. While all models use irreducible representations of SO(3) (irreps) for the internal representation of atomic features, to ensure equivariance, tensor field network and MACE can only apply non-linearities over the scalar features of the irreps. That is, only a small portion of the irreps features can be processed with non-linearities. 

On the other hand, the eSCN convolution maps the irreps to an equivalent spherical function and then apply point-wise spherical non-linearities, then maps the spherical functions back to irreps when the convolution finishes. This allows non-linearities over higher order tensor features while preserving equivariance. We hypothesize non-linearities over higher-order features makes eSCN a powerful architecture for the charge density prediction task. We believe our task is very challenging, and requires very expressive networks because the rich charge density information needs to be compressed into a rather compact set of basis set coefficients -- a highly complex mapping. We believe this task requires more expressive model than ML potentials (which only needs to learn a scalar energy and vector forces on the atoms), or probe-based charge density prediction (which doesn't have to compress the charge density to highly compact representations).

Thank you again for your time and considerations. We are happy to discuss further if you have additional feedback.
Do you have some qualitative, rough idea on why switching form eSCN to another backbone would cause such a dramatic drop in the performance?

Other than that, all my questions are sufficiently addressed. Thank you again for your rebuttal.
Thank you for your response! We have completed the experiments for the two other molecules (ethane and malonaldehyde):

| Molecule      | SCDP (Ours) | GPWNO | InfGCN | CNN  | DeepDFT | DeepDFT2 | EGNN  | DimeNet | DimeNet++ | GNO   | FNO   | LNO   |
|---------------|-------------|--------|---|------|---------|----------|-------|---------|-----------|-------|-------|-------|
| ethanol       | **2.40 ± 0.26**         | 4.00 | 8.43   | 13.97| 7.34    | 8.83     | 13.90 | 13.99   | 14.24     | 82.35 | 31.98 | 43.17 |
| benzene       | **1.15 ± 0.06**       | 2.45 | 5.11   | 11.98| 6.61    | 5.49     | 13.49 | 14.48   | 14.34     | 82.46 | 20.05 | 38.82 |
| phenol        | **1.32 ± 0.07**        | 2.68 | 5.51   | 11.52| 9.09    | 7.00     | 13.59 | 12.93   | 12.99     | 66.69 | 42.98 | 60.70 |
| resorcinol    | **1.38 ± 0.08**        | 2.73 | 5.95   | 11.07| 8.18    | 6.95     | 12.61 | 12.04   | 12.01     | 58.75 | 26.06 | 35.07 |
| ethane        | **2.10 ± 0.13**        | 3.67 | 7.01   | 14.72| 8.31    | 6.36     | 15.17 | 13.11   | 12.95     | 71.12 | 26.31 | 77.14 |
| malonaldehyde | **2.77 ± 0.63**      | 5.32 | 10.34  | 18.52| 9.31    | 10.68    | 12.37 | 18.71   | 16.79     | 84.52 | 34.58 | 47.22 |

We find our method ( 4-layer eSCN, $L_{\mathrm{max}} = 3$, $\beta = 1.5$, with virtual nodes) significantly outperforms all baseline models on all six molecules in the MD dataset. Next, we plan to experiment with scaling factor finetuning for the MD molecules. We are also happy to further extend our experiments to a materials dataset such as NMC or Cubic in the final version. 

Thank you again for your time and considerations. We are happy to discuss further if you have additional feedback.
I really appreciate the authors' effort for the rebuttal. I am basically satisfied with the author's responses. 

Concerning W1, it is helpful if the authors can tell me which dataset the authors are going to try further for the camera-ready version, in addition to QM9 and MD. It would help me to decide the reconsidered evaluation score.