BSA: Ball Sparse Attention for Large-scale Geometries

2025/06/14 by Catalin E. Brita, Brita, Catalin E., Hieu Nguyen +5
Computer Science · Engineering · #Artificial Intelligence (cs.AI) #Augmented Reality Applications #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #Image and Object Detection Techniques #Machine Learning (cs.LG) #Medical Imaging and Analysis

paper · pdf · doi:10.48550/arxiv.2506.12541

openalex publication_date 2025/06/14 · openalex created_date 2025/10/13 · openalex updated_date 2026/07/28

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.

Citations

• Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
• Resolution-Agnostic Transformer-based Climate Downscaling
• Latent Neural Operator for Solving Forward and Inverse PDE Problems
• Accurate structure prediction of biomolecular interactions with AlphaFold 3
• Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators
• Locality-Sensitive Hashing-Based Efficient Point Transformer with Applications in High-Energy Physics
• Transolver: A Fast Transformer Solver for PDEs on General Geometries
• Point Transformer V3: Simpler, Faster, Stronger
• Improved Operator Learning by Orthogonal Attention
• Transformers are efficient hierarchical chemical graph learners
• Geometry-Informed Neural Operator for Large-Scale 3D PDEs
• GNOT: A General Neural Operator Transformer for Operator Learning
• Solving High-Dimensional PDEs with Latent Spectral Models
• FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer
• Linear attention coupled Fourier neural operator for simulation of three-dimensional turbulence
• Transformer for Partial Differential Equations' Operator Learning
• Fourier Neural Operator for Parametric Partial Differential Equations
• Big Bird: Transformers for Longer Sequences
• GLU Variants Improve Transformer
• Root Mean Square Layer Normalization
• Attention Is All You Need
• PointNet: Deep Learning on Point Sets for 3D Classification and\n Segmentation

Related

• Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasets
• Scatterbrain: Unifying Sparse and Low-rank Attention Approximation
• Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light
• SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference
• SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space