# Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 65 references that are in the map.
- First posted on arXiv 24 Feb 2025.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Point Transformer V3: Simpler, Faster, Stronger](https://mlsanity.com/paper/2312.10035/) (2023)
- [Fourier Neural Operator with Learned Deformations for PDEs on General Geometries](https://mlsanity.com/paper/2207.05209/) (2022)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [GNOT: A General Neural Operator Transformer for Operator Learning](https://mlsanity.com/paper/2302.14376/) (2023)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [Geometric and Physical Quantities Improve E(3) Equivariant Message Passing](https://mlsanity.com/paper/2110.02905/) (2021)
- [Scalable Transformer for PDE Surrogate Modeling](https://mlsanity.com/paper/2305.17560/) (2023)
- [OctFormer: Octree-based Transformers for 3D Point Clouds](https://mlsanity.com/paper/2305.03045/) (2023)
- [Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries](https://mlsanity.com/paper/2502.02414/) (2025)
- [SWFormer: Sparse Window Transformer for 3D Object Detection in Point Clouds](https://mlsanity.com/paper/2210.07372/) (2022)
- [Learned Coarse Models for Efficient Turbulence Simulation](https://mlsanity.com/paper/2112.15275/) (2021)
- [Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations](https://mlsanity.com/paper/2407.07218/) (2024)
- [Latent Neural Operator for Solving Forward and Inverse PDE Problems](https://mlsanity.com/paper/2406.03923/) (2024)
- [FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer](https://mlsanity.com/paper/2301.08739/) (2023)
- [Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics](https://mlsanity.com/paper/2302.00600/) (2023)
- [Point Cloud Mamba: Point Cloud Learning via State Space Model](https://mlsanity.com/paper/2403.00762/) (2024)
- [Improved Operator Learning by Orthogonal Attention](https://mlsanity.com/paper/2310.12487/) (2023)

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

- [Appa: Bending Weather Dynamics with Latent Diffusion Models for Global Data Assimilation](https://mlsanity.com/paper/2504.18720/) (2025)
- [BLIPs: Bayesian Learned Interatomic Potentials](https://mlsanity.com/paper/2508.14022/) (2025)
- [Linear Attention with Global Context: A Multipole Attention Mechanism for Vision and Physics](https://mlsanity.com/paper/2507.02748/) (2025)
- [(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models](https://mlsanity.com/paper/2604.16429/) (2026)
- [Geometric Hyena Networks for Large-scale Equivariant Learning](https://mlsanity.com/paper/2505.22560/) (2025)
- [MSPT: Efficient Large-Scale Physical Modeling via Parallelized Multi-Scale Attention](https://mlsanity.com/paper/2512.01738/) (2025)
- [Point-wise Diffusion Models for Physical Systems with Shape Variations: Application to Spatio-temporal and Large-scale system](https://mlsanity.com/paper/2508.01230/) (2025)
- [BSA: Ball Sparse Attention for Large-scale Geometries](https://mlsanity.com/paper/2506.12541/) (2025)
- [CORDS: Continuous Representations of Discrete Structures](https://mlsanity.com/paper/2601.21583/) (2026)
- [Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasets](https://mlsanity.com/paper/2508.10758/) (2025)
- [Conditional Clifford-Steerable CNNs for PDE Modeling](https://mlsanity.com/paper/2510.14007/) (2025)
- [PIBNet: a Physics-Inspired Boundary Network for Multiple Scattering Simulations](https://mlsanity.com/paper/2512.02049/) (2025)
- [When Does Equivariance Help? Canonical Alignment in Neural Fluid Surrogates](https://mlsanity.com/paper/2605.18816/) (2026)
- [An Exploratory Study into using Machine-Learning for Fast Step-by-step Emulation of Numerical Mechanical Thrombectomy Simulations for Ischemic Stroke](https://mlsanity.com/paper/2606.00892/) (2026)
- [Revisiting Neural Processes via Fourier Transform and Volterra Series](https://mlsanity.com/paper/2606.01172/) (2026)
- [Higher-Order Fourier Neural Operator: Explicit Mode Mixer for Nonlinear PDEs](https://mlsanity.com/paper/2606.28122/) (2026)
- [Physics Transformer: Tailoring Transformer for General PDE Prediction](https://mlsanity.com/paper/2607.24513/) (2026)
- [Distance-Aware Attention and Wall-Distance Expert Routing for Transformer-Based 3D Flow Prediction](https://mlsanity.com/paper/2609.07222/) (2026)
- [ScaGNN: a Graph Neural Network for Multiple Scattering Simulations](https://mlsanity.com/paper/2609.37509/) (2026)
- [Geometry Meets Physics: Data-Efficient Pre-Training for Unstructured Neural PDE Solvers](https://mlsanity.com/paper/2610.03363/) (2026)
- [Cross-Domain Pretraining for Steady-State Neural CFD Surrogates](https://mlsanity.com/paper/2610.10398/) (2026)
- [NEMORA: Neural Equivariant Multipole Operators for Long-Range Atomistic Learning](https://mlsanity.com/paper/2610.10776/) (2026)

## Abstract (arXiv)

Large-scale physical systems defined on irregular grids pose significant scalability challenges for deep learning methods, especially in the presence of long-range interactions and multi-scale coupling. Traditional approaches that compute all pairwise interactions, such as attention, become computationally prohibitive as they scale quadratically with the number of nodes. We present Erwin, a hierarchical transformer inspired by methods from computational many-body physics, which combines the efficiency of tree-based algorithms with the expressivity of attention mechanisms. Erwin employs ball tree partitioning to organize computation, which enables linear-time attention by processing nodes in parallel within local neighborhoods of fixed size. Through progressive coarsening and refinement of the ball tree structure, complemented by a novel cross-ball interaction mechanism, it captures both fine-grained local details and global features. We demonstrate Erwin's effectiveness across multiple domains, including cosmology, molecular dynamics, PDE solving, and particle fluid dynamics, where it consistently outperforms baseline methods both in accuracy and computational efficiency.