BSA：用于大规模几何图的球稀疏注意力
机器学习
2025-06-17 v1 人工智能
计算机视觉与模式识别
摘要
自注意力机制的计算复杂度随输入规模呈二次增长，这限制了其在大规模物理系统中的应用。虽然稀疏注意力机制提供了可行的替代方案，但它们主要为文本或图像等规则结构设计，难以适用于不规则几何图。本文提出球稀疏注意力（BSA），将本土稀疏注意力（NSA）（Yuan等，2025）适用于无序点集，通过引入Erwin Transformer（Zhdanov等，2025）中的球树结构来实现规则化。我们修改了NSA的组件以适用于基于球的邻域，实现了亚二次计算成本下的全局感受野。在气流压力预测任务上，我们在保持与完全注意力相当精度的同时，显著降低了理论计算复杂度。我们的实现已公开可用，链接为https://github.com/britacatalin/bsa。
引用
@article{arxiv.2506.12541,
title = {BSA: Ball Sparse Attention for Large-scale Geometries},
author = {Catalin E. Brita and Hieu Nguyen and Lohithsai Yadala Chanchu and Domonkos Nagy and Maksim Zhdanov},
journal= {arXiv preprint arXiv:2506.12541},
year = {2025}
}
备注
Long Context Foundation Models Workshop @ ICML 2025