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

TL;DR

E2Former-V2 tackles the scalability bottlenecks of equivariant GNNs by replacing edge-centric tensor materialization with a node-centric, hardware-aware design. It combines Equivariant Axis-Aligned Sparsification (EAAS), which reduces dense SO(3) contractions to sparse re-indexing in an SO(2) basis, with a fused On-the-Fly Equivariant Attention kernel that streams neighbor interactions without materializing edge tensors. The approach delivers roughly a 20× improvement in TFLOPS and maintains comparable predictive performance on SPICE and OMol25 benchmarks, enabling large-scale molecular modeling on standard GPUs. This work demonstrates that large equivariant transformers can be trained efficiently with practical hardware by rethinking arithmetic and memory layouts around graph edges and nodes.

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.

Figures (5)

• Figure 1: Latency vs. number of atoms ($N$) for Traditional EGNNs, FlashAttention, and Ours. FlashAttention consistently improves over Traditional EGNNs across all system sizes. The advantage of our method becomes increasingly pronounced as $N$ grows. See Appendix \ref{['app:latency']} for detailed experimental settings.
• Figure 2: Key components of E2Former-V2. (a) EAAS. E2Former-V2 aligns features with $D_R$, applies the sparse EAAS re-indexing operator $\mathcal{P}$ (Eq. 13) in the axis-aligned frame, and inverse-aligns with $D_{R^{-1}}$. The visualization shows the re-indexing pattern for $\ell_i=1$ and $\ell_f=1$. (b) On-the-fly equivariant attention. E2Former-V2 computes attention by streaming over neighbors and accumulating the output on the fly, avoiding explicit materialization of edge-level intermediates.
• Figure 3: Forward pass time comparison between our EAAS SO(2)-based tensor product and e3nn's SO(3) tensor product. Benchmarked with $\ell_{\max}=2$ and 128 channels across varying numbers of tensor product operations.
• Figure 4: Performance benchmarks of our on-the-fly equivariant attention kernels on H20 GPU.(a)(b) Computational throughput (TFLOPS) as a function of the number of neighbors $K$ and atoms $N$, respectively. (c)(d) Peak GPU memory usage (GB) as a function of the number of neighbors $K$ and atoms $N$, respectively.
• Figure 5: Oxygen-Oxygen radial distribution function (RDF) comparison. E2Former-V2 is compared against MACE-OFF and experimental data on bulk water.

Theorems & Definitions (3)

• Lemma 4.1: Pole Sparsity of Solid Spherical Harmonics
• Definition 4.2: EAAS Re-indexing Operator
• Proposition 4.3: Equivariant Axis-Aligned Sparsification