Rotational Sampling: A Plug-and-Play Encoder for Rotation-Invariant 3D Molecular GNNs

TL;DR

This paper introduces a novel rotational sampling encoder for 3D molecular GNNs that achieves rotation invariance and improves prediction accuracy, robustness, and efficiency in molecular property prediction tasks.

Contribution

It proposes a plug-and-play rotational sampling module that approximates rotation invariance by expectation over SO(3), with a post-alignment strategy for strict invariance, enhancing 3D molecular GNN performance.

Findings

Outperforms existing methods on QM9 and C10 datasets.

Achieves superior accuracy, robustness, and generalization.

Maintains low computational complexity and interpretability.

Abstract

Graph neural networks (GNNs) have achieved remarkable success in molecular property prediction. However, traditional graph representations struggle to effectively encode the inherent 3D spatial structures of molecules, as molecular orientations in 3D space introduce significant variability, severely limiting model generalization and robustness. Existing approaches primarily focus on rotation-invariant and rotation-equivariant methods. Invariant methods often rely heavily on prior knowledge and lack sufficient generalizability, while equivariant methods suffer from high computational costs. To address these limitations, this paper proposes a novel plug-and-play 3D encoding module leveraging rotational sampling. By computing the expectation over the SO(3) rotational group, the method naturally achieves approximate rotational invariance. Furthermore, by introducing a carefully designed post-alignment strategy, strict invariance can be achieved without compromising performance. Experimental evaluations on the QM9 and C10 Datasets demonstrate superior predictive accuracy, robustness, and generalization performance compared to existing methods. Moreover, the proposed approach maintains low computational complexity and enhanced interpretability, providing a promising direction for efficient and effective handling of 3D molecular information in drug discovery and material design.