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

Machine learning models for 3D molecular property prediction typically rely on atom-based representations. However, these representations 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 CNNs: atom types, raw electron density, and its gradient magnitude, across two key molecular tasks: protein–ligand binding affinity (PDBbind) and quantum property prediction (QM9). On PDBbind, all representations perform similarly with full data, but in low-data regimes, density-based inputs outperform atom types. Interestingly, a shape-based baseline performs comparably—suggesting that shape occupancy dominates in this task. On QM9, where labels are computed using Density Functional Theory (DFT) but input densities are generated using a less accurate quantum chemistry method (XTB), density-based inputs still outperform atom-based ones at scale. This indicates that the benefit arises not from perfect input–label alignment, but from the rich structural and electronic signal encoded in the density. These results highlight the task- and regime-specific strengths of density-derived inputs, improving data efficiency in affinity prediction and accuracy in quantum property modeling.