Title: ELECTRA: A Cartesian Network for 3D Charge Density Prediction with Floating Orbitals
Venue: NeurIPS 2025
Accepted: yes
Keywords: charge density

quantum chemistry

equivariance

graph neural networks

molecular representation
Abstract: We present the Electronic Tensor Reconstruction Algorithm (ELECTRA) - an equivariant model for predicting electronic charge densities using floating orbitals. Floating orbitals are a long-standing concept in the quantum chemistry community that promises more compact and accurate representations by placing orbitals freely in space, as opposed to centering all orbitals at the position of atoms. Finding the ideal placement of these orbitals requires extensive domain knowledge, though, which thus far has prevented widespread adoption. We solve this in a data-driven manner by training a Cartesian tensor network to predict the orbital positions along with orbital coefficients. This is made possible through a symmetry-breaking mechanism that is used to learn position displacements with lower symmetry than the input molecule while preserving the rotation equivariance of the charge density itself. Inspired by recent successes of Gaussian Splatting in representing densities in space, we are using Gaussian orbitals and predicting their weights and covariance matrices. Our method achieves a state-of-the-art balance between computational efficiency and predictive accuracy on established benchmarks. Furthermore, ELECTRA is able to lower the compute time required to arrive at converged DFT solutions - initializing calculations using our predicted densities yields an average 50.72 % reduction in self-consistent field (SCF) iterations on unseen molecules.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces the novel machine learning model ELECTRA for the prediction of the electron density in molecular systems. Contrary to previous approaches, the density is not represented in an atom-centered basis expressed in spherical harmonics, but through floating orbitals using ordinary Gaussian functions. The positions and covariance matrices of these floating orbitals are learned by a neural network alongside the corresponding coefficients. Since the neural network is equivariant ELECTRA also introduces a symmetry breaking mechanism in order to allow for expressive densities also in molecular systems which exhibit a high degree of symmetry. The model achieves state of the art results, while being an order of magnitude faster than models with similar performance.\n\nSTRENGTHS AND WEAKNESSES:\n- Strengths:\n    - The paper is written in a clear and understandable fashion. I also found the blue boxes highlighting key takeaways helpful.\n    - The model shows very good results in particular with respect to inference speed at comparable accuracy.\n- Weaknesses\n    - Although the need for a symmetry breaking mechanism is convincingly demonstrated by the conducted ablation studies, I am not sure about the claim that this is due to local symmetries within the molecules. The mathematical arguments provided in Appendix C are perfectly fine for global symmetries of molecules, but break down for local symmetries in a message passing framework where messages are propagated beyond the locally symmetric regions.  The problem I see is that in the proof it is assumed that the output features of nodes within a sub graph depend only on input features within the same sub graph. For a general message passing network this is however not the case.\n    - The final densities are constructed using the ReLU function to ensure positivity everywhere. However this can result in regions of unphysical densities where they are constant 0. This could reduce the performance in downstream tasks, but might also be problematic during training, e.g. when the position of a Gaussian is within such a region of vanishing density. In such a case also the gradient signal on the parameters of the respective Gaussian would become very small.\n    - In the comparison of evaluation metrics on the MD dataset, ELECTRA is compared only to the fastest SCDP model, but given the results from the QM9 dataset I would expect the SCDP+BO (L = 6) model to perform much closer (maybe even better) compared to ELECTRA.\n    - Since the densities are not represented in a LCAO basis the resulting densities cannot straightforwardly be used as initial guess in conventional Kohn Sham DFT which is one of the most common use cases for machine learned electron densities for molecular systems.\n\nQUESTIONS:\n- On page 5 it is mentioned that ELECTRA uses an atomic embedding function inspired by SpookyNet, but I could not find any details on the exact implementation. Could the authors provide details on this in the Appendix?\n- I would be interested in how well a Gaussian mixture model with the same number of Gaussians as in ELECTRA is able to approximate the target densities in principle in order to know how much of the density error can be attributed to the choice of basis and how much to the limited accuracy of the model.\n- If possible I would like to see a comparison to the SCDP+BO (L = 6) model on the MD dataset.\n\nI am willing to increase my score if the above points are adequately addressed.\n\nLIMITATIONS:\nLimitations are addressed. Not being able to directly use the density prediction in conventional Kohn-Sham SCF procedures is another point which might still be added.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Electronic Tensor Reconstruction Algorithm (ELECTRA), which uses floating orbitals to predict electronic charge densities. Inspired by the Gaussian Splatting representation, ELECTRA replaces costly traditional linear combinations of fixed basis functions with floating orbitals. Since designing such orbitals usually requires expert intuition, ELECTRA learns them in a purely data-driven manner, enabling high-resolution charge density prediction. Experiments on well-known benchmarks demonstrate its high accuracy.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The narrative is clearly explained. The paper introduces floating orbitals and justifies their data-driven prediction via machine learning, which makes the paper easy to follow.\n- The authors pinpoint the directional bias that arises in HotPP message passing and propose a dedicated Debiasing Layer to address it.\n- Experimental results show that ELECTRA achieves fast and accurate charge-density prediction.\n\nWeaknesses:\n- The model largely inherits purely the equivariance property from the existing HotPP framework.\n- L.224: \"However, we find, empirically, that this does not hinder strong performance.\" Since equivariance is essential to the method, substantial empirical or theoretical justification is necessary.\n- A key contribution of Gaussian Splatting is its tight link to efficient, differentiable rendering and optimization. However, ELECTRA does not exploit this property; it merely outputs a set of Gaussians. Alternative representations, such as point clouds or sparse voxels, should be considered. Alternatively, a loss function tailored to the Gaussian representation must be provided to justify this design choice.\n\nQUESTIONS:\n- Why does the method compute directions via eigenvectors? In highly symmetric cases, these eigenvectors are unstable, and noise can easily switch their order or reverse their orientation. How is this handled?\n- Debiasing is an important contribution, yet the origin of the bias itself is unclear. Which HotPP component introduces this bias? Understanding the root cause may be more valuable than an ad hoc debiasing fix.\n- This approach seems heavily tied to HotPP. While choosing an appropriate backbone is reasonable, what about other equivariant GNNs—or even backbones that ignore symmetry? Since the core idea is to break symmetry within an otherwise symmetric network, broader validation would strengthen the claim.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes predicting electronic charge density using floating orbitals inspired by Gaussian splatting. By leveraging an equivariant neural network architecture, symmetry-breaking initialization, and debiasing layers, the model aims to predict charge densities directly from molecular structures.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Utilizing a Gaussian splatting-inspired idea for electron density prediction is clever and leads to superior performance and significantly reduced inference time.\n\nWeaknesses:\n- The manuscript makes a major assumption regarding the rotational invariance of the electron density, which is conceptually flawed.\n  - If the charge density were strictly rotationally invariant, then $\\rho(\\mathbf{r}_1) = \\rho(\\mathbf{r}_2)$ would hold for all vectors with the same magnitude, i.e., $\\|\\mathbf{r}_1\\|=\\|\\mathbf{r}_2\\|$. This is clearly not the case for most molecular charge densities.\n  - In my view, the authors appear to confuse invariance of the physical quantity under rigid-body rotation of the molecule with the concept of group action on function spaces (regular representations). For example, as in [1] (see Equation 2), a rotated density transforms as $\\rho(\\mathbf{R} \\mathbf{r}) = (T_{\\mathbf{R}} \\cdot \\rho)(\\mathbf{r})$, which is not invariance.\n- A more detailed introduction to Cartesian tensors (e.g., covariances) is necessary, especially given that the topic lies at the intersection of quantum chemistry and machine learning. Many ML researchers may not be familiar with the geometric tensor concepts; the background and appendices do not sufficiently support understanding.\n- While using Cartesian tensor-based basis functions is an interesting idea, the manuscript does not provide adequate justification for why decomposing functions in this way is natural. Additional explanations comparing conventional basis choices—such as Gaussian-type orbitals, Slater-type orbitals, or plane-wave basis functions—to Cartesian basis functions would greatly aid comprehension.\n- The statement that “Gaussians are equivalent to traditional Cartesian orbital functions with angular quantum number $l=2$, an extra nonlinearity, and simplified radial dependency” is unconvincing. The relationship between spherical harmonics-based basis functions and Cartesian functions is not clearly described. Moreover, while the manuscript refers to higher-order tensors (e.g., vectors for dipoles or matrices for polarizabilities), the connection between geometric tensor rank and quantum angular momentum is fundamentally different and requires clarification.\n  - Specifically, the relationship between Equation (15) and Equation (16) is ambiguous. There is no explanation of why the angular quantum number $l$ in both equations should be interpreted identically.\n\n\nTypos and Minor Issues:\n- Equation (2): $\\mathbf{r}_i$ is not defined.\n- Inconsistent notation between \\mathbf{} and \\boldsymbol{} for vectors, e.g., $\\mathbf{f}$ vs. $\\boldsymbol{f}$, and $\\mathbf{v}$ vs. $\\boldsymbol{v}$.\n- Duplicated results for InfGCN in Table 1.\n- Line 192: Inconsistent use of “figure” and “fig”; consistent referencing with \\ref or \\Cref in LaTeX is recommended.\n- Equation (11): The definition of $\\hat{\\cdot}$ is missing.\n- Line 208: The notation k=1..N is informal; it should be $k=1,\\ldots,N$.\n- Line 538: Inconsistent use of $\\top$ and $T$ for transposition.\n- Line 555: Potential typo, where $g$ should be $G$.\n- Equation (18): $\\Phi_{l=2}$ should be $\\Phi_{\\alpha, l=2}$ for clarity.\n\nQUESTIONS:\n1.\tInterpretation of Figures 1 and 2: What is the frame or vector space represented in these figures? Specifically, are the vectors shown in Figures 1a and 1c the predicted mean positions $\\boldsymbol{\\mu}_{A,j}$ in Equation (6)? A detailed explanation of what each vector represents in these figures would greatly improve clarity.\n2.\tGPU Memory Requirements: How much GPU memory was used during training and inference? Are there specific trade-offs in inference performance compared to prior work?\n3.\tChannel Dimension Limitation: Line 184 specifies $N_c = 8 \\cdot M_e$. Does this imply that the model can only handle atoms with fewer than 8 valence electrons? How generalizable is the model to elements beyond the first two rows of the periodic table?\n4.\tDensity Grid Sampling: How did you handle sampling of the density during training? Did you employ specific sampling algorithms over the reference density grids in the dataset?\n5.\tPractical Use Cases: Could you elaborate on whether ELECTRA’s predicted electron densities can improve computational physics simulations, such as providing initial guesses for SCF iterations or being directly used in orbital-free DFT calculations?\n\n---\nupdate Quality 1 -> 2, Clarity 1 -> 2\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: SUMMARY:
The paper introduces ELECTRA, a model for predicting electronic charge density using floating orbitals, a concept rooted in quantum chemistry but adapted in a data-driven manner via an Equivariant backbone neural network. Incorporates a symmetry-breaking mechanism using eigenvectors of the moment of inertia tensor to enhance expressivity while preserving rotational equivariance, addressing limitations of highly symmetric inputs. Compared to prior work (e.g., Fu et al., 2024; Rackers et al., 2023), ELECTRA’s use of floating orbitals and symmetry-breaking is distinct, as it avoids reliance on fixed basis sets or bond-centered orbitals, offering a fully data-driven solution (Page 9). This is a significant departure from traditional DFT methods and other ML-based approaches like SCDP or Charge3Net, which often use atom-centered or spherical harmonic representations.

STRENGTHS AND WEAKNESSES:
Strengths - Incorporating floating orbitals with an Equivariant backbone neural network is novel. 
Weaknesses - The author should also address when this would fail in certain cases.

QUESTIONS:
Can you also address some examples of molecules where this approach fails or has less advantages?

LIMITATIONS:
N/A
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02