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 efficiency 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}, an efficient architecture that integrates algebraic sparsity with hardware-aware execution. We introduce \textbf{Equivariant Axis-Aligned Sparsification (EAAS)}, which leverages an $\mathrm{SO}(3)\!\rightarrow\!\mathrm{SO}(2)$ change of basis to convert dense Wigner-$6j$ tensor contractions into sparse, parity-based re-indexing operations. Building on this representation, we propose \textbf{On-the-Fly Equivariant Attention}, a fully node-centric mechanism implemented via a fused Triton kernel. By eliminating materialized edge tensors and maximizing SRAM utilization, our kernel achieves up to \textbf{20$\times$ higher TFLOPS} than standard implementations. Experiments on SPICE and OMol25 show that E2Former-V2 preserves predictive accuracy while substantially accelerating inference, demonstrating the practical feasibility of large equivariant transformers on commodity GPUs. Our released code can be found at \url{https://github.com/IQuestLab/UBio-MolFM/tree/main}

Lay Summary Equivariant Graph Neural Networks (EGNNs) are powerful AI models for simulating molecules and materials, but existing methods are often computationally expensive because they build large geometric tensors for every atomic interaction. This limits their scalability on modern GPUs. We introduce E2Former-V2, a new equivariant transformer architecture that combines mathematical sparsity with hardware-aware GPU execution. Our method, called Equivariant Axis-Aligned Sparsification (EAAS), converts expensive geometric tensor operations into lightweight sparse indexing operations. We further design an On-the-Fly Equivariant Attention mechanism implemented with fused Triton GPU kernels, avoiding costly intermediate tensor materialization. Together, these optimizations achieve up to 20× higher effective TFLOPS than standard implementations while preserving predictive accuracy on the SPICE and OMol25 benchmarks. Our results show that large equivariant transformers can become significantly more practical and efficient on commodity GPUs. Our released code can be found at \url{https://github.com/IQuestLab/UBio-MolFM/tree/main}