Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ELECTRA: A Cartesian Network for 3D Charge Density Prediction with Floating Orbitals

Jonas Elsborg, Luca Thiede, Alán Aspuru-Guzik, Tejs Vegge, Arghya Bhowmik

TL;DR

ELECTRA tackles the expensive SCF bottleneck in DFT by learning a direct, data-driven prediction of electronic charge densities using floating Gaussian orbitals. It deploys a rotation-equivariant Cartesian tensor backbone (based on HotPP) to predict Gaussian weights, means, and covariances, with symmetry-breaking and debiasing mechanisms enabling flexible, lower-symmetry orbital placements while preserving density invariance. The model represents density as a Gaussian mixture, normalizes to the correct electron count, and is trained with a normalized MAE objective. On QM9 and MD datasets, ELECTRA achieves state-of-the-art or competitive accuracy with significantly faster inference, and it reduces SCF iterations by about 50% on unseen molecules, illustrating substantial practical impact for accelerating DFT-based workflows. This data-driven approach shifts toward learned density representations with floating orbitals, offering avenues for extension to periodic systems and hybrid architectures that blend floating and atom-centered orbitals for optimal efficiency and accuracy.

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.

ELECTRA: A Cartesian Network for 3D Charge Density Prediction with Floating Orbitals

TL;DR

ELECTRA tackles the expensive SCF bottleneck in DFT by learning a direct, data-driven prediction of electronic charge densities using floating Gaussian orbitals. It deploys a rotation-equivariant Cartesian tensor backbone (based on HotPP) to predict Gaussian weights, means, and covariances, with symmetry-breaking and debiasing mechanisms enabling flexible, lower-symmetry orbital placements while preserving density invariance. The model represents density as a Gaussian mixture, normalizes to the correct electron count, and is trained with a normalized MAE objective. On QM9 and MD datasets, ELECTRA achieves state-of-the-art or competitive accuracy with significantly faster inference, and it reduces SCF iterations by about 50% on unseen molecules, illustrating substantial practical impact for accelerating DFT-based workflows. This data-driven approach shifts toward learned density representations with floating orbitals, offering avenues for extension to periodic systems and hybrid architectures that blend floating and atom-centered orbitals for optimal efficiency and accuracy.

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.

Paper Structure

This paper contains 41 sections, 40 equations, 5 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Charge density prediction
  4. Use $\rho(\boldsymbol{r})$ as an initial guess:
  5. Use $\rho(\boldsymbol{r})$ directly:
  6. Equivariance and Cartesian tensors
  7. Methods
  8. The ELECTRA model
  9. Equivariant backbone neural network.
  10. Atomic embeddings and variable basis set size.
  11. Symmetry-breaking
  12. Debiasing layers.
  13. Density construction
  14. Normalization.
  15. Objective function
  16. ...and 26 more sections

Figures (5)

  • Figure 1: (a) Model output without symmetry-breaking: Equivariant neural networks constrain their output to have the same symmetry as the input. If the input molecule is highly symmetric, this leads to highly constrained Gaussian positions. (b) To solve this issue, ELECTRA initializes each atom's $l=1$ vector features with the eigenvectors of the moment of inertia tensor as calculated in that specific atom, which breaks the input symmetry but retains rotational equivariance. (c) Model output after first linear layer with symmetry-breaking: The model can learn its own set of symmetry-breaking vectors, allowing output to not be constrained by the symmetry of the input molecule.
  • Figure 2: (a) The initial symmetry-breaking objects of the NH3 molecule (Ammonia). (b) The output of a HotPP model with symmetry-breaking but without debiasing layers: The message passing induces a directional bias that concentrates vectors along certain directions, (c) The output of our model with debiasing layers: The output vectors don't show any visible bias.
  • Figure 2: NMAE[%] on MD test sets for ELECTRA, SCDP fu2024recipe, GPWNO kim2024gaussian and InfGCN cheng2024equivariant, as well as the $N=1$ Gaussian Overfit experiment. ELECTRA models were trained for 10 GPU hours.
  • Figure 3: (a) Predicted density for C7H9NO using ELECTRA (NMAE = 0.19 %, red = high density, blue = lower density). (b) Ground-truth density. (c) Gaussian placements (red: $w_{A,j}>0$, blue: $w_{A,j}<0$). (d) 0.001 $\mathrm{e}/\mathrm{bohr}^3$ error isosurfaces (blue = over-prediction, red = under-prediction).
  • Figure 4: Gaussian placements vs. ground-truth electron density for Benzene.

Theorems & Definitions (1)

  • Definition D.1