🤖 AI

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

E2Former-V2 is a scalable equivariant graph neural network architecture that overcomes the computational bottlenecks of mainstream models by introducing Equivariant Axis-Aligned Sparsification and an On-the-Fly Equivariant Attention mechanism, achieving a 20-fold TFLOPS improvement while maintaining predictive performance on 3D atomistic systems.

Original authors: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Published 2026-06-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict how a massive crowd of people (atoms) will move and interact in a room. In the world of chemistry and physics, this is called "atomistic modeling." For a long time, scientists have used a special type of AI called an Equivariant Graph Neural Network (EGNN) to do this. These networks are great because they understand that if you rotate the whole room, the physics inside it shouldn't change.

However, the current versions of these networks have a major problem: they are clumsy and slow when the crowd gets big.

Here is a simple breakdown of the paper's solution, E2Former-V2, using everyday analogies.

The Problem: The "Edge-Centric" Traffic Jam

Imagine a party where everyone wants to talk to their neighbors.

Old Method (Traditional EGNNs): The system creates a giant, physical list of every single conversation happening at the party. If there are 1,000 people, and each talks to 50 neighbors, the system has to write down 50,000 separate conversation notes on a massive whiteboard (memory) before it can do any math.
The Bottleneck: As the party grows (more atoms), this whiteboard gets so huge that the computer runs out of space. It spends all its time just shuffling these notes around rather than actually thinking. This is why old models crash or slow down to a crawl when simulating large molecules like proteins.

The Solution: E2Former-V2

The authors built a new system that stops writing down every single conversation note. Instead, it uses two clever tricks to make the process lightning-fast.

Trick 1: The "Magic Compass" (Equivariant Axis-Aligned Sparsification)

In the old system, calculating how atoms interact is like trying to solve a puzzle where every piece can be rotated in any direction. It's messy and requires a lot of math.

The Analogy: Imagine you have a complex 3D puzzle. Instead of trying to solve it while it's spinning, you use a "Magic Compass" to rotate the whole puzzle so that one specific piece is always pointing straight up.
The Result: Once the puzzle is aligned, most of the pieces become empty or zero. You don't need to do math on the empty parts anymore. The paper calls this EAAS. It turns a messy, dense calculation into a simple "re-indexing" game (just moving pieces around), which is much faster.

Trick 2: The "Streaming Chef" (On-the-Fly Equivariant Attention)

This is the biggest breakthrough.

The Old Way (The Buffet): The old system prepares the entire buffet (all 50,000 conversation notes) and puts it on a table before anyone can eat. If the table is too small, the food spills, and the party stops.
The New Way (The Streaming Chef): The new system, E2Former-V2, acts like a chef who cooks one plate at a time, right in front of the customer, and serves it immediately.
- It doesn't write down the notes.
- It doesn't store the whole list.
- It calculates the interaction between two people, adds the result to the final score, and then immediately forgets the details to make room for the next pair.
The Result: This is called "On-the-Fly" processing. Because it never builds the giant whiteboard, it uses almost no memory. It fits everything into the computer's super-fast internal cache (SRAM), making it 20 times faster than the old methods.

What They Actually Achieved

The paper tested this new system on two main things:

Small Molecules (SPICE): They tested it on drug-like molecules. The new system was just as accurate as the best existing models but much faster.
Huge Molecules (OMol25): They tested it on massive chemical structures.
- The Limit: Old models crashed (ran out of memory) when trying to simulate systems with more than 10,000 atoms.
- The Win: E2Former-V2 successfully simulated systems with 100,000 atoms without crashing.

Real-World Tests Mentioned

The authors didn't just run numbers; they simulated real physics:

Liquid Water: They simulated a box of water molecules. The new model predicted how the water molecules arrange themselves (the "radial distribution function") almost perfectly, matching real-world experiments better than previous models.
Proteins: They simulated a large protein (about 30,000 atoms) floating in water. The model kept the protein stable and helped it fold into a shape that matched real experimental data, proving it can handle complex biological structures.

Summary

E2Former-V2 is like upgrading from a system that writes every single conversation on a giant, slow whiteboard to a system where a smart chef cooks and serves conversations one by one, instantly. This allows scientists to simulate massive molecular systems (like huge proteins) on standard computers without running out of memory, while keeping the physics accurate.

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

1. Problem Statement

Equivariant Graph Neural Networks (EGNNs) have become the standard for modeling 3D atomistic systems due to their ability to preserve rotational and translational symmetries. However, mainstream architectures (e.g., MACE, EquiformerV2, eSCN) suffer from critical efficiency bottlenecks rooted in their edge-centric design. These models explicitly construct geometric features or perform dense tensor products on every edge, leading to:

Memory Complexity: $O(|E|)$ activation memory, where $|E| \approx kN$ (with $k$ being the average neighbor count). For large systems, this results in massive memory footprints that exceed GPU capacity.
Hardware Inefficiency: The explicit materialization of edge tensors forces data movement between High Bandwidth Memory (HBM) and on-chip SRAM, creating a memory wall that limits throughput.
Scalability Limits: As system size ( $N$ ) increases, latency scales poorly, making large-scale molecular dynamics (MD) simulations computationally prohibitive on commodity hardware.

While recent advances in Transformers (e.g., FlashAttention) have solved similar bottlenecks for sequence data via streaming, tile-based execution, these principles have remained unexplored in SO(3)-equivariant architectures.

2. Methodology

E2Former-V2 addresses these issues by integrating algebraic sparsification with a hardware-aware, node-centric execution model. The architecture consists of two core innovations:

2.1. Equivariant Axis-Aligned Sparsification (EAAS)

Standard SO(3) tensor contractions involve dense summations over Clebsch–Gordan coefficients, resulting in $O(L^6)$ complexity and irregular memory access patterns. E2Former-V2 introduces EAAS to transform these dense operations into sparse, permutation-based operations:

SO(3) $\to$ SO(2) Change of Basis: The method aligns the global coordinate system with the local vector $\vec{r}$ (the axis connecting two atoms). In this aligned frame, solid spherical harmonics become maximally sparse (Lemma 4.1), where only the $m=0$ component is non-zero.
Re-indexing Operator: By commuting rotations through the tensor product, the dense contraction collapses into a deterministic re-indexing rule (Eq. 13). This operator maps input features to output features via sparse blockwise linear maps, eliminating the need for dense summation over magnetic indices.
Result: This reduces the convolution stage complexity and achieves a $\sim 1.5\times$ speedup over standard implementations.

2.2. On-the-Fly Equivariant Attention

To eliminate the $O(|E|)$ memory bottleneck, the authors propose a fully fused, node-centric attention mechanism implemented via a custom Triton kernel:

Node-Centric Factorization: Instead of materializing edge features, the attention mechanism decouples the computation. Query and Key vectors are constructed using SO(3)-invariant scalar projections (Eq. 6), while equivariant geometric information is propagated only through the value path.
Streaming Reduction: The kernel processes neighbors in a streaming fashion. It iterates over neighbors $j \in \mathcal{N}(i)$ , computing attention scores and accumulating value updates on-the-fly within SRAM.
Elimination of Edge Materialization: Intermediate tensors of shape $N \times K \times H \times D$ are never written to HBM. The kernel maintains running statistics (max, normalization, and value accumulators) on-chip, reducing activation memory complexity from $O(|E|)$ to $O(|V|)$ .
Online Softmax: To ensure numerical stability during streaming, an online softmax formulation is used to update accumulators incrementally.

3. Key Contributions

Equivariant Axis-Aligned Sparsification (EAAS): A novel algebraic technique that unifies Wigner-6j recoupling with SO(2) axis alignment, converting dense SO(3) tensor contractions into sparse re-indexing operations.
On-the-Fly Equivariant Attention: A fully fused Triton kernel that avoids edge materialization, maximizing SRAM reuse and achieving up to 20 $\times$ higher TFLOPS compared to standard implementations.
Scalable Performance: Demonstrated ability to handle systems with up to 100,000 atoms on a single GPU, a regime where previous equivariant Transformers fail due to Out-Of-Memory (OOM) errors.
Physical Fidelity: Validation in molecular dynamics showing accurate reproduction of liquid water structures and stable refinement of AlphaFold-predicted protein structures ( $\sim$ 30k atoms).

4. Experimental Results

The authors evaluated E2Former-V2 on the SPICE (medicinal chemistry) and OMol25 (foundation model scale) datasets, as well as in MD simulations.

Predictive Accuracy:
- On SPICE, E2Former-V2 achieved the lowest Energy and Force Mean Absolute Errors (MAE) across most subsets, outperforming MACE-Large by 48% in Energy MAE on Dimers.
- On OMol25, the conservative variant achieved an aggregate Energy MAE of 1.27 meV/atom, matching eSEN-small and significantly outperforming UMA-S (3.62 meV/atom).
- On OC20 S2EF-2M, it achieved accuracy comparable to strong baselines like EquiformerV2 and eSCN.
Efficiency and Scalability:
- Throughput: The fused kernel achieved an average 20 $\times$ speedup in TFLOPS over naive PyTorch implementations.
- Memory: Peak memory usage was reduced by 16.9 $\times$ to 19.5 $\times$ compared to baselines.
- Scalability: While prior models (e.g., EquiformerV2, MACE-Large) encountered OOM failures beyond 10k–50k atoms, E2Former-V2 scaled reliably to 100k atoms.
- Inference Speed: At 100k atoms, E2Former-V2 (Conservative) reached 0.12 steps/s, nearly 3 $\times$ faster than UMA-S. In the Direct setting, it was an order of magnitude faster than Allegro and EquiformerV2 at smaller scales.
Molecular Dynamics:
- Simulations of periodic bulk water (216 molecules) showed E2Former-V2 closely matching experimental radial distribution functions (RDF), outperforming MACE-OFF.
- A 2 ns simulation of a solvated protein ( $\sim$ 30k atoms) maintained high TM-scores relative to the native structure, demonstrating stable dynamics.

5. Significance and Claims

The paper claims that E2Former-V2 demonstrates the practical feasibility of large equivariant transformers on commodity GPUs. By shifting the execution paradigm from edge-centric to node-centric and leveraging hardware-aware streaming, the method resolves the fundamental memory bottleneck that has limited the scalability of equivariant models.

The authors emphasize that this is not merely an algorithmic improvement but a system-level realization that bridges the gap between theoretical factorization (reducing complexity from $O(|E|)$ to $O(|V|)$ ) and hardware efficiency. The work suggests that large-scale molecular foundation models and long-timescale MD simulations, previously computationally prohibitive, are now accessible using standard GPU infrastructure. The code is released to facilitate further adoption in the community.