English

E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products

Machine Learning 2025-11-04 v4 Materials Science

Abstract

Equivariant Graph Neural Networks (EGNNs) have demonstrated significant success in modeling microscale systems, including those in chemistry, biology and materials science. However, EGNNs face substantial computational challenges due to the high cost of constructing edge features via spherical tensor products, making them impractical for large-scale systems. To address this limitation, we introduce E2Former, an equivariant and efficient transformer architecture that incorporates the Wigner 6j6j convolution (Wigner 6j6j Conv). By shifting the computational burden from edges to nodes, the Wigner 6j6j Conv reduces the complexity from O(E)O(|\mathcal{E}|) to O(V) O(| \mathcal{V}|) while preserving both the model's expressive power and rotational equivariance. We show that this approach achieves a 7x-30x speedup compared to conventional SO(3)\mathrm{SO}(3) convolutions. Furthermore, our empirical results demonstrate that the derived E2Former mitigates the computational challenges of existing approaches without compromising the ability to capture detailed geometric information. This development could suggest a promising direction for scalable and efficient molecular modeling.

Keywords

Cite

@article{arxiv.2501.19216,
  title  = {E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products},
  author = {Yunyang Li and Lin Huang and Zhihao Ding and Chu Wang and Xinran Wei and Han Yang and Zun Wang and Chang Liu and Yu Shi and Peiran Jin and Tao Qin and Mark Gerstein and Jia Zhang},
  journal= {arXiv preprint arXiv:2501.19216},
  year   = {2025}
}
R2 v1 2026-06-28T21:27:47.721Z