Beyond Atoms: Evaluating Electron Density Representation for 3D Molecular Learning

TL;DR

This paper evaluates electron density maps as a physically grounded alternative to atom-based representations for 3D molecular learning, demonstrating their advantages in data efficiency and accuracy across different tasks and regimes.

Contribution

It systematically compares voxel-based electron density inputs with atom types for molecular property prediction, highlighting their benefits in low-data and quantum property tasks.

Findings

Density-based inputs outperform atom types in low-data regimes for affinity prediction.

Electron density inputs improve quantum property prediction accuracy at scale.

Shape-based baseline performs similarly to density inputs in certain tasks.

Abstract

Machine learning models for 3D molecular property prediction typically rely on atom-based representations, which may overlook subtle physical information. Electron density maps -- the direct output of X-ray crystallography and cryo-electron microscopy -- offer a continuous, physically grounded alternative. We compare three voxel-based input types for 3D convolutional neural networks (CNNs): atom types, raw electron density, and density gradient magnitude, across two molecular tasks -- protein-ligand binding affinity prediction (PDBbind) and quantum property prediction (QM9). We focus on voxel-based CNNs because electron density is inherently volumetric, and voxel grids provide the most natural representation for both experimental and computed densities. On PDBbind, all representations perform similarly with full data, but in low-data regimes, density-based inputs outperform atom types, while a shape-based baseline performs comparably -- suggesting that spatial occupancy dominates this task. On QM9, where labels are derived from Density Functional Theory (DFT) but input densities from a lower-level method (XTB), density-based inputs still outperform atom-based ones at scale, reflecting the rich structural and electronic information encoded in density. Overall, these results highlight the task- and regime-dependent strengths of density-derived inputs, improving data efficiency in affinity prediction and accuracy in quantum property modeling.