mlsanityData to 7 Oct 2026

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

arXiv 2502.17019 ↗First posted 24 Feb 2025cs.LG cs.AI cs.CVcited by 21

Not on the map. Why?

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.

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

Builds on

42 of 65 references.

References · 42

  1. 2017Attention Is All You Needcited by 59,932
  2. 2021Swin Transformer: Hierarchical Vision Transformer using Shifted Windowscited by 12,082
  3. 2017Graph Attention Networkscited by 8,854
  4. 2017PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Spacecited by 4,384
  5. 2017Neural Message Passing for Quantum Chemistrycited by 3,460
  6. 2020Fourier Neural Operator for Parametric Partial Differential Equationscited by 2,806
  7. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  8. 2020GLU Variants Improve Transformercited by 1,724
  9. 2016Group Equivariant Convolutional Networkscited by 1,154
  10. 2020Learning Mesh-Based Simulation with Graph Networkscited by 840
  11. 2021E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentialscited by 701
  12. 2023Point Transformer V3: Simpler, Faster, Strongercited by 579
30 more
  1. 2022Fourier Neural Operator with Learned Deformations for PDEs on General Geometriescited by 416
  2. 2021Choose a Transformer: Fourier or Galerkincited by 333
  3. 2023GNOT: A General Neural Operator Transformer for Operator Learningcited by 321
  4. 2024Transolver: A Fast Transformer Solver for PDEs on General Geometriescited by 299
  5. 2023Geometry-Informed Neural Operator for Large-Scale 3D PDEscited by 293
  6. 2022Transformer for Partial Differential Equations' Operator Learningcited by 270
  7. 2021Geometric and Physical Quantities Improve E(3) Equivariant Message Passingcited by 225
  8. 2023Scalable Transformer for PDE Surrogate Modelingcited by 123
  9. 2023OctFormer: Octree-based Transformers for 3D Point Cloudscited by 105
  10. 2022SWFormer: Sparse Window Transformer for 3D Object Detection in Point Cloudscited by 92
  11. 2021Learned Coarse Models for Efficient Turbulence Simulationcited by 90
  12. 2025Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometriescited by 89
  13. 2024Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equationscited by 87
  14. 2024Latent Neural Operator for Solving Forward and Inverse PDE Problemscited by 76
  15. 2023FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformercited by 72
  16. 2023Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamicscited by 72
  17. 2024Point Cloud Mamba: Point Cloud Learning via State Space Modelcited by 66
  18. 2023Improved Operator Learning by Orthogonal Attentioncited by 58
  19. 2023Eagle: Large-Scale Learning of Turbulent Fluid Dynamics with Mesh Transformerscited by 54
  20. 2025AB-UPT: Scaling Neural CFD Surrogates for High-Fidelity Automotive Aerodynamics Simulations via Anchored-Branched Universal Physics Transformerscited by 49
  21. 2024Generative Modeling of Molecular Dynamics Trajectoriescited by 48
  22. 2021H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequencescited by 47
  23. 2023Geometric Algebra Transformercited by 46
  24. 2022Efficient Learning of Mesh-Based Physical Simulation with BSMS-GNNcited by 43
  25. 2022Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-Scale Graph Networkscited by 23
  26. 2023MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Designcited by 23
  27. 2024Grounding Continuous Representations in Geometry: Equivariant Neural Fieldscited by 19
  28. 2024Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Costcited by 12
  29. 2024NeuralDEM -- Real-time Simulation of Industrial Particulate Flowscited by 10
  30. 2024A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processingcited by 5

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 10; 2026: 11.

21 papers.

Sort:

By area:

  1. 2025Appa: Bending Weather Dynamics with Latent Diffusion Models for Global Data Assimilationcited by 17
  2. 2025BLIPs: Bayesian Learned Interatomic Potentialscited by 6
  3. 2025Linear Attention with Global Context: A Multipole Attention Mechanism for Vision and Physicscited by 5
  4. 2025Geometric Hyena Networks for Large-scale Equivariant Learningcited by 4
  5. 2026(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Modelscited by 4
  6. 2025MSPT: Efficient Large-Scale Physical Modeling via Parallelized Multi-Scale Attentioncited by 3
  7. 2025Point-wise Diffusion Models for Physical Systems with Shape Variations: Application to Spatio-temporal and Large-scale systemcited by 2
  8. 2025BSA: Ball Sparse Attention for Large-scale Geometriescited by 1
  9. 2026CORDS: Continuous Representations of Discrete Structurescited by 1
  10. 2025Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasetscited by 0
  11. 2025Conditional Clifford-Steerable CNNs for PDE Modelingcited by 0
  12. 2025PIBNet: a Physics-Inspired Boundary Network for Multiple Scattering Simulationscited by 0
9 more
  1. 2026When Does Equivariance Help? Canonical Alignment in Neural Fluid Surrogatescited by 0
  2. 2026An Exploratory Study into using Machine-Learning for Fast Step-by-step Emulation of Numerical Mechanical Thrombectomy Simulations for Ischemic Strokecited by 0
  3. 2026Revisiting Neural Processes via Fourier Transform and Volterra Seriescited by 0
  4. 2026Higher-Order Fourier Neural Operator: Explicit Mode Mixer for Nonlinear PDEscited by 0
  5. 2026Physics Transformer: Tailoring Transformer for General PDE Predictioncited by 0
  6. 2026Distance-Aware Attention and Wall-Distance Expert Routing for Transformer-Based 3D Flow Predictioncited by 0
  7. 2026ScaGNN: a Graph Neural Network for Multiple Scattering Simulationscited by 0
  8. 2026Cross-Domain Pretraining for Steady-State Neural CFD Surrogatescited by 0
  9. 2026NEMORA: Neural Equivariant Multipole Operators for Long-Range Atomistic Learningcited by 0