Introducing Virtual Points in Equivariant Networks by Extending Atom Representation for Effective Prediction

Recent equivariant models embed a molecule as a set of atoms fixed in three-dimensional space, which is analogous to a ball-and-stick view. This perspective provides a concise view of molecular configurations; however, these representations may be limited in incorporating the surrounding environments of atomic nuclei, including electron configurations. To overcome this limitation, we propose neural polarization (NP), a novel method that extends equivariant networks by embedding each atom as a pair of an atom and virtual points. Motivated by electron density configurations, NP represents each atom as a pair comprising the original fixed atom and a virtual atom whose position is updated through parameterization during model training. NP can be flexibly applied to most types of existing equivariant models. We showed that NP can improve the prediction performance of existing models over a wide range of targets, including electron density. Our experimental results on various benchmarks suggest new insights, indicating that the extended atomic representations can improve the overall molecular tasks.