Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems (2025-02-24T00:00:00.000000Z)

TL;DR

Erwin is 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 and demonstrates Erwin's effectiveness across multiple domains, including cosmology, molecular dynamics, PDE solving, and particle fluid dynamics.

Abstract

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.

Authors

Maksim Zhdanov

1 papers

Max Welling

1 papers

Jan-Willem van de Meent

1 papers

TL;DR

Abstract

Authors

References65 items

Computational reconstruction of mental representations using human behavior

Geometric Hyena Networks for Large-scale Equivariant Learning

Learning Distributions of Complex Fluid Simulations with Diffusion Graph Networks

AB-UPT: Scaling Neural CFD Surrogates for High- Fidelity Automotive Aerodynamics Simulations via Anchored- Branched Universal Physics Transformers

Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries

Euclidean Fast Attention -- Machine Learning Global Atomic Representations at Linear Cost

Probabilistic weather forecasting with machine learning

NeuralDEM - Real-time Simulation of Industrial Particulate Flows

A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processing

Generative Modeling of Molecular Dynamics Trajectories

Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations

Grounding Continuous Representations in Geometry: Equivariant Neural Fields

Latent Neural Operator for Solving Forward and Inverse PDE Problems

Point Cloud Mamba: Point Cloud Learning via State Space Model

Transolver: A Fast Transformer Solver for PDEs on General Geometries

Point Transformer V3: Simpler, Faster, Stronger

Learning skillful medium-range global weather forecasting

Improved Operator Learning by Orthogonal Attention

MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design

Geometry-Informed Neural Operator for Large-Scale 3D PDEs

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Geometric Algebra Transformer

Scalable Transformer for PDE Surrogate Modeling

OctFormer: Octree-based Transformers for 3D Point Clouds

GNOT: A General Neural Operator Transformer for Operator Learning

Eagle: Large-Scale Learning of Turbulent Fluid Dynamics with Mesh Transformers

Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics

FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer

SWFormer: Sparse Window Transformer for 3D Object Detection in Point Clouds

Efficient Learning of Mesh-Based Physical Simulation with Bi-Stride Multi-Scale Graph Neural Network

Fourier Neural Operator with Learned Deformations for PDEs on General Geometries

Transformer for Partial Differential Equations' Operator Learning

Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-scale Graph Networks

Learned Coarse Models for Efficient Turbulence Simulation

Geometric and Physical Quantities improve E(3) Equivariant Message Passing

H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequences

Choose a Transformer: Fourier or Galerkin

E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials

Fourier Neural Operator for Parametric Partial Differential Equations

Multiscale Forecasting of High-Impact Weather: Current Status and Future Challenges

Learning Mesh-Based Simulation with Graph Networks

Targeted sequence design within the coarse-grained polymer genome

GLU Variants Improve Transformer

Learning three-dimensional flow for interactive aerodynamic design

Decoupled Weight Decay Regularization

Graph Attention Networks

Attention is All you Need

PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space

Neural Message Passing for Quantum Chemistry

Group Equivariant Convolutional Networks

Scikit-learn: Machine Learning in Python

Many-body tree methods in physics

A hierarchical O(N log N) force-calculation algorithm

Aurora: A Foundation Model of the Atmosphere

Fast Multipole Attention: A Divide-and-Conquer Attention Mechanism for Long Sequences

Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation

HT-Net: Hierarchical Transformer based Operator Learning Model for Multiscale PDEs

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

U-net: Convolu-tional networks for biomedical image segmentation

Computer Simulation Using Particles

A FAST ADAPTIVE MULTIPOLE ALGORITHM FOR PARTICLE SIMULATIONS *

Bridging operator learning and conditioned neural fields: A unifying perspective

Universal physics transformers: A framework for efficiently scaling neural operators

Erwin Transformer

Transferable

Field of Study

Journal Information

Name

Volume

Venue Information

Name

Type

URL

Alternate Names