# BSA: Ball Sparse Attention for Large-scale Geometries

arXiv: https://arxiv.org/abs/2506.12541 · HTML: https://mlsanity.com/paper/2506.12541/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 30 references that are in the map.
- First posted on arXiv 14 Jun 2025.

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Point Transformer V3: Simpler, Faster, Stronger](https://mlsanity.com/paper/2312.10035/) (2023)
- [GNOT: A General Neural Operator Transformer for Operator Learning](https://mlsanity.com/paper/2302.14376/) (2023)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [Solving High-Dimensional PDEs with Latent Spectral Models](https://mlsanity.com/paper/2301.12664/) (2023)
- [Latent Neural Operator for Solving Forward and Inverse PDE Problems](https://mlsanity.com/paper/2406.03923/) (2024)
- [FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer](https://mlsanity.com/paper/2301.08739/) (2023)
- [Improved Operator Learning by Orthogonal Attention](https://mlsanity.com/paper/2310.12487/) (2023)
- [Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators](https://mlsanity.com/paper/2402.12365/) (2024)
- [Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems](https://mlsanity.com/paper/2502.17019/) (2025)
- [Locality-Sensitive Hashing-Based Efficient Point Transformer with Applications in High-Energy Physics](https://mlsanity.com/paper/2402.12535/) (2024)
- [Transformers are efficient hierarchical chemical graph learners](https://mlsanity.com/paper/2310.01704/) (2023)
- [Resolution-Agnostic Transformer-based Climate Downscaling](https://mlsanity.com/paper/2411.14774/) (2024)

## Built on by (most cited first, 1 of 1)

- [(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models](https://mlsanity.com/paper/2604.16429/) (2026)

## Abstract (arXiv)

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.