Equivariant Spherical Transformer for Efficient Molecular Modeling

Equivariant Graph Neural Networks (GNNs) have significantly advanced the modeling of 3D molecular structure by leveraging group representations. However, their message passing, heavily relying on strictly equivariant operations, suffers from restricted expressiveness due to the limited non-linearity and low degree of group representations. To overcome this, we introduce the Equivariant Spherical Transformer (EST), a novel plug-and-play module that applies a Transformer-based architecture to the Fourier spatial domain of group representations. The integration of EST enhances the model's expressiveness while preserving the crucial equivariant inductive bias through a uniform sampling strategy of spherical Fourier transforms. As demonstrated by our experiments on challenging benchmarks like OC20, MPtrj and QM9, EST-based models achieve state-of-the-art performance. For the complex molecular systems within OC20, small models empowered by EST can outperform some larger models and those using additional data. In addition, we provide both theoretical and experimental validation of EST's equivariance as well, paving the way for new research in this area.