Block sparse attention speeds up long text processing in language models
Block Sparse Attention with Log-Linear Complexity
Machine Learning
Summary
Long pieces of text are hard for computers to fully analyze because the usual method of paying attention to every part with every other part gets very slow. The authors created a way to break the text into blocks and smartly pick only the most important blocks to focus on, making the process much faster. They do this by organizing the blocks in levels from coarse to fine, quickly narrowing down which parts to pay attention to without checking everything. Their new method works faster and keeps similar quality when tested on language understanding and retrieval tasks.
What this means in practice
- •For natural language processing engineers: Improve efficiency of language models on long documents by using hierarchical block sparse attention with log-linear speedup.
- •For information retrieval teams: Accelerate retrieval tasks by focusing attention mechanisms on the most relevant blocks of a query, improving result quality and speed.
Authors
Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu
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.