E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

TL;DR

E2Former-V2 introduces a scalable, efficient equivariant graph neural network architecture that leverages algebraic sparsity and custom GPU kernels to improve computational efficiency while maintaining accuracy in 3D molecular modeling.

Contribution

The paper presents E2Former-V2, a novel architecture combining algebraic sparsity with hardware-aware execution, enabling scalable and efficient equivariant attention for 3D atomistic systems.

Findings

Achieves 20× TFLOPS improvement with custom kernels.

Maintains comparable predictive performance on molecular datasets.

Demonstrates scalable training of large equivariant transformers.

Abstract

Equivariant Graph Neural Networks (EGNNs) have become a widely used approach for modeling 3D atomistic systems. However, mainstream architectures face critical scalability bottlenecks due to the explicit construction of geometric features or dense tensor products on \textit{every} edge. To overcome this, we introduce \textbf{E2Former-V2}, a scalable architecture that integrates algebraic sparsity with hardware-aware execution. We first propose \textbf{E}quivariant \textbf{A}xis-\textbf{A}ligned \textbf{S}parsification (EAAS). EAAS builds on Wigner-$6j$ convolution by exploiting an $\mathrm{SO}(3) \rightarrow \mathrm{SO}(2)$ change of basis to transform computationally expensive dense tensor contractions into efficient, sparse parity re-indexing operations. Building on this representation, we introduce \textbf{On-the-Fly Equivariant Attention}, a fully node-centric mechanism implemented via a custom fused Triton kernel. By eliminating materialized edge tensors and maximizing SRAM utilization, our kernel achieves a \textbf{20$\times$ improvement in TFLOPS} compared to standard implementations. Extensive experiments on the SPICE and OMol25 datasets demonstrate that E2Former-V2 maintains comparable predictive performance while notably accelerating inference. This work demonstrates that large equivariant transformers can be trained efficiently using widely accessible GPU platforms. The code is avalible at https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2.