BSA: Ball Sparse Attention for Large-scale Geometries
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912429932281856 |
|---|---|
| author | Brita, Catalin E. Nguyen, Hieu Chanchu, Lohithsai Yadala Nagy, Domonkos Zhdanov, Maksim |
| author_facet | Brita, Catalin E. Nguyen, Hieu Chanchu, Lohithsai Yadala Nagy, Domonkos Zhdanov, Maksim |
| contents | Self-attention scales quadratically with input size, limiting its use for large-scale physical systems. Although sparse attention mechanisms provide a viable alternative, they are primarily designed for regular structures such as text or images, making them inapplicable for irregular geometries. In this work, we present Ball Sparse Attention (BSA), which adapts Native Sparse Attention (NSA) (Yuan et al., 2025) to unordered point sets by imposing regularity using the Ball Tree structure from the Erwin Transformer (Zhdanov et al., 2025). We modify NSA's components to work with ball-based neighborhoods, yielding a global receptive field at sub-quadratic cost. On an airflow pressure prediction task, we achieve accuracy comparable to Full Attention while significantly reducing the theoretical computational complexity. Our implementation is available at https://github.com/britacatalin/bsa. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_12541 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | BSA: Ball Sparse Attention for Large-scale Geometries Brita, Catalin E. Nguyen, Hieu Chanchu, Lohithsai Yadala Nagy, Domonkos Zhdanov, Maksim Machine Learning Artificial Intelligence Computer Vision and Pattern Recognition Self-attention scales quadratically with input size, limiting its use for large-scale physical systems. Although sparse attention mechanisms provide a viable alternative, they are primarily designed for regular structures such as text or images, making them inapplicable for irregular geometries. In this work, we present Ball Sparse Attention (BSA), which adapts Native Sparse Attention (NSA) (Yuan et al., 2025) to unordered point sets by imposing regularity using the Ball Tree structure from the Erwin Transformer (Zhdanov et al., 2025). We modify NSA's components to work with ball-based neighborhoods, yielding a global receptive field at sub-quadratic cost. On an airflow pressure prediction task, we achieve accuracy comparable to Full Attention while significantly reducing the theoretical computational complexity. Our implementation is available at https://github.com/britacatalin/bsa. |
| title | BSA: Ball Sparse Attention for Large-scale Geometries |
| topic | Machine Learning Artificial Intelligence Computer Vision and Pattern Recognition |
| url | https://arxiv.org/abs/2506.12541 |