BSA: Ball Sparse Attention for Large-scale Geometries
This work addresses a bottleneck for researchers and practitioners in computational physics or engineering dealing with irregular geometries, but it is incremental as it builds on existing sparse attention methods.
The paper tackles the problem of self-attention's quadratic scaling for large-scale physical systems by proposing Ball Sparse Attention (BSA), which adapts sparse attention to irregular geometries using a Ball Tree structure, achieving accuracy comparable to Full Attention on an airflow pressure prediction task while reducing computational complexity.
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.