等变球面Transformer用于高效分子建模
Equivariant Spherical Transformer for Efficient Molecular Modeling
Junyi An, Xinyu Lu, Yunfei Shi, Peijia Lin, Licheng Xu, Qianwei Tang, Chao Qu, Fenglei Cao 等 9 位作者
AI 速读 基于摘要
提出可插拔的等变球面Transformer(EST)模块,在球面空间表示上施加Transformer注意力以混合不同采样方向的信息,并在OC20、MPtrj、QM9等基准上报告了先进性能。
cs.LG · cs.AI
查看原文摘要
Equivariant Graph Neural Networks (GNNs) have significantly advanced the modeling of 3D molecular structure by leveraging group representations. However, message passing built only from exactly equivariant operations can have restricted practical expressiveness at a fixed representation degree. We introduce the Equivariant Spherical Transformer (EST), a plug-and-play module that applies Transformer attention to spherical spatial representations obtained from group representations. EST explicitly mixes information across sampled orientations, unlike prior orientation-wise spherical operations, while retaining a strong approximate equivariant inductive bias through near-uniform spherical sampling. 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. We characterize the ideal closed-domain condition under which the continuous operation is exactly equivariant and empirically quantify the finite-sampling accuracy--equivariance--memory trade-off used in practice.