BSA: Ball Sparse Attention for Large-scale Geometries

Catalin E. Brita, Hieu Nguyen, Lohithsai Yadala Chanchu, Domonkos Nagy, Maksim Zhdanov
6/14/2025
cs.LGcs.AIcs.CV

Official code: britacatalin/bsa (13 stars, MIT, last commit 2025-06-10).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

britacatalin/bsaOfficial100%
MIT
130C++, ShellMay 31, 20251 years agoMIT

Similar Papers

Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasets

85% match1 repo32025

SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space

81% match1 repo322025

Re-ttention: Ultra Sparse Visual Generation via Attention Statistical Reshape

81% match1 repo72025

Towards Understanding the Nature of Attention with Low-Rank Sparse Decomposition

81% match1 repo302025

SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference

80% match2026

SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference

80% match1 repo1.0k2025

Cite this paper

@article{brita2025bsa,
  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},
  year   = {2025},
  eprint = {2506.12541},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2506.12541v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2506.12541.svg)](https://codeofpaper.com/paper/2506.12541)

Discussion