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

arXiv: https://arxiv.org/abs/2601.16622 · HTML: https://mlsanity.com/paper/2601.16622/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 43 references that are in the map.
- First posted on arXiv 23 Jan 2026.

## Builds on (most cited first, 24 of 24)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [UMA: A Family of Universal Models for Atoms](https://mlsanity.com/paper/2506.23971/) (2025)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [SE(3)-equivariant prediction of molecular wavefunctions and electronic densities](https://mlsanity.com/paper/2106.02347/) (2021)
- [The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains](https://mlsanity.com/paper/2410.24169/) (2024)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking](https://mlsanity.com/paper/1804.06826/) (2018)
- [An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks](https://mlsanity.com/paper/2501.13986/) (2025)
- [The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products](https://mlsanity.com/paper/2506.13523/) (2025)
- [E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products](https://mlsanity.com/paper/2501.19216/) (2025)
- [Tensor Decomposition Networks for Fast Machine Learning Interatomic Potential Computations](https://mlsanity.com/paper/2507.01131/) (2025)

## Built on by (most cited first, 3 of 3)

- [Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials](https://mlsanity.com/paper/2607.10664/) (2026)
- [UBio-MolFM: A Universal Molecular Foundation Model for Bio-Systems](https://mlsanity.com/paper/2602.17709/) (2026)
- [E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU](https://mlsanity.com/paper/2609.35099/) (2026)

## Abstract (arXiv)

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 E2Former-V2, a scalable architecture that integrates algebraic sparsity with hardware-aware execution. We first propose Equivariant Axis-Aligned Sparsification (EAAS). EAAS builds on Wigner-6j convolution by exploiting an SO(3) rightarrow 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 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 20× 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.