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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Li, Yunyang, Huang, Lin, Ding, Zhihao, Wang, Chu, Wei, Xinran, Yang, Han, Wang, Zun, Liu, Chang, Shi, Yu, Jin, Peiran, Qin, Tao, Gerstein, Mark, Zhang, Jia
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866917055882592256
author Li, Yunyang
Huang, Lin
Ding, Zhihao
Wang, Chu
Wei, Xinran
Yang, Han
Wang, Zun
Liu, Chang
Shi, Yu
Jin, Peiran
Qin, Tao
Gerstein, Mark
Zhang, Jia
author_facet Li, Yunyang
Huang, Lin
Ding, Zhihao
Wang, Chu
Wei, Xinran
Yang, Han
Wang, Zun
Liu, Chang
Shi, Yu
Jin, Peiran
Qin, Tao
Gerstein, Mark
Zhang, Jia
contents 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 $6j$ convolution (Wigner $6j$ Conv). By shifting the computational burden from edges to nodes, the Wigner $6j$ Conv reduces the complexity from $O(|\mathcal{E}|)$ to $ 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 $\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.
format Preprint
id arxiv_https___arxiv_org_abs_2501_19216
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products
Li, Yunyang
Huang, Lin
Ding, Zhihao
Wang, Chu
Wei, Xinran
Yang, Han
Wang, Zun
Liu, Chang
Shi, Yu
Jin, Peiran
Qin, Tao
Gerstein, Mark
Zhang, Jia
Machine Learning
Materials Science
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 $6j$ convolution (Wigner $6j$ Conv). By shifting the computational burden from edges to nodes, the Wigner $6j$ Conv reduces the complexity from $O(|\mathcal{E}|)$ to $ 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 $\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.
title E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products
topic Machine Learning
Materials Science
url https://arxiv.org/abs/2501.19216