Research map / Architectures and efficiency / RNNs and normalization
E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the Wigner 6j convolution, which factorizes spherical-harmonic tensor products into node-local terms to cut SO(3) convolution cost from edge-scaling to node-scaling, and builds the E2Former transformer on it.
Problem
Spherical equivariant GNNs (SO(3)/SE(3) transformers) are expensive because edge-wise spherical tensor products scale with the number of edges and angular cutoff L, limiting use on large molecular systems.
Approach
Use a binomial local expansion of edge spherical harmonics R^(l)(r_ij) into node-local terms, then apply Wigner 6j recoupling so that the tensor products are performed per node rather than per edge, yielding O(|V|) tensor-product count while keeping exact equivariance and expressivity.
Key results
- Wigner 6j convolution gives about 7x speed-up over SO(3) convolution on 1000-node graphs at L_max=6, and 7x–30x overall per the abstract.
- On OC20-2M, E2Former 67M reaches 270 meV energy MAE and 20.50 meV/Å force MAE, while E2Former 33M is the fastest model at 62 samples/sec with 800 GPU hours.
- On the 6400-atom system, E2Former is the only tested model that runs (80 GB), where MACE-Large and EquiformerV2 fail.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use E2Former 67M instead of EquiformerV2 | OC20-2M · OC20-2M validation | Energy MAE (meV): 285 -> 270 | Helped |
| Use E2Former 67M instead of EquiformerV2 | OC20-2M · OC20-2M validation | Force MAE (meV/Å): 20.46 -> 20.50 | Hurt |
| Use E2Former 67M instead of EquiformerV2 | OC22 S2EF-Total · ID validation | Energy MAE (meV): 433 -> 491 | Hurt |
| Use E2Former 67M instead of EquiformerV2 | OC22 S2EF-Total · OOD validation | Energy MAE (meV): 629 -> 724 | Hurt |
| Use E2Former 67M instead of EquiformerV2 | OC22 S2EF-Total · ID validation | Force MAE (meV/Å): 22.88 -> 25.98 | Hurt |
| Use E2Former 67M instead of EquiformerV2 | OC22 S2EF-Total · OOD validation | Force MAE (meV/Å): 30.70 -> 36.45 | Hurt |
| Use E2Former 33M instead of MACE Large | SPICE · SPICE test, All | Energy MAE (meV/atom): 0.77 -> 0.60 | Helped |
| Use E2Former 33M instead of MACE Large | SPICE · SPICE test, All | Force MAE (meV/Å): 12.26 -> 7.46 | Helped |
| Use Wigner 6j convolution instead of SO(3) convolution | Convolution runtime · 1000-node graphs, angular momentum cutoff up to L_max = 6 | speedup: approximately a 7× speed-up | Helped |
| Use E2Former-33M instead of MACE-Large-33M | Large-scale system training · 3200 atoms | Training efficiency (samples/second): 0.41 -> 1.2 | Helped |
| Scale E2Former-33M to 6400 atoms | Large-scale system training · 6400 atoms | Training efficiency (samples/second) | Helped |
Surprising findings
- E2Former is not strictly linear in input size despite the claimed O(|V|) tensor-product scaling, since attention remains edge-level.
Limitations
- Architecture is not strictly linear because attention weights remain edge-dependent.
- OC20 experiments use the 2M subset and results are taken from prior publications rather than re-run, limiting direct comparison.
Experimental setup
Datasets: OC20 OC22 SPICE QMugs DeShaw GEMS
Benchmarks: OC20-2M OC22 S2EF SPICE MD power spectrum (253-atom amino acid, 6000-atom water cluster)
Metrics: energy MAE (meV) force MAE (meV/Å) inference speed (samples/sec) training GPU hours memory cost (GB) runtime speed-up
Baselines: GemNet-dT GemNet-OC SCN eSCN EquiformerV2 MACE-Large GFN2-xTB MADFT
Builds on: EquiformerV2 MACE eSCN Clebsch-Gordan tensor product
Closest papers by meaning
- Equivariant Spherical Transformer for Efficient Molecular Modeling (NeurIPS 2026 · similarity 0.92)
- EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations (ICLR 2024 · similarity 0.90)
- GeoMFormer: A General Architecture for Geometric Molecular Representation Learning (ICML 2024 · similarity 0.89)
- Equivariant Masked Position Prediction for Efficient Molecular Representation (ICLR 2025 · similarity 0.86)
- Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning (ICML 2024 · similarity 0.86)
- Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.