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

要約

等変量グラフニューラルネットワーク（EGNN）は、化学、生物学、材料科学などのマイクロスケールシステムのモデリングにおいて大きな成功を収めている。しかし、EGNNは、球面テンソル積によってエッジ特徴を構成するコストが高いため、計算上の大きな課題に直面しており、大規模システムには実用的ではない。この限界に対処するため、我々は、ウィグナー$6j$畳み込み（Wigner $6j$ Conv）を組み込んだ、等変量で効率的な変換器アーキテクチャであるE2Formerを紹介する。Wigner $6j$ Convは、計算負担をエッジからノードにシフトすることで、モデルの表現力と回転等変数を保持しながら、複雑度を$O(|mathcal{E}|)$から$O(| \mathcal{V}|)$に低減する。このアプローチにより、従来の$mathmathrm{SO}(3)$畳み込みと比較して、7倍から30倍のスピードアップを達成することを示す。さらに、我々の実証結果は、詳細な幾何学的情報を捉える能力を損なうことなく、導出されたE2Formerが既存のアプローチの計算上の課題を軽減することを示している。この開発は、スケーラブルで効率的な分子モデリングの有望な方向性を示唆するものである。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL