A Recipe for Charge Density Prediction
Authors: Xiang Fu
Andrew Scott Rosen
Kyle Bystrom
Rui Wang
Albert Musaelian
Boris Kozinsky
Tess Smidt
Tommi Jaakkola
Venue: NeurIPS 2024 poster (2024)
Keywords: AI for science
density functional theory
charge density
equivariant
molecule
Primary area: machine_learning_for_physical_sciences
TL;DR: A method that predicts charge density from atomic coordinates that outperforms state-of-the-art while being more than an order of magnitude faster.

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.

A Recipe for Charge Density Prediction
Xiang Ful
Andrew Rosen 2.3
Kyle Bystrom"
Rui Wang'
Albert Musaelian"
Boris Kozinsky4.
Tess Smidt
Tommi Jaakkola*
Massachusetts Institute of Technology
UC Berkeley
Lawrence Berkeley National Laboratory
Harvard John A. Paulson School Of Engineering and Applied Sciences
'Robert Bosch Research and Technology Center
Abstract
In density functional theory, the charge density is the core attribute of atomic systems from which all chemical properties can be derived. Machine learning methods are promising as means of significantly accelerating charge density predictions, yet existing approaches either lack accuracy or scalability. We propose a recipe that can achieve both. particular, we identify three key ingredients: (1) representing the charge density with atomic and virtual orbitals (spherical fields centered atom/virtual coordinates); (2) using expressive and learnable orbital basis sets (basis functions 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 order of magnitude faster than existing methods. Furthermore, our method enables flexible efficiency-accuracy trade-offs by adjusting the model and/or basis set sizes.
Introduction
Density functional theory (DFT) is computational quantum chemistry method that has enabled countless advancements in the chemical sciences by providing a tractable means calculate the electronic structure of molecules and materials [1]. The central concept in DFT is the charge density, fundamental quantity from which all derivable ground-state physicochemical properties of system, such as energy and forces, can, in principle, be derived. The most widely used Kohn-So Shamm formalism [2] 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 (NN)) where N₂ 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 reliant on an iterative calculation to identify the charge density that minimizes the potential energy functional for given atomic configuration. This process, known converging the self-consistent field, is the main computational expense within DFT. With a machine learning (MI.) 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. accurate enough, a machine-learmed charge density could also be used to directly predict electronic structure properties, such as the band gap, band structure, and electronic density of states material. Furthermore, the charge density itself can provide an enormous amount of insight into a molecule material. From the charge density, partial
"Conrespondence to Xiang Fu (xiangfu@csail mattaailll........mmmmmmmm..ee edu).
edd)..
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(5)
(b)
High
Low
Figure I: (a) Illustration of the orbital-based method for charge density representation for an example molecule (indole, C₂H,N), The overall charge density is represented as sum over spherical- harmonics-bas based atomic orbital basis functions (spherical fields) centered at each atom. (b) Left: Illustration of the probe- based method for charge density representation. The charge density is represented as a voxel where each grid point (probe node) represents a scalar density at that coordinate. The voxel for the example molecule is of size 108 x 96 40. Grid points with very small charge densities (< 0.05) are not visualized. Right: For mooee-based machine learning prediction model, the voxel contains too many grid points to be processed simultaneously. Sampling of the voxel points is needed during training and inference. All charge densities use the same colormap scale at the right-most side of the figure. Atom color code: H (white), C (gray), N (blue). The charge density is from the QM9 charge density dataset [7].
atomic charges, dipole moments, atomic spin densities, and effective bond orders can all be directly computed through one of several population analysis methods [3, 4]. For some materials discovery tasks, the charge density can also be a crucial descriptor depending on the application area [5, 6]. 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 rep- resented as voxels with a grid resolution of around 0.1 Á [7, 8J. 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 a representation that too coarse can have a substantial impact on energy and other derivable properties. This need for both efficiency and accuracy creates a significant challenge for ML methods.
The existing literature has mainly focused on two approaches to learning to predict 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 [9, 10, 11, 12, 13]. 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 model complex electronic structures between atoms.
The second approach (probe-based) [14, 7, 15, 16], 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 [7].
This paper aims to address this tations and ML models for charge density prediction. We identify three key ingredients:
dilemma with a new recipe for building represen-
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 domain-informed and expressive basis sets. In particular, we construct 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 securacy-efficieee We make the basis set exponents learnable to further improve expressivity.
trade-offs.
3. We use high-capacity equivariant neural network architecture (eSCN [17]), which enables efficient training and inference with features of high tensor order for large dataset.
We apply our recipe the widely used QM9 charge density benchmark [18, 19, 7J. Our method outperforms existing state-of-the-art methods while being around 30x faster. Furthermore, we can flexibly trade off accuracy and efficiency by adjusting the model/basis size; in doing so, we achieve up to 171x x efficiency compared to state-of-the-arr methods with only 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 justify the significance of each proposed ingredient.
2 Related Works
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 [9] and graph neural networks [10, 11, 12, 13] in small molecules, water, and materials 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 [18, 19, 7] and the Materials Project [21, 8] charge density datasets. trobe-based methods, on the other hand, predict the charge density by neural message passing between the atoms and probe nodes at all grid points. These methods [14, 7, 15, 16] have shown superior accuracy in both molecules and materials but suffer from poor scalability, as they require neural processing of millions of probe nodes for molecule/material structures of tens of atoms. Recent works also explored a combination of atomic orbitals and probe-based methods [12] or plane-wave basis sets [22]. However, both methods still require neural message passing with 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.
systems. [20] used Jacobi-Legendre expansion
Equivariant neural networks. Equivariant neural networks [23, 24, 25, 26, 27, 28, 29, 17] use equivariant representations and processing layers that can preserve rotational and 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 [30, 31, 32, 33] addition to charge density prediction tasks [7, 15]. 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 [17] to improve the accuracy and scalability of charge density prediction.
3 Methods
Our recipe for building ML charge density prediction capabilities involves two complementary aspects: the charge density representation and the prediction model.
3.1 Charge Density Representation
Gaussian-type orbitals (GTOs) are widely used as basis sets for representing electron cnnigurattions in quantum chemistry [34]. They are spherical Gaussian functions centered at atomic coordinates. For an atom at coordinate r., a GTO basis function with exponent a, angular momentum quantum number (also called tensor order or degree) I, and magnetic quantum number m is given by the
3

---

(8)
600
dett BER-RCHT
Bade sat stbe
500
B=2.0
400
pm333
300
200
100
N
Figure 2: (a) Two example molecules (left: indole, C&H,N; right: methanol CH;OH), before and after the bond-midpoint-basse H (white), C (gray), N (bluc), O (red), virtual nodes (small, black). (b) 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 B, which controls the number of basis functions as described in Equation (3).
virtual coordinates (small black points) are inserted. Atom color code:
following expression:
@ La.lmm (r) = R
I exp(-ar"))'y,
Pi
(I)
the distance from query coordinate r to the atom coordinate r, and Yi. are real spherical harmonics. A is normalizing constant, such that | dV For the purpose of developing a machine learning model based on GTOs, we choose to represent the charge density, p. of an atomic system via linear combination of many basis functions:
where 1-r-- r:
N
N
L.
p(r) =
C.j.m T Pa aa (r),
where N is the number of atoms (including virtual ones when applicable). N is the number of I values (a values) for atom i. It should be noted that the charge density in Kohn-Sham DFT is not computed this way; rather, Equation (2) 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 @ Pa are chosen first as the basis set, with a fixed set of 7 and a values for each element. For example, the values for | and a for hydrogen in the def2-QZVPH basis set [35] are presented in
Appendix A, Table 2. The number of basis functions for atom can be derived as Σ (2- 4., +1), because m can be an integer from -1 to 1. A higher I 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 this section) is included in Figure 2. Atoms with more complex electronic structures are often represented with more basis functions. The number of basis functions, Is, and a values are carefully chosen in existing basis sets such as def2-QZVPPD. We refer interested readers to the original papers [35, 36, 37] 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 Cij.m are then fit such that Equation (2) best represent the charge density. GTOs have been studied in several previous works [9, 11, 12] 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 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 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 defining the overall electronic structure, we insert virtual nodes into the midpoint of all chemical bonds for given molecule (illustrated Figure 2 (a)). With this method, the coordinates to insert the virtual nodes are SEuuieaiuiaaa with regard to the input atom coordinates. Therefore, as long the prediction of the basis set coefficients is -quiivaauuyaiii 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 5. After the virtual nodes are created, one must decide which basis functions use for the virtual orbitals. In this work, we use the basis functions of element O for the virtual nodes, which offers balance in accuracy and efficiency based on preliminary experiments.
Even-tempered Gaussian basis. The number of basis functions 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 I and a that improve expressivity effectively. As an example, the def2-QZVPPD basis set for hydrogen already contains basis functions with 1 1 and
2.292. Extending this basis set with basis functions with I = 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-tent tempered Gaussian basis set [38]. Based on a reference atomic orbital basis set (e.g., def2-QZVPPD), the even-tempered basis set constructs a series of GTOs with set of angular momentum quantum numbers determined by the atomic number and exponents a given by:
Or = a - 3 for 0, 0,,12...... N₁.
For each spherical harmonics degree 1, a and N₁ are chosen such that the exponents in the reference atomic orbital basis set are well-covered.. 8 controls the number of basis functions smaller creates a more expressive basis set with denser exponents. The use of an even-tempered basis set allows us smoothly control the number of basis functions N effectively. Figure 2 (b) shows how the number of orbital basis functions for elements H. C. N. O, and F grows with smaller 3 for the even-tente reenteempered Gaussian basis derived from the def2-QZVPPD basis set.
Scaling factors for orbital exponents. existing orbital-based models [9, 10, 11, 12, 13], while the coefficients for the basis functions are predicted by the MI. model, the exponents fixed for each atom type and not trainable. However, atoms in different local atomic environments can exhibit significantly different charge density 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 > O, such that Equation (I) becomes:
o Palmm (r, 8) 2a.1s exp( ar²))"
(T- (T-Ti (-777
r
where 2a,1, a normalizing constant such that 12 represented with coefficients Ci.jom and scaling factors si, as:
The charge density is now
N
Ng
ha
((r) =
Σ
Σ
Cijj. T Par (r, 8t.j),
(5)
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 fine-tuning approach detailed in Section 3.2.
3.2 Prediction Model
Using the atomic orbital basis set representation of charge density, the prediction model aims to predict the basis set coefficients Cjj.m and the scaling factors $i,, 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 l...., NN and coordinates R {rili = I.....N) of all real and virtual nodes:
{crgi.m Styli = 1,........... N:j #..j=,.... 1...., Ng. m m=
F(A,
(6)
Backbone architecture. Our construction of the ML prediction model is motivated by the following:
We adopt the implementation of PySCF [39] and refer interested readers to the original paper/code for more details on the construction of even-tempered Gaussian basis.
5

---

Charge density is model that can preserve this symmetry is desired. Concretely, the basis set coefficients C...o are SE(()-equivariant with regard to the input atom coordinates. The scaling factors 8ig are SF(3))invariant with regard to the input atom coordinates.
with regard to the input atom coordinates. An equivariant
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): spherical harmonic coefficients. A higher degree of representation (I.) is desired for building high- capacity models. Previous works have employed the PaiNN architecture [27, 7] that based on Cartesian features (equivalent to L 1) or architectures based on irreps of SO(3) and tensor products [10, 15, 12, 28] 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 (())) 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 [17] as our model backbone. While using SE(3)-equivariant representations and processing layers, the convolution layers eSCN reduce the SO(3) convolutions [41] or tensor products [23, 31] to convolutions in SO(2) that are mathematically equivalent. It reduces the complexity of the convolution operation from ( to O Further, the use of point-wise, spherical non-linearity in eSCN also distinguishes itself from e3nn-based 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 [15, 23] and MACE [42]. Using eSCN as the backbone architecture, we get the last-layer latent features x, for all real/virtual nodes:
(a. li = 1, N3 = eSCN R).
Prediction layers. The features x; are encoded using multi-channel spherical harmonic coefficients (irreps). Note that the prediction target, basis set coefficients C.j.m. are also encoded as multi-channel spherical harmonic coefficients. For example, for an eSCN with I, = 3 and a latent dimension of 128, the last-layer latent atom features will be 128x0e 128x10 128x2e + 128x3o. For the (uncontracted) - basis set of hydrogen described in Table 2, its irreps are 7x0e + 4x1e 2x2e 1x3e (even parity as charge density is rellection-invariant)) The scaling factors St,j are -Evariann and can be seen scalar features of multi-chann 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 through a fully connected tensor product layer over the atom features and additional processing:
{cidms helj = L, I,..., Nf. = -4,, -4....., 4.3 Fully halll/oonnneeeeeeerrrrttteeeeeee x.) (8) (st 1, ... N:) = C1 + exxp(-Lnneate
In C2)) + C3-
The basis set coefficients are directly obtained through the fully connected tensor product. The tensor product also produces scalar features h; (128x0e for 128-channel eSCN), which are used for predicting the scaling factors. The parameterization of Equation (9) allows the prediction to range from (C3, C1 + Cs), and the scaling factors will be C1/(1 + C2) + C3 when the linear network in Equation (9) is zero-ine irrroinitialized. By setting C₁ = 1.5, C2 = 2, and C3- = 0.5, we can limit the range of the scaling factors be (0.5, 2) (at most halve double exponent) and let initial scaling factors be 1 with a zero initialization of the linear layer in Equation (9).
With the predicted coefficients and scaling factors, the charge density prediction can be obtained efficiently by evaluating Equation (5) at all grid coordinates of the charge density voxel. We train the model end-to-end with mban-absolute error loss L over the charge density:
We refer interested readers to [28] and [40] for more information on equivariant geometric neural networks.

---

Table 1:
QM9 charge density prediction error and efficiency on the test set. Metrics for baseline models are from previous papers whenever possible and skipped (-) when unavailable. The metrics (L means lower the better, f means higher the better) of the best-performing model are bold. The metrics are reported with corresponding standard errors when available. For SCDP models, K' is the number of interaction layers in the eSCN backhone, T. the tensor order of the feature representation in the eSCN backbone, and 3 controls the expressiveness of the even-tempered Gaussian basis set. higher K, higher L, or lower / indicates a more expressive model. eSCN + VO indicates that virtual orbitals are used. NMAE stands for normalized mean absolute error. Efficiency is measured by molecule per minute (mol. per min.).
NMAE [%]
NMAF, Split [8]]
Mol. per min. [min 1
D-DceDDFT [7]
0.357 + 0.001
e-DeepDFT [7]
0.284 0.001
ChargE3Net [15]
0.196 0.001
0.203 0.003
3.95
InfGCN [12]
0.869 + 0.002
0.93
72.00
InGGCN, GTO only [12]
3.72
GPWNO [22)
0.73
SCDP models (Ours)
eSCN, 4,4 3,8 2,0
0.504 : 0.001
0.514 0.003
675.47
eSCN, =8, -0.0=0,-200
0.431 0.006
0.452 0.017
567.19
eSCN, K 5.8-6,8-1
6.8-1.5 1.5
0.381 + 0.001
0.391 4: 0.002
442.25
sSCN VO. K =88,166,88 2.0
0.237 0.001
0,250 0.002
231.21
eSCN =
0.206 0.001
0.220 0.002
177.14
eSCN + VO, K =8, 1.3
0.196 0.001
0.209 + 0.002
136.92
SCDP fine-tuned with scaling factors (Ours)
eSCN, K =, 1, 8.
0.432 + 0.001
0.438 + 0.003
644.00
eSCN, A =
6, 2.0
0.309 + 0.007
0,386 0.018
514.56
cSCN, A == -= -
0.346 0.001
0.354 0.002
419.57
eSCN VO. K =8.L- -..= 2.0
0,207 1:07:::00 0.001
0,220 0.002
221.19
eSCN VO, K- =8.L=6,9= 1.5
0.187 = 0.001
0.200 - 0.002
164.94
eSCN VO, K 8,L 6,3 1.3
0.178 1 0.001
0,191 + 0.002
125.29
L = EreData [Ip(r) -(()))
(10)
Fine-tuning for scaling factor prediction. The scaling factors the exponents lead to significant training instability when the network is trained from scratch. Therefore, we use 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 fine-tune the prediction model with small learning rate with the learning for scaling factors enabled. To achieve this, we aero-initialize the linear layer in Equation (9) and freeze its weights until the fine-tuning stage.
Experiments
Our experiments 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.
Dataset and metrics. The QM9 charge density dataset [18, 19, 7] 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 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 [7, 15], benchmark the prediction accuracy with the normalized mean absolute error, defined as:
NMAF()) = SEslP(r) - dV Sgolp(r) dv
(11)
7

---

Accuracy (1-NMAE [%]]
99.8
k=4,233,==20
K=8,1=6.f=22
99.7
8,
vo, 8,,
vo, 6,,
99.66
vo, K= -== 11L=6,8111.
with scaling factors
99.5
without scaling factors
100
200
300
400
500
600
700
Efficiency (Mol. per min. [min-1])
Figure 3:
trade-off for SCDP models. The models with scaling factor fine-tuning form the Pareto front.
where the integration is approximated by summing over the full charge density voxel. benchmark the efficiency of different methods by the number of molecules predicted per minute (Mol. per min.) on a single NVIDIA A100-80GB-PCC PCIe GPU for the QM9 test split.
In addition to the QM9 charge density dataset, we also benchmark our method on the MD charge density dataset [43, 44, 12] and the Cubic charge density dataset [45]. We use the same data splits as previous works on these benchmarks [12, 22]. Experimental results and comparisons to baselines are included Appendix
Baseline Models. We compare SCDP to several previous works on the QM9 charge density pre- diction benchmark [18, 19, 7]. i-DeepDFT, e-DeepDFT [7], and ChargE3Net [15] are probe-based methods with different backbone architectures: i-DeepDFT uses SchNet [27], e-DeepDFT uses PaiNN [27], while ChargE3Net uses higher-order equivariant features under the tensor field network framework [23, 28]. InfGCN [12] combines GTOs and a shallow network for probe-based inference. also has more efficient but less accurate GTO-only variant. GPWNO [22] combines GTOs and plane-wave basis sets but still requires large number (64,000) of probe nodes for constructing the plane wave prediction. NMAE results for F-FeepDFT, e-DeepDFT, and ChargE3Net are from [15]. NMAE results for InfGC and GPWNO are also from the original papers, which uses 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 E3Net and InfGCN). We 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 each batch instead of the default setting of 2,500 probes per batch, and for InfGCN, process 40,000 probes in each batch with batch size of 4.
significant advance in both accuracy and efficiency. The metrics for all methods are presented in Table 1. We have a series of SCDP models with different model sizes, basis set sizes, as well as options on the inclusion of virtual orbitals and scaling factors. Our best-performing model uses the virtual orbitals described in Section 3.1, an eSCN of 8 layers and feature representation of order L 6, an even-tempered Gaussian basis with B 1.3, and scaling factor fine-tuning. This model achieves an NMAE of 0.178 on the QM9 charge density test set, outperforming the state-of-the-as method Char [15]- a probe-based method. While being more accurate, our best model also significantly outperforms ChargF33Nt by 31.7x in efficiency. Other configurations of our model with smaller model sizes, basis set sizes, and models without virtual orbitals can trade off accuracy for further gains in efficiency. The trade-off curves are visualized in Figure 3. Compared to more efficient baseline model, InfGCN, all benchmarked configurations of our method are more efficient and significantly outperform in accuracy. These results convincingly demonstrate significant advance in the eecurrcc--fficieene trade-off in ML methods for charge density prediction. Figure in Appendix A shows the convergence of validation NMAE during pretraining and fine-tuning of SCDP models. More details on the hyperparameters for model construction and training are included in Appendix A, Table 3.
Alltion Analysis. We discuss the effectiveness ingredients through an ablation analysis of the performance of different SCDP models. Starting from the most lightweight model with

---

VO
VO+ Scaling
VO Scaling
Knd K=Ri=ape20
K=LL=kf=20
K=RL=G,p=20
0.05
0005
7.25
0.05
Figure 4: Visualization of the reference charge density and prediction errors for select SCDP models with two representative test molecules (top: C2HęNO, and bottom: C). The first column the ground truth charge density with the corresponding color scale. The next five columns are prediction errors from various models which all use the same color scale in the rightmost for error magnitude. The prediction errors significantly reduce with larger model size, virtual orbitals, orbital exponent scaling, and larger basis set. VO stands for virtual orbitals. Scaling stands for scaling factor fine-tuning. The virtual orbitals significantly reduce errors around chemical bonds. Atom color code: H (white), C (gray), N (blue), O (red), virtual nodes (small, black).
K = 4, = 2.0 and no virtual orbitals, we first observe that increasing the model size 8, significantly improves the performance, reducing the NMAE from 0.504% to 0.434%. Next, we increase the basis set size by adjusting 8 from 2.0 to 1.5, which further reduces the NMAE 0.381%. The introduction of the virtual orbitals renders significant gain in accuracy by reducing the error from 0.434% to 0.237% for 3 = 2.0 and from 0.381 1 % to 0. 196% for 3 = 1.5. In particular, the charge density near chemical bonds significantly more accurate after introducing the virtual orbitals, as visualized in Figure 4. On the other hand, for models with higher capacity, improved accuracy comes at the cost efficiency. As shown Table 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 slight cost on efficiency for all models. As shown in Figure 3, models with scaling factors from the Pareto front 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 means of greatly accelerating DFT by circumventing the iterative procedure used to find the ground-state charge density given set of atomic coordinates. In this paper, we propose recipe that combines three ingredients:
(1) virtual nodes; (2) expressive basis sets; and (3) high sigh-capacity equivariant networks that collectively outperform state-of 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 [46] and diffusion- [47, 48] 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 [49], an automated method for placing a higher density virtual nodes near sites of chemical relevance may also be worthwhile to pursue. Second, can basis functions beyond Gaussian-type orbitals (e.g., Slater-typs orbitals [50, 51] or non-decay radial basis functions [52]) that may require fewer functions 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- perso-rrroming 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 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 [53]. (2) While our approach achieves state-of-the-art performance on the QM9 charge density prediction benchmark, its effectiveness in crystalline materials [21, 8] has major room for improvement. The GTOs and the equívariant network can be applied to materials without modification. The bond-midpoint-bass to crystals through crystal graph construction algorithm, such as CrystalNN [54]. 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 [55]. The virtual nodes 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 energy and forces [7, 56, 15], will be highly valuable.
virtual node assignment for molecules can be generalized
Acknowledgments and Disclosure of Funding
We thank Teddy Koker, Chaoran Cheng, and Aria Mansouri Tehrani for their helpful discussions and insights. This work was supported by the GIST-MIT Research Collaboration grant funded by GIST and the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium. A.S.R. acknowledges support via a Miller Research Fellowship from the Miller Institute for Basic Research in Science, University of California, Berkeley.
References
[1] Anubhav Jain, Yongwoo Shin, and Kristin A Persson. Computational predictions of energy materials using density functional theory. Nature Reviews Materials, 1(1):1113, 2016.
[2] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1 110A11133 133, 1965.
[3] Célia Fonseca Guerra, Jan-Willem Handgraaf, Evert Jan Baerends, and F Matthias Bickelhaupt. Voronoi deformation density (VDD) charges: Assessment of the Mulliken, Bader, Hirsh- feld, Weinhold, and VDD methods for charge analysis. Journal of computational chemistry, 25(2):189-210, 2004.
[4] Nidia Gabaldon Limas and Thomas A Manz. Introducing DDEC6 atomic population analysis: part efficient parallel computation of net atomic charges, atomic spin moments, bond orders. and more. RSC advances, 8(5): 2018.
[5] Huam Shen, Haoming Howard Li, Ann Rutt, Matthew K Horton, and Kristin A Persson. Topological graph-based analysis of solid-state ion migration. npj Computational Materials, 9(1):99, 2023.
[6] Zhikun Yao, Yanzhen Zhao, Wenjun Zhang, and Lee Alan Burton. Assessing the design rules of electrides. Journal of Materials Chemistry C, 2024.
[7] Peter Bjøm Jørgensen and Arghya Bhowmik. Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids. npj Computational Materials, 8(1):183, 2022.
[8] Jimmy-Xuse Xuan Shen, Jason M Munro, Matthew K Horton, Patrick Huck, Shyam Dwaraknath, and Kristin A Persson. A representation iisieetttttooonnneeeeeeee electronic charge density database for crystalline materials. Scientific data, 9()):661, 2022.
[9] Alberto Fabrizio, Andrea Grisafi, Benjamin Meyer, Michele Ceriotti, and Clemence Cormin- boeuf. Flectron density learning of non-covalent systems. Chemical science, 10(41):9424 943244424 2019.
10

---

[10] Zhuoran Qiao, Anders S Christensen, Matthew Welborn, Frederick R Manby, Anima Anand- kumar, and Thomas F Miller III. Informing geometric deep learning with electronic inter- actions to accelerate quantum chemistry. Proceedings of the National Academy of Sciences, 119(31))e2202221119 2022.
[1]] Joshua Rackers, Lucas Tecot, Mario Geiger, and Tess Smidt. recipe for cracking the quantum scaling limit with machine learned electron densities. Machine Learning: Science and Technology, 4( (1))0055 2023.
[12] Chaoran Cheng and Jian Peng. Equivariant neural operator learning with graphon convolution. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
[13] Beatriz G del Rio, Brandon Phan, and Rampi Ramprasad. A deep learning framework to emulate density functional theory. upj Computational Materials, 9(1):188, 2023.
[14] Sheng Gong, Tian Xie, Taishan Zhu, Shuo Wang, Eric R Fadel, Yawei Li, and Jeffrey C Grossman. Predicting charge density distribution of materials using local-enviromment-bass graph convolutional network. Physical Review B, 100(18):
2019.
[15] Teddy Koker, Keegan Quigley, Eric Taw, Kevin Tibbetts, and Lin Li. Higher-order equivariant neural networks for charge density prediction in materials. arXiv preprint arXiv:2312.05388, 2023.
[16] Phillip Pope and David Jacobs. Towards combinatorial generalization for catalysts: A kohn- sham charge-density approach. Advances Neural Information Processing Systems, 36, 2023.
[17] Saro Passaro and C Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In International Conference on Machine Learning, pages 27420-27438 PMLR, 2023.
[18] Lars Ruddigkeit, Ruud Van Deursen, Lorenz Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database GDB-17. Journal of chemical information and modeling, 52(1 1) 1):2864 2275, 2012.
[19] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[20] Bruno Focassio, Michelangelo Domina, Urvesh Patil. Adalberto Fazzio, and Stefano San- vito. Linear jacohi-legendre expansion of the charge density for machine carning-acceleree electronic structure calculations. npj Computational Materials, 91):877 2023.
[21] Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, and Kristin Persson. Commentary: The Materials Project: A materials genome approach to accelerating materials innovation. APL materials, 1(1), 2013.
[22] Seongsu Kim and Sungsoo Ahn. Gaussian plane-wave wave neural operator for electron density estimation. arXiv preprint -xXx:22402.04275 2024.
[23] Nathaniel Thomas, Tess Smidt, Steven Keames, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equiivrrann neural networks for 3D point clouds. arXiv preprint arXiv::8022 /802.082/90 19, 2018
[24] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems, 32, 2019.
[25] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3D steerable CNNs: Learning rotationally equivariant features in volumetric data. Advances in Neural Information Processing Systems, 31, 2018.
[26] Johannes Gasteiger, Florian Becker, and Stephan Günneemann Gemnet: Universal directional graph neural networks for molecules. Advances Neural Information Processing Systems, 6867990-802,,0 2021.
11

---

[27] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.
[28] Mario Geiger and Tess Smidt. arX1v:200..94553, 2022.
e3nn: Euclidean neural networks.
arXiv preprint
[29] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3D atomistic graphs. arXiv preprint MXii:22006.19900 2022.
[30] Oliver T' Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schütt,, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemi- cal Reviews, 121(16)::0142-10155 2021.
[31] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neu- ral networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
[32] Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafael Gomez-Bombarelli and Tommi S. Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. Transactions Machine Learning Research, 2023. Survey Certification.
[33] Vaibhav Bihani, Sajid Mannan, Utkarsh Pratiush, Tao Du, Zhimin Chen, Santiago Miret, Matthieu Micoulaut, Morten M Smedskjaer, Sayan Ranu, and NM Anoop Krishnan. EGraFF- Bench: evaluation of equivariant graph neural network force fields for atomistic simulations. Digital Discovery, 3(4):759-768, 2024.
[34] Karin Eichkorn, Oliver Treutler, Holger Öhm, Marco Häser, and Reinhart Ahlrichs. Auxiliary basis sets to approximate coulomb potentials. Chemical physics letters, 240(4):283 1995.
[35] Florian Weigend and Reinhart Ahlrichs. Balanced basis sets of split valence, triple zeta valence and quadruple zeta valence quality for H to Rn: Design and assessment of accuracy. Phys. Chem. Chem. Phys., 7:3297, 2005.
[36] Karen Schuchardt, Brett Didier, Todd Elsethagen, Lisong Sun, Vidhya Gurumoorthi, Jared Chase, Jun Li, and Theresa L Windus. Basis set exchange: a community database for computational sciences. Journal of chemical information and modeling. 47()):0055-1022. 2007.
[37] Benjamin P Pritchard, Doaa Altarawy, Brett Didier, Tara D Gibson, and Theresa L Windus. New basis set exchange: An open, up-to-date resource for the molecular sciences community. Journal of chemical information and modeling, 59(1 11:4414444200, 2019.
[38] Richard D Bardo and Klaus Ruedenberg. Even-tempered atomic orbitals. vi. optimal orbital exponents and optimal contractions of gaussian primitives for hydrogen, carbon, and oxygen in molecules. The Journal of Chemical Physics, 60(3):918-93) 1974.
[39] Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao, Mare Barbry, Nick S Blunt, Nikolay A Bogdanov, George Booth, Jia Chen, Zhi-Hao Cui, et al. Recent developments the PySCF program package. The Journal of chemical physics, 153(2), 2020.
[40] Alexandre Duval, Simon V Mathis, Chaitanya K Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D Malliaros, Taco Cohen, Pietro Liò, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric GNNs for 3D atomic systems. arXiv preprint arXiv:2312.075111 2023.
[41] Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood. Spherical channels for modeling atomic interactions. Advances in Neural Information Processing Systems, 35:8054-800 8067, 2022.
12

---

[42] Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant messuge passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1422-114366 2022.
[43] Felix Brockherde, Leslie Vogt, Li, Mark E Tuckerman, Kieron Burke, and Klaus- Müller. Bypassing the kohn-sham equations with machine learning. Nature communications, 8(1):872, 2017.
[44] Mihail Bogojeski, Leslie Vogt-Maranto, Mark E Tuckerman, Klaus-Robert Müller, and Kieron Burke. Quantum chemical accuracy from density functional approximations via machine learning. Nature communications, 11(1):5223, 2020.
[45] Fancy Qian Wang, Kamal Choudhary, Yu Liu, Jianjun Hu, and Ming Hu. Large scale dataset of real space electronic charge density of cubic inorganic materials from density functional theory (dft) calculations. Scientific Data, 9():55, 2022.
[46] Ameya Daigavane, Song Kim, Mario Geiger, and Tess Smidt. Symphony: Symmetry-cquivarte point-centered spherical harmonics for molecule generation. arXiv preprint arXiv:2311.1619,, 2023.
[47] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3D. In International conference machine learning, pages 8867-8887. PMLR, 2022.
[48] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.0292. 2022.
[49] Emil Prodan and Walter Kohn. Nearsightedness of electronic matter. Proceedings of the National Academy of Sciences, 1035333: 2005.
[50] John C Slater. Atomic shielding constants. Physical review, 36(1):57, 1930.
[51] Delano P Chong, Erik Van Lenthe, Stan Van Gisbergen, and Evert Jan Baerends. Even-tempered slater-type orbitals revisited: From hydrogen to krypton. Journal of computational chemistry, 1030):003- 2004.
[52] Bowen Jing, Tommi S. Jaakkola, and Bonnie Berger. Equivariant scalar fields for molecular docking with fast fourier transforms. In The Twelfth International Conference on Learning Representations, 2024.
[53] Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt. EquiformerV2: Improved equivari- ant transformer for scaling to higher-degree representations. arXiv preprint arXiv:2306.12055 12059, 2023.
[54] Nils ER Zimmermann and Anubhav Jain. Local structure order parameters and site fingerprints for quantification of coordination environment and crystal structure similarity. RSC advances, 10(10):6063 1611166636668,,,,, 2020.
[55] Mare Alexa, Johannes Behr, Daniel Cohen-Or, Shachar Fleishman, David Levin, and Claudio Silva. Computing and rendering point set surfaces. IEEE Transactions on visualization and computer graphics, 91):3-15, 2003.
[56) Ethan M Sunshine, Muhammed Shuaibi, Zachary W Ulissi, and John R Kitchin. Chemical properties from graph neural network-pres prredcctee electron densities. The Journal of Physical Chemistry C, (27(4)):2455-2344 2023.
[57] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 44/2.9900 2014.
[58] Florian Weigend. Hartree-fock exchange fitting basis sets for h m. Journal of computational chemistry, 177 2008
13

---

[59] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga. al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[60] John Nickolls, lan Buck, Michael Garland, and Kevin Skadron. Scalable parallel programming with CUDA: Is CUDA the parallel programming model that application developers have been waiting for? Queue, 6(2):40-53 2008.
[61] Lowik Chanussot, Abhishek Das. Siddharth Goyal. Thibaut Lavril, Muhammed Shuaibi, Mor- gane Riviere, Kevin Tran, Javier Heras-Doming Caleb Ho, Weihua Hu, al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11110):059-60722 2021.
[62] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Vir- tanen, David Coumapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Hal- dane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-362 September 2020.
Kevin
[63] Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E Castelli, Rune Christensen, Marcin Dulak, Jesper Friis, Michael N Groves, Bjørk Hammer, Cory Hargus, al. The atomic simulation environment a Python library for working with atoms. Journal of Physics: Condensed Matter, 29(27):273002, 2017.
[64] Shyue Ping Ong, William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L Chevrier, Kristin Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): robust, open-source python library for materials analysis. Computational Materials Science, 68:314-319, 2013.
[65] Lukas Biewald. Experiment tracking with weights and biases, 2020. Software available from wandb.com.
[66] J. D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering, 9(3):90-95, 2007.
[67] Plotly Technologies Inc. Collaborative data science, 2015.
14

---

A
Appendix
Pretraining
Scaling fine -uninn
1.0
0.45
Validation NMAE (96
0.8
0.40
K- 1-,11-3,--21
0.35
K-8,1-6,3-20
0.6
K-
0.4
VO, 6.
0.25
VO, -0.
0.20
K0,K-
Training step
x10
Training step
Figure 5: Convergence of validation NMAE during pretraining and finetuning.
Table 2: The (uncontracted) def2-QZVPPD basis set for H.
ax
0
190.6916900
28. 6.653200
1.8412455 0
6. 6.5095943
0.59853725 0
0.21397624
0.0803 6286 0
2.29200000
0.83800000
0.29200000
0.084063199222
2.06200000
2
0.66200000
1.39700000
Algorithm Pseudo code for the charge density prediction procedure
1: Input: Atomic numbers A (N, 1), atom positions R (N, grid positions R. (M, 3)
2: Output: charge density at the grid positions C (M, 1)
3: obtain atomic irreps features: X [æsli li 1, ...., N} = eSCN(A, R)
4: obtain basis set coefficients {Ci.j.m h 1,..., N, Ng; - and
scaling factors {st = ...... N.j 1,..., N2 51/(1+xxp(-Linerr +In C2)) + C33
following Equation (8) and (9)
5: for R. do
A> in practice, we use batched inference
6: obtain p(r) following Equation (4) and (5)
7: end for
8: return C = (p(r)rr E Re}
Experimental results on MD and Cubic. We benchmark the proposed SCDP models on MD [43, 44, 12] and the Cubic charge density dataset [45] in Table 4. We find the SCDP models significantly outperform bascline models. Hyperparameters used for the MD and Cubic experiments summarized in Table 5 and Table 6. The MD models trained on 4 GPUs, while the Cubic model is trained on 8 GPUs. For the molecules the MD dataset, we use bond centers as coordinates for virtual nodes. For the materials in the Cubic dataset, we iteratively insert virtual nodes up to the number atoms in the unit cell using an algorithm based on Voronoi diagrams [55].
Software. Basis-set-exchae www.seeeexxhaaeeOO.... 1 [36, 37] and PySCF-v2.5.0 [39] are used to build the orbital basis sets. E3NN-v0.5.1 [28], PyTorch-v1.133 [59], and CUDA-v11, 66 [60] are used to build the SCDP models. The eSCN [17] implementation is adopted from the Open Catalyst Project [61]. We also acknowledge Numpy [62], ASE [63], Pymatgen [64], wandb [65], Matplotlib [66], and Plotly [67].
15

---

225
Mol. per min. [min
200
175
150
125
12
14
16
18
20
Number of atoms
Figure 6: Efficiency function of molecular size for our most expressive model (eSCN VO, K 8, I. = B = 1.3, with scaling factors). We measure the efficiency by running inference over 500 sampled molecules from the QM9 charge density dataset for given number of atoms.
B Broader Impact
This paper proposes ML method for accelerating charge density prediction, crucial task in computational chemistry. The adoption of our method is useful for scientific discovery and can yield positive negative repercussions, contingent the applications. The proposed method should be used for materials and drug discovery research that benefit our society.
16

---

Table 3: Hyperparameters for SCDP models the QM9 dataset. 'The cutoff distance used for building the message passing graph. "hh cutoff distance for computing the charge density using Equation (5). An orbital basis function only influences all grid coordinates within this distance.
Hyperparameter
Value
eSCN
interaction layers
[4, 8]
Lmax
[3. 6]
Mmax
2
sphere channels
128
hidden channels
256
edge channels # sphere samples radius cutoff!
128
128
6 A
Basis set
reference basis set
def2-QZVPPD [35]
3
[2.0, 1.5, 1.3]
orbital inference cutoff?
5 A
Training
batch size
4
# grid point samples (training, without VO)
100, 000
# grid point samples (validation/testing. without VO)
200, 000
# grid point samples (training, with VO)
60, 000
# grid point samples (validation/testing, with VO)
120, 000
precision
32
gradient clipping
0.5
# training steps (pretraining)
500, 000
# training steps (fine-tuning)
300, 000
optimizer
Adam [57]
Adam 81
0.9
Adam 83
0.999
Adam
1 x 10-8
weight decay
0
initial learning rate (pretraining)
0.001
initial learning rate (fine-tuning)
2 x 10-5
learning rate scheduler
exponential (LR initial LR 0.96step/C)
terminal learning rate (pretraining)
1 x 10-b
terminal learning rate (fine-tuning)
2 x 10 6
Inference
batch size (without VO)
8
batch size (with VO)
4
Max grid points in forward pass for Equation (5)
eSCN, K 4, 1. 3, 9 2.0
2, 000,,000
<SCN, K 8, L 6, 3 2.0
1, 1,000,000
eSCN, K 8, =6,3 1.5
1, 000,0000
eSCN VO, K L. =6,3 2.0
600, 000
SCN VO, K S, 6,9 1.5
400, 000
eCNN+ VO, K 8, 6, 1.3
400, 000
17

---

Table 4: Benchmark results (NMAE) the MD and Cubic datasets.
Molecule
SCDP (Ours)
GPWNO [22]
InfGCN [12]
MDD-ethanol
2.34 + 0.25
4.00
8.43
MDD-bnnene
1.13 0.06
2.45
5.11
MD MD-ppenol
1.29 + 0.07
2.68
5.51
MD-resorcinol
1.35 - 0.08
2.73
5.95
MD-ethane
2.05 0.12
3.67
7.01
MDD-malnaaldehyde
2.71 + 0.60
5.32
10.34
Cubic
2.59 0.25
7.69
8.98
Table Hyperparameters for SCDP models on the MD dataset. Parameters that are the same as the QM9 models are omitted in this table.
Hyperparameter
Value
eSCN
# interaction layers Lmax
Mmax
Basis set
reference basis set
def2-QZVPPD [35]
3
1.5
orbital inference cutoff
5 Å
Training
batch size
# grid point samples (Iraining, with VO)
125, 000
# grid point samples (validation/testingg, with VO)
125, 000
# training steps (pretraining)
250, 000
# training steps (fine-tuning)
50, 000
Inference
batch size (with VO)
Max grid points in a forward pass for Equation (5)
eSCN, K = 4, L = 3, B = 1.5
1, 000000 000
18

---

Table Hyperparameters for SCDP models on the Cubic dataset. Parameters that the same as the QM9 models are omitted in this table.
Hyperparameter
Value
eSCN
interaction layers
8
Lunax
4
Mmaax
2
Basis set
reference basis set
def2-umivereal--JIII (58)
1.5
orbital inference cutoff
4 Á
Training
batch size
2
grid point samples (training, with VO)
25,000
# grid point samples (validation/eesting, with VO)
35, 000
training steps (pretraining)
500, 000
training steps (fine-tuning)
0
Inference
batch size (with VO)
2
Max grid points in forward pass for Equation (5)
eSCN, K 8, L = 4,8
100, 000
Table Hyperparameters for baseline model architectures (3 2.0, all other hyperparameters are kept the same in Table 3).
Hyperparameter
Value
Charge3Net backbone
# interaction layers
Lmax
3
Feature irreps
167x0o
167x0e + 56x10 56x1e + 33x20 + 33x2e
MACE backbone
# interaction layers
4
Liax
3
Hidden irreps
64x0e
64x1o
64x2e
MLP irreps
128x0e
Common
Max grid points in forward pass
2, 0000 000
Initial learning rate
1 10 2
Terminal learning rate
1 10
19

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] provided a proper justification given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] " is not grounds for rejection. While the questions are phrased in binary way, acknowledge that the true answer is often more nuanced, so please just your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section 3, Section 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 5
20

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: This paper does not present theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 3, Section 4, Appendix A. Code is available at com com/kyynofx/scddp
Guidelines:
The answer NA means that the paper does not include experiments.
21

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
The data used in this paper is publicly available [7]. Code is available at Justification: attps://githir com/ /tt///oooofo///c/o///oo////oo/// x/cdp.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: /tppsscccppblic//geg.ece guidiiii/ggiie//eee/ee/ /CodeS//omddmmm issiisssion ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
22

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer:
[Yes]
Justification: Section 3, Section Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section 4, Section 5, Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
23

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttp::////eeu....eeeee wwterrissccpppupiiae ww/ppccicpc////ccccccccccccclllt Ceeeiii///iiiiieee wtthicssuieeleeeeeeeeeeeeeee
Answer: [Yes]
Justification: The authors believe this paper conforms, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Appendix B
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The authors believe the data and models presented in this work do not have high risk for misuse.
24

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Appendix A
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwooee...theeee.....o.oooooeoee com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: license.
Code is available at httppp///ggtttt...cooooott//ooot under an MIT
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
25

---

Answer:
[NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "sBsGRlmQOb", "reviewer": "Reviewer_paLd", "summary": "This application-oriented paper uses equivariant GNNs to predict charge density, represented orbital basis sets.", "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": "The limitations are sufficiently discussed in the paper.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- the idea of this paper is sound. Representing the charge density using atomic orbital basis set sounds more efficient than the voxel-based methods.\n- the performance of this new method is evidenced by well-designed experiments\n- the presentation of this paper is excellent", "weaknesses": "- I feel like the architectures are a simple variant of popular backbones. But this shouldn't affect the novelty of this application paper.\n- the computation cost is still considerable", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer paLd for helpful feedback and comments. We address each of the reviewer’s concerns below.\n\n> I feel like the architectures are a simple variant of popular backbones. But this shouldn't affect the novelty of this application paper.\n\nWe agree the model backbone architecture is a simple variant of the popular eSCN backbone. We believe the novelty of this paper lies in the combination of several representation, prediction, and training techniques – and most importantly, the significant improvements in model performance, which unlocks many directions for future research.\n\n> the computation cost is still considerable\n\nWe agree the training cost is considerable. The inference speed is more than one order of magnitude faster than existing state-of-the-art, but can potentially benefit from further improved basis set, architecture, and other factors. We are excited to explore the directions proposed in the discussion section to further accelerate model training and inference in future works. \n\n> 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.\n\nThank you for the constructive feedback. We provide the performance metrics for a tensor field network backbone used in Charge3Net [1] and MACE [2], with no virtual nodes and $\\beta = 2.0$. We compare these models to the small eSCN model with no virtual nodes and $\\beta = 2.0$, so the basis set expressive power is the same. The hyperparameters for the baseline models are adopted from previous works and are included in the rebuttal PDF.\n\n| Model | NMAPE | Mol. per Min. |\n|-----------------|---------------|---------------| \n| Charge3Net | 4.79 ± 0.026 | 505.10 | \n| MACE | 5.14 ± 0.024 | 660.25| \n| eSCN | 0.504 ± 0.001 | 675.47 |\n\nWe conclude the performance of the eSCN model is significantly better. We are also submitting our code, which can be used to reproduce all results in the paper, for peer review. Following NeurIPS 2024 instructions, we have sent an official comment to the AC that includes an anonymous link to the code base.\n\nWe look forward to further discussions if you have additional questions or suggestions.\n\n[1] Koker, Teddy, et al. \"Higher-order equivariant neural networks for charge density prediction in materials.\" npj Computational Materials 10.1 (2024): 161.\n\n[2] Batatia, Ilyes, et al. \"MACE: Higher order equivariant message passing neural networks for fast and accurate force fields.\" Advances in Neural Information Processing Systems 35 (2022): 11423-11436.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q8tHDaWheF", "reviewer": "Reviewer_L6ft", "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.", "questions": "Page 5. \"where z_alph... =1\" should be in the preceding sentence.\nPage 6, equation 9. Are the hidden scalar features h_i the direct scalar output of eSCN? Or the output of FCTensorProd.?\nPage 6, lines 230. If I am not mistaken, should be \"c3, c1+c3\" and \"c1/(1+c2)+c3\"", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Motivations are clear. Technical details are well presented and easy to follow. The experiments are also clear.", "weaknesses": "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer L6ft for helpful feedback and comments. We address each of the reviewer’s concerns below.\n\n> 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. \n\nThank you for your appreciation of our work as a valuable contribution. We believe our paper makes several significant technical contributions from an ML perspective which enable a significant improvement in model performance. We also believe charge density prediction has many impactful applications. We further elaborate on these two points below. \n\n> 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.\n\nIn the context of charge density prediction, the trade-off between accuracy and efficiency is a long-standing problem. We are driven to resolve this challenge and combine several well-motivated novel ideas. Our contributions include many design choices on how to build the GTO basis sets, the use of virtual nodes for charge density representations, expressive prediction networks, and an end-to-end training/finetuning procedure. These methods required a significant amount of trial and error to set up correctly and combined together. \n\nThese novel technical contributions lead to a final method that is simple yet very effective, and we see the simplicity of our final model as an advantage. Our experiments also reveal novel insights into understanding the bottleneck of ML-based charge density prediction, as well as model considerations when modeling a data-rich modality. Further, we believe the accuracy, efficiency and flexibility make our proposed method a good stepping stone for future research. We are happy to discuss further if your concerns remain.\n\n> 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.\n\nIn this work, we focus on striking a favorable accuracy-efficiency trade-off as it requires significant domain knowledge and efforts for downstream applications that are independent of the ML contributions. Here we elaborate on the applications of ML-based charge density prediction, and why we believe an accurate+efficient model is particularly impactful.\n\n- **Charge density is the core of DFT from which all properties can be derived.** There are many important molecular/materials properties, such as the electronic band structures, dipole moments, atomic spin densities, and effective bond orders, that can be directly computed from the charge density but not from an ML potential. An ML potential can be seen as a coarse-graining of the charge density prediction task where only energy/forces are considered.\n- **Self-consistency-field (SCF) from predicted charge density guarantees DFT accuracy.** ML-predicted charge density can be used as an initialization for the SCF calculation of DFT. This can reduce the computational cost of DFT and correctness is guaranteed. An ML potential does not have any performance guarantee.\n- **Charge density as pretraining.** Charge density is a data-rich modality that contains all the information from the quantum mechanical calculations, while energy/forces are a small portion of this information. We believe pre-training on charge density is a promising future direction to improve ML potentials, property prediction, and other atomistic modeling tasks.  \n- **Multi-modality in atomistic modeling.** Modeling of charge is still an active research topic in building better ML potentials, especially for reactive systems. Modeling the charge density is a promising direction for incorporating charge information into ML potentials. Efficiency is especially important when applying the ML potential in heavy workflows such as molecular dynamics simulations.\n- **An efficient+accurate model unlocks a major roadblock in learning with charge densities.** While the several directions stated above are potentially very impactful, a major roadblock for leveraging the charge density is the lack of a method that is both efficient and accurate. Our contribution addresses this challenge and can potentially make charge density a much more accessible modality for future research.\n\n> 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?\n\nOur method and ChargE3Net are fundamentally different in how we represent charge density, so the batch size is not directly comparable. We report all our hyperparameters in the appendix. Because ChargE3Net is a probe-based model, efficiency is limited by the fact that it needs to conduct neural message passing between all atoms and all positions of the charge density grid, which is very expensive as there are usually hundreds of thousands of grid positions for QM9 molecules.\n\nWhen comparing efficiency, we apply all models to the same set of molecules (test set of the QM9 charge density dataset). We use the pretrained QM9 model released by the ChargE3Net authors (therefore the original hyperparameters), and use optimized inference parameters to maximize the utilization of our GPU. For ChargE3Net, we process 20,000 probes in each batch which is ~1.3x faster than the default setting of 2,500 probes per batch from the ChargE3Net authors. All model efficiencies are benchmarked on the same machine and GPU.\n\n\n> Typos/clarifications on page 5 and 6\n\nThank you for pointing out the typos on line179 and 230. We will correct it in the final version. The hidden scalar features $h_i$ is the output of the final FCTensorProduct layer. We will clarify this in the final version.\n\nWe look forward to further discussions if you have additional questions or suggestions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ipyqYNOM2P", "reviewer": "Reviewer_SS9K", "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.", "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. \n\n2. At line 102, roto-translational --> rotation-translational     (the former seems not strictly accepted term by all the community). \n\n3. 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. \n\n[2]Li, Tianbo, et al. \"D4FT: A deep learning approach to Kohn-Sham density functional theory.\" arXiv preprint arXiv:2303.00399 (2023).", "limitations": "Yes. It is provided in the \"Discussion\" section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. By combining domain domain-knowledge (GTOs, Even-tempered Gaussian basis, etc) and new machine learning techniques (virtual orbitals, scaling factors for orbital exponents), the method achieves around 10% error reduction with around 30 times better efficacy in comparison to the previous state-of-the-art model (ChargE3Net).\n\n2. \"Methods\" are well-described with detailed explanations from physical and engineering point of view.", "weaknesses": "The following (at least major) weaknesses should be approached before acceptance: \n\nMajor weakness (the upper, the higher priority): \n\n1. 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).\n\n2. 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?\n\n3. 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. \n\n[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\n\nMinor (for further improvement): \n\n1. The explanation in \"Prediction layers\" can be improved. It would be strongly recommended to provide more information in the Appendix using more space.  \n\n2. 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. \n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer SS9K for helpful feedback and comments. We address each of the reviewer’s concerns below.\n\n> The paper only performed the experiments on QM9...\n\nThank you for your feedback. We follow your suggestions and extend our experiments to the MD dataset used in [1]. Due to limited time and compute resources, we were only able to complete experiments for four out of six molecules in the MD dataset. We report the NMAPE performance below (baseline metrics adopted from [1]):\n\n| Molecule   | SCDP (Ours) | InfGCN | CNN  | DeepDFT | DeepDFT2 | EGNN  | DimeNet | DimeNet++ | GNO   | FNO   | LNO   |\n|-----|------|--------|------|---------|----------|-------|---------|-----|-------|-------|-------|\n| ethanol    | **2.40 ± 0.26**       | 8.43   | 13.97| 7.34    | 8.83     | 13.90 | 13.99   | 14.24     | 82.35 | 31.98 | 43.17 |\n| benzene    | **1.15 ± 0.06**        | 5.11   | 11.98| 6.61    | 5.49     | 13.49 | 14.48   | 14.34     | 82.46 | 20.05 | 38.82 |\n| phenol     | **1.32 ± 0.07**        | 5.51   | 11.52| 9.09    | 7.00     | 13.59 | 12.93   | 12.99     | 66.69 | 42.98 | 60.70 |\n| resorcinol | **1.38 ± 0.08**        | 5.95   | 11.07| 8.18    | 6.95     | 12.61 | 12.04   | 12.01     | 58.75 | 26.06 | 35.07 |\n\nOur method significantly outperforms baseline models. In all experiments, we use an SCDP model with an eSCN of 4 layers, $L_{\\mathrm{max}}=3$, and $\\beta=1.5$. We train for 250000 steps while keeping all other training hyperparameters unchanged. A smaller model and shorter training schedule are used because the MD datasets are smaller in size. We will update results on all molecules in the MD dataset once training is completed. We appreciate your understanding of the time limitation of the rebuttal period.\n\n> Although not so crucial, the code is not provided in supplemental material.\n\nWe planned to do further improvement on the clarity of our codebase at the time of submission and we are happy to submit our code for peer review. Following NeurIPS 2024 instructions, we have sent an official comment to the AC that includes an anonymous link to the code base.\n\n> Even though the author mentions at line 214-215 that eSCN outperforms tensor field networks and MACE…explicit indication of the experiment would be desirable.\n\nThank you for the constructive feedback. We find that eSCN performs better initially in small-scale experiments. We provide the performance metrics for a tensor field network backbone used in Charge3Net and MACE, with no virtual nodes and $\\beta = 2.0$. We compare these models to the small eSCN model with no virtual nodes and $\\beta = 2.0$, so the basis set expressive power is the same. The hyperparameters for the baseline models are adopted from previous works and are included in the rebuttal PDF.\n\n| Model | NMAPE | Mol. per Min. |\n|-|-|-| \n| Charge3Net | 4.79 ± 0.026 | 505.10 | \n| MACE | 5.14 ± 0.024 | 660.25| \n| eSCN | 0.504 ± 0.001 | 675.47 |\n\nWe conclude the performance of the eSCN model is significantly better. The source code for the Charge3Net and MACE backbones are also included in our code submission. \n\n> The explanation in \"Prediction layers\" can be improved.\n\nThank you for your suggestion. We include pseudo code for the prediction process in the rebuttal PDF and would love to hear it if you have further feedback.\n\n> From the description of \"Prediction layers\", this method is a little difficult to apply to other models than eSCN.\n\nOur method can be directly applied 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. To predict the GTO coefficients in an equivariant way, the model needs to output a set of Irreps, that can be obtained through a tensor product of the atom feature irreps.\n\n>The inference time scaling in terms of the molecular size (atom number) is missing.\n\nThank you for the constructive feedback. In theory, the run time scales linearly ($O(N)$) with regard to the number of atoms (as opposed to an $O(N^3)$ scaling for DFT). We add additional experiments on the inference time scaling in terms of the molecular size using our most accurate model. The results are shown below:\n\n| Number of Atoms | Mol. per Min. | \n|---|---| \n| 12 | 230.31 | \n| 14 | 183.79 | \n| 16 | 155.50 | \n| 18 | 134.17 | \n| 20 | 119.76 |\n\nWe also include a plot in the rebuttal PDF. The result shows the runtime indeed scales linearly with regard to the number of atoms.\n\n> The idea of the scaling factors...a relation with the D4FT [2] Equations (20) and (21).\n\nThank you for pointing out this related work that we were unaware of. Both our method and D4FT aim to enhance the Gaussian basis set through learnable scaling factors – D4FT uses it in representing the wave function for KS-DFT; we use it in representing the charge density. We will update the paper/reference in the final version.\n\n> At line 102, roto-translational --> rotation-translational\n\nThank you for pointing this out, we will modify it in the final version.\n\n> In Table 1, the \"Mol. per min.\" shows around 10% increases in SCDP when adding scaling factors. Why? \n\nThis is a great question. There are two sources of additional compute when the scaling factors are introduced: (1) the network has an extra layer for predicting the scaling factors; (2) In equation (4), when the scaling factors are introduced, the normalization $z_{\\alpha, l, s}$ depends on the scaling factor $s$, which now changes in every inference steps. Therefore, the normalization needs to be recomputed in every forward pass with non-negligible computational cost. Without the scaling factors, the normalization only needs to be computed once when initializing the GTOs. This is reflected in L112-L115 of `gtos.py` of our source code.\n\nWe look forward to further discussions if you have additional questions or suggestions.\n\n[1] Cheng, Chaoran, and Jian Peng. \"Equivariant neural operator learning with graphon convolution.\" NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ViJzW9lSvQ", "reviewer": "Reviewer_nCiL", "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.\n\nSection 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.\n\nSection 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.", "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.\n* Can you include traditional DFT throughput numbers in the table for reference? This is to get a sense of the speedup.\n* Can you describe in more detail why an equivariant network is required? Have you tried evaluating non-equivariant networks?", "limitations": "The paper includes a discussion on limitations in Section 5.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes a new method for a relevant problem in AI for scientific applications. Its strengths include:\n* A new method for calculating charge density and integrating ML methods into computational chemistry workflows. [Originality, Significance]\n* A detailed description of the method, relevant background and experimental details. [Quality, Clarity]\n* Experimental results show convincing improvement compared to the state-of-the art. [Quality, Significance]", "weaknesses": "The paper could be improved by providing more detail along the following:\n* 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]\n* Provide greater clarity on how their proposed method for Charge Density Representation differs from prior approaches, such as a summary table. [Clarity, Significance]\n* Address the questions below.\n\n-----------\n\nAssessment updated during discussion period.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer nCiL for helpful feedback and comments. We address each of the reviewer’s concerns below.\n\n> 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]\n\nThank you for the constructive feedback. This is a great point to include in the paper. We briefly summarize several key differences here:\n\n- **Charge density is the core of DFT from which all properties can be derived.** There are many important molecular/materials properties, such as the electronic band structures, dipole moments, atomic spin densities, and effective bond orders, that can be directly computed from the charge density but not from an ML potential. An ML potential can be seen as a coarse-graining of the charge density prediction task where only energy/forces are considered. The charge density represents all the information from DFT calculations.\n- **Self-consistency-field (SCF) from predicted charge density guarantees DFT accuracy.** ML-predicted charge density can be used as an initialization for the SCF calculation of DFT. This can reduce the computational cost of DFT and correctness is guaranteed. An ML potential does not have any performance guarantee.\n- **Detailed charge information matters.** Modeling of charge is still an active research topic in building better ML potentials, especially for reactive systems. Charge density encapsulates this crucial information for high-fidelity modeling.\n- **Charge density Pros**: the finest-grained information from DFT; all chemical properties can be derived; can be combined with conventional DFT for accuracy guarantee; can model complex systems more accurately.\n- **Charge density Cons**: more costly than ML potentials; not required for many relaxation/simulation tasks that don’t require very high accuracy.\n\nWe will include these discussions in the final version.\n\n> Provide greater clarity on how their proposed method for Charge Density Representation differs from prior approaches, such as a summary table. [Clarity, Significance]\n\nThank you for the great suggestion. Our method builds on top of atomic orbital representations of charge density and introduces even-tempered basis, virtual nodes/orbitals, and learnable scaling factors to improve its expressivity while preserving its efficiency. These elements are also highly tunable to trade-off accuracy and efficiency.\n\nWe summarize a comparison of the different representations below:\n\n| Method | Expressive? | Efficient? | Flexible trade-off? | \n|---------------------------|-------------|------------|----------------------| \n| Orbital-based (e.g., [1]) | No | Yes | No | \n| Probe-based (e.g., [2]) | Yes | No | No | \n Ours | Yes | Yes | Yes |\n\n\nWe will include these further details in the final version.\n\n> 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?\n\nIn early small-scale experiments, we have considered other basis sets such as `def2-universal-JKFIT` and `aug-cc-pVQZ`. We choose def2-QZVPPD because it’s one of the most expressive and popular basis sets and its promising initial results. We are unable to complete this experiment at the time of rebuttal due to limitations on time and compute but would be happy to include an ablation study in the final version.\n\n> Can you include traditional DFT throughput numbers in the table for reference?\n\nUnfortunately, the DFT runtime was not included in the QM9 charge density dataset [3]. The throughput of conventional DFT depends a lot on the level of theory, convergent criterion, and the size of the system as it scales as $O(N^3)$. Practical calculations of a small molecule usually take a few minutes to complete, and the run time grows rapidly when the number of atoms increases. ML-based charge density prediction methods scale linearly with the number of atoms, and therefore are promising to enable larger-scale calculations at higher levels of theories.\n\n> Can you describe in more detail why an equivariant network is required? Have you tried evaluating non-equivariant networks?\n\nAn equivariant model is highly desired because the charge density is an equivariant quantity. That is, when the molecule is rotated/translated, the charge density should rotate/translate accordingly. When predicting the charge density through the orbital basis representation, the basis set coefficients correspond to different components of spherical harmonics. An equivariant model is able to predict the charge density equivariantly through equivariant prediction of the basis set coefficient. We only evaluated equivariant networks because a non-equivariant network breaks the physical constraints of the charge density prediction task. Different from an ML potential where the energy is an invariant quantity and the forces can be obtained through taking its derivative with regard to atomic positions, the basis set coefficients that we predict are not scalars nor derivatives of scalars, and can be of higher tensor order. Equivariant models are highly suitable for predicting these quantities.\n\nWe look forward to further discussions if you have additional questions or suggestions.\n\n[1] Rackers, Joshua A., et al. \"A recipe for cracking the quantum scaling limit with machine learned electron densities.\" Machine Learning: Science and Technology 4.1 (2023): 015027.\n\n[2] Koker, Teddy, et al. \"Higher-order equivariant neural networks for charge density prediction in materials.\" npj Computational Materials 10.1 (2024): 161.\n\n[3] Jørgensen, Peter Bjørn, and Arghya Bhowmik. \"Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids.\" npj Computational Materials 8.1 (2022): 183.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Area Chairs and Reviewers,

Thank you for your time and consideration in reviewing our paper. Following the suggestions from the reviewers, here we summarize our responses and several improvements we aim to make in the next version.

- We report the performance metrics of alternative model backbone architectures to support our choice of eSCN as our main backbone architecture.
- We elaborate on the significance and importance of learning to predict charge density in terms of downstream applications, ML applications, and novel insights of our paper. In particular, we emphasize the importance of efficiency in unblocking future research in this topic.
- We conduct additional experiments on the MD dataset [1,2] used in [3] (detailed in response to reviewer SS9k) and will add the new experiment results to the final version.
- We report the inference time scaling with regard to molecular size results to verify the theoretical linear scaling of our proposed method.
- We include more detailed descriptions of our design choices and model.
- We clarify the difference between charge density prediction models and ML potentials.
- We submit our code through an anonymous link for peer review. Following the NeurIPS 2024 guideline, we sent an official comment containing the link to the AC.

Our rebuttal PDF contains the following:

- A figure that demonstrates the inference time scaling with regard to molecular size.
- Pseudo code for the charge density prediction procedure of our proposed method.
- Hyperparameters for alternative model backbone architectures.

Thank you again for your feedback which greatly helped us improve the clarity and thoroughness of our paper. We look forward to further discussions if you have additional questions or suggestions.

[1] Bogojeski, Mihail, et al. "Quantum chemical accuracy from density functional approximations via machine learning." Nature communications 11.1 (2020): 5223.

[2] Brockherde, Felix, et al. "Bypassing the Kohn-Sham equations with machine learning." Nature communications 8.1 (2017): 872.

[3] Cheng, Chaoran, and Jian Peng. "Equivariant neural operator learning with graphon convolution." Advances in Neural Information Processing Systems 36 (2024).

Meta-review: 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.
Decision: Accept (poster)