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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Zhdanov, Maksim, Welling, Max, van de Meent, Jan-Willem
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866912407836688384
author Zhdanov, Maksim
Welling, Max
van de Meent, Jan-Willem
author_facet Zhdanov, Maksim
Welling, Max
van de Meent, Jan-Willem
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_17019
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
Zhdanov, Maksim
Welling, Max
van de Meent, Jan-Willem
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
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.
title Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2502.17019