用树结构加速物理系统建模,兼顾精度与效率
Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
- 基于球树分层组织计算,实现线性时间注意力
- 在宇宙学、分子动力学等任务中精度和速度均超越基线
- 适合大规模不规则网格物理系统的高效建模
定义在不规则网格上的大规模物理系统给深度学习带来显著可扩展性挑战,尤其在长程相互作用和多尺度耦合情况下。传统方法如注意力机制需计算所有节点对之间的交互,随节点数呈二次增长,计算成本过高。我们提出Erwin,一种受计算多体物理方法启发的层次化变压器,结合树状算法的高效性与注意力机制的表达力。Erwin采用球树划分组织计算,通过在固定大小的局部邻域内并行处理节点,实现线性时间注意力。通过球树结构的逐步粗化与细化,辅以新颖的跨球交互机制,有效捕捉细粒度局部细节与全局特征。我们在宇宙学、分子动力学、偏微分方程求解及粒子流体动力学等多个领域验证了Erwin的有效性,其在准确性和计算效率上持续优于基线方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。