Neural Polarization: Toward Electron Density for Molecules by Extending Equivariant Networks

TL;DR

Neural Polarization extends equivariant neural networks to include electron density representation, improving molecular property predictions by capturing polarization effects and electron distribution more accurately.

Contribution

It introduces a novel embedding method that incorporates electron density into equivariant models, enhancing their expressiveness and predictive performance.

Findings

Improved prediction accuracy across various molecular targets.

Enhanced expressiveness and equivariance in neural network models.

Applicable to most existing equivariant models.

Abstract

Recent SO(3)-equivariant models embedded a molecule as a set of single atoms fixed in the three-dimensional space, which is analogous to a ball-and-stick view. This perspective provides a concise view of atom arrangements, however, the surrounding electron density cannot be represented and its polarization effects may be underestimated. To overcome this limitation, we propose \textit{Neural Polarization}, a novel method extending equivariant network by embedding each atom as a pair of fixed and moving points. Motivated by density functional theory, Neural Polarization represents molecules as a space-filling view which includes an electron density, in contrast with a ball-and-stick view. Neural Polarization can flexibly be applied to most type of existing equivariant models. We showed that Neural Polarization can improve prediction performances of existing models over a wide range of targets. Finally, we verified that our method can improve the expressiveness and equivariance in terms of mathematical aspects.