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

TL;DR

Erwin introduces a hierarchical, tree-based transformer that efficiently models large-scale physical systems with irregular grids, capturing multi-scale interactions while maintaining linear computational complexity.

Contribution

The paper presents Erwin, a novel hierarchical transformer combining tree-based algorithms with attention, enabling scalable and accurate modeling of complex physical systems.

Findings

Outperforms baseline methods in accuracy and efficiency

Effective across diverse domains like cosmology and molecular dynamics

Achieves linear-time attention through tree-based partitioning

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.