E2Former: A Linear-time Efficient and Equivariant Transformer for Scalable Molecular Modeling

  • 2025-02-03 18:46:30
  • Yunyang Li, Lin Huang, Zhihao Ding, Chu Wang, Xinran Wei, Han Yang, Zun Wang, Chang Liu, Yu Shi, Peiran Jin, Jia Zhang, Mark Gerstein, Tao Qin
  • 0

Abstract

Equivariant Graph Neural Networks (EGNNs) have demonstrated significantsuccess in modeling microscale systems, including those in chemistry, biologyand materials science. However, EGNNs face substantial computational challengesdue to the high cost of constructing edge features via spherical tensorproducts, making them impractical for large-scale systems. To address thislimitation, we introduce E2Former, an equivariant and efficient transformerarchitecture that incorporates the Wigner $6j$ convolution (Wigner $6j$ Conv).By shifting the computational burden from edges to nodes, the Wigner $6j$ Convreduces the complexity from $O(|\mathcal{E}|)$ to $ O(| \mathcal{V}|)$ whilepreserving both the model's expressive power and rotational equivariance. Weshow that this approach achieves a 7x-30x speedup compared to conventional$\mathrm{SO}(3)$ convolutions. Furthermore, our empirical results demonstratethat the derived E2Former mitigates the computational challenges of existingapproaches without compromising the ability to capture detailed geometricinformation. This development could suggest a promising direction for scalableand efficient molecular modeling.

 

Quick Read (beta)

loading the full paper ...