E3STO: Orbital Inspired SE(3)-Equivariant Molecular Representation for Electron Density Prediction

TL;DR

The paper tackles the challenge of predicting molecular electron densities, a task tied to solving the many-electron Schrödinger equation and expensive for DFT. It introduces an SE(3)-equivariant framework that learns molecule-dependent Slater-Type Orbital–inspired basis functions to construct orbital-like representations and compute $\rho(\mathbf{r})$ via $\Psi(\mathbf{r})$; density is obtained as $\rho(\mathbf{r}) = ||\Psi(\mathbf{r})||^2$. The approach leverages Tensor Field Networks with dynamic radial coefficients and higher-order spherical harmonics to capture complex symmetries, achieving state-of-the-art density predictions, notably large gains on MD datasets (up to 30–70% over baselines). Limitations include difficulties with lattice-based periodic systems and scaling to very large molecules, suggesting future work on periodicity and scalability. Overall, the work demonstrates that combining physics-informed STO-like basis expansions with $SE(3)$-equivariant representations yields highly accurate, efficient electron-density predictions with strong potential for quantum-mechanical modeling.

Abstract

Electron density prediction stands as a cornerstone challenge in molecular systems, pivotal for various applications such as understanding molecular interactions and conducting precise quantum mechanical calculations. However, the scaling of density functional theory (DFT) calculations is prohibitively expensive. Machine learning methods provide an alternative, offering efficiency and accuracy. We introduce a novel SE(3)-equivariant architecture, drawing inspiration from Slater-Type Orbitals (STO), to learn representations of molecular electronic structures. Our approach offers an alternative functional form for learned orbital-like molecular representation. We showcase the effectiveness of our method by achieving SOTA prediction accuracy of molecular electron density with 30-70\% improvement over other work on Molecular Dynamics data.

Figures (9)

• Figure 1: A) Schematic of total electron density prediction given molecular structure. B) $Invariants$ algorithm for extracting meaningful invariant features from higher order equivariant features. C) $Embed$ function which uses a simple lookup table by atom type. D) $Convolution$ which mixes features of neighboring atoms and creates richer representation for the molecular orbitals.
• Figure 2: Top and bottom rows are the error and prediction of electronic density of two different molecules. Left and right errors are NMAE=0.47% and 0.49%.
• Figure 3: Left: prediction, Right: error. NMAE=0.49%
• Figure 4: Left: prediction, Right: error. NMAE=0.57%
• Figure 5: Left: prediction, Right: error. NMAE=0.56%