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

arXiv:2601.16622v2 Announce Type: replace-cross
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 textitevery edge. To overcome this, we introduce textbfE2Former-V2, a scalable architecture that integrates algebraic sparsity with hardware-aware execution. We first propose textbfEquivariant textbfAxis-textbfAligned textbfSparsification (EAAS). EAAS builds on Wigner-$6j$ convolution by exploiting an $mathrmSO(3) rightarrow mathrmSO(2)$ change of basis to transform computationally expensive dense tensor contractions into efficient, sparse parity re-indexing operations. Building on this representation, we introduce textbfOn-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 textbf20$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.