arXiv preprint ·
Block Sparse Attention with Log-Linear Complexity
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
Scaling language models to long contexts is limited by the quadratic cost of self-attention. Block sparse attention offers an efficient alternative, but selecting the retained blocks remains a bottleneck. Conventional block selection requires scoring all query-block pairs and therefore remains quadratic in sequence length. To address this issue, we propose PISA, a block-sparse attention mechanism that employs a pyramid Top-$K$ selection strategy. The main idea is to gradually narrow down the candidates across different levels, making it more efficient to find the most relevant keys. Specifically, we construct a coarse-to-fine hierarchy of keys and perform selection from the coarsest level. At each level, LogSumExp scoring is applied to a bounded candidate set to select candidates for the next finer level, continuing until the finest level is reached. Through pooling, we construct $O(\log N)$ levels of keys, yielding an overall complexity of $O(N\log N)$, where $N$ denotes the sequence length. We develop hardware-aware Triton kernels for both training and inference, fusing hierarchical routing and LogSumExp scoring without materializing the query-key score matrix. We further evaluate our method on language modeling tasks. Compared with the baseline, our method achieves comparable performance on benchmarks such as commonsense reasoning while delivering better results on retrieval tasks.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
MARCEDES: Score-based causal discovery under non-Gaussianity with continuous optimization
Emergent One-Third Scaling Law as Attention Tries to Concentrate
Causal Retention in Interactive Agents: Interface Factorization and Selective Adaptation
LLM Unlearning Evaluation with TRIAGE
DiffusionShadow: Diffusion-based Shadow Caching for Neural Volume Rendering
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.