Equivariant Spherical Transformer for Efficient Molecular Modeling

TL;DR

The paper introduces the Equivariant Spherical Transformer (EST), a novel architecture that enhances the expressiveness of molecular modeling by applying Transformer-like mechanisms to group representations while maintaining equivariance, leading to state-of-the-art results.

Contribution

EST is a new plug-and-play framework that improves the expressiveness of equivariant GNNs for molecular modeling by integrating Fourier transforms and Transformer architecture.

Findings

EST achieves state-of-the-art performance on OC20 and QM9 benchmarks.

Small EST-based models outperform larger models with more data.

Theoretical and experimental validation confirms EST's equivariance properties.

Abstract

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 Clebsch-Gordan tensor product convolutions, 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 framework that applies a Transformer-like architecture to the Fourier spatial domain of group representations. EST achieves higher expressiveness than conventional models 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 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 to demonstrating such strong expressiveness,we provide both theoretical and experimental validation of EST's equivariance as well, paving the way for new research in this area.