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

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 $\mathit{\text{every}}$ edge. To overcome this, we introduce $\mathbf{\text{E2Former-V2}}$, a scalable architecture that integrates algebraic sparsity with hardware-aware execution. We first propose $\mathbf{\text{E}}$quivariant $\mathbf{\text{A}}$xis-$\mathbf{\text{A}}$ligned $\mathbf{\text{S}}$parsification (EAAS). EAAS builds on Wigner-$6j$ convolution by exploiting an $\mathrm{S}\mathrm{O}(3)\to \mathrm{S}\mathrm{O}(2)$ change of basis to transform computationally expensive dense tensor contractions into efficient, sparse parity re-indexing operations. Building on this representation, we introduce $\mathbf{\text{On-the-Fly}}\mathbf{\text{Equivariant}}\mathbf{\text{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 $\mathbf{\text{20}}\times \mathbf{\text{improvement}}\mathbf{\text{in}}\mathbf{\text{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.