Research map / Architectures and efficiency / Attention and transformers
Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
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.
Erwin is a hierarchical transformer built on ball trees that yields linear-time ball attention with progressive coarsening/refinement and rotated-tree cross-ball connections, evaluated across cosmology, molecular dynamics, PDE benchmarks, ShapeNet-Car and EAGLE turbulence.
Problem
Transformers on large-scale irregular point clouds and meshes suffer quadratic attention cost, while long-range and multi-scale interactions must still be captured.
Approach
Organize attention via ball trees built as perfect binary trees with contiguous storage, so attention is computed independently within fixed-size balls at multiple tree levels; alternate original and rotated trees for cross-ball interaction and use tree coarsening/refinement in a UNet-style encoder-decoder.
Key results
- On EAGLE, Erwin achieves V RMSE 0.044 and P RMSE 0.31 at +1 horizon, beating EAGLE (0.053/0.46), with 11 ms inference vs 30 ms for EAGLE.
- Ball tree attention runtime scales with fitted exponent beta=1.054 (R^2=0.999), with ball tree construction under 5% of total runtime.
- Erwin achieves state-of-the-art on 2 of 4 PDE benchmark tasks, but underperforms markedly on Airfoil.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase ball size from 32 to 256 (ablation) | Erwin-S · Cosmology | test loss: 0.620 -> 0.595 | Helped |
| Increase ball size from 32 to 256 (ablation) | Erwin-S · Cosmology | runtime (ms): 126.0 -> 229.6 | Hurt |
| Add MPNN embedding (ablation) | Erwin · MD | test loss: 0.738 -> 0.720 | Helped |
| Add MPNN embedding (ablation) | Erwin · ShapeNet-Car | test loss: 30.39 -> 30.49 | Hurt |
| Add relative position embedding (ablation) | Erwin with MPNN embedding · MD | test loss: 0.720 -> 0.715 | Helped |
| Add relative position embedding (ablation) | Erwin with MPNN embedding · ShapeNet-Car | test loss: 30.49 -> 30.02 | Helped |
| Add rotating-tree cross-ball connections (ablation) | Erwin with MPNN and RPE · MD | test loss: 0.715 -> 0.712 | Helped |
| Add rotating-tree cross-ball connections (ablation) | Erwin with MPNN and RPE · ShapeNet-Car | test loss: 30.02 -> 15.85 | Helped |
| Use Erwin instead of EAGLE | EAGLE benchmark · EAGLE +1 | velocity RMSE: 0.053 -> 0.044 | Helped |
| Use Erwin instead of EAGLE | EAGLE benchmark · EAGLE +1 | pressure RMSE: 0.46 -> 0.31 | Helped |
| Use Erwin instead of EAGLE | EAGLE benchmark · EAGLE +50 | velocity RMSE: 0.349 -> 0.281 | Helped |
| Use Erwin instead of EAGLE | EAGLE benchmark · EAGLE inference | runtime (ms): 30 -> 11 | Helped |
Surprising findings
- On the Airfoil PDE task Erwin dramatically underperforms, attributed to varying mesh density across balls.
- On the MD task, model choice has minimal impact on test loss despite large runtime differences.
Limitations
- Perfect binary tree padding with virtual nodes adds overhead for non-coarsened ball attention layers.
- Erwin is neither permutation nor rotation equivariant.
- Cross-ball interaction needs multiple steps for signals to propagate, and coarsening sacrifices fine-grained detail.
Experimental setup
Datasets: Cosmology Molecular dynamics polymer (coarse-grained) Elasticity Plasticity Airfoil Pipe ShapeNet-Car EAGLE
Benchmarks: Cosmology Molecular dynamics Elasticity Plasticity Airfoil Pipe ShapeNet-Car EAGLE
Metrics: MSE relative L2 error negative log-likelihood RMSE runtime memory
Baselines: MPNN NequIP SEGNN PTv3 PointNet++ MeshGraphNet GAT DilResNet EAGLE Transolver UPT GINO
Builds on: Swin Transformer Transformer PointTransformer v3 MPNN UNet
Closest papers by meaning
- SEA: State-Exchange Attention for High-Fidelity Physics Based Transformers (NeurIPS 2024 · similarity 0.85)
- PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations (ICML 2025 · similarity 0.85)
- Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators (NeurIPS 2024 · similarity 0.84)
- Locality-Sensitive Hashing-Based Efficient Point Transformer with Applications in High-Energy Physics (ICML 2024 · similarity 0.84)
- Attention on the Sphere (NeurIPS 2025 · similarity 0.83)
- SPARTAN: A Sparse Transformer World Model Attending to What Matters (NeurIPS 2025 · similarity 0.83)
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.