In plain words
Language models can take longer to answer when they must look back at a very long conversation or document. One way to save work is to let the model look at only some earlier tokens, but choosing only the highest-scoring ones can miss useful information.
This paper introduces Stratified Sparse Attention, which keeps high-scoring tokens and samples other tokens using estimates of their importance.
On one long-context test, it beat a practical method that samples the remaining tokens uniformly, though an exact-score selection method did better in one reported condition.
The research
When a language model handles a long context, sparse attention can reduce repeated reading of its growing cache, but choosing only top-scoring tokens may miss useful attention and uniform sampling can waste selections.
Can a method preserve more of the useful context without giving up the efficiency of attending to only a fraction of tokens?
Stratified Sparse Attention combines high-score token retention with score-guided sampling to capture more useful context under sparse attention.
The authors ask how sparse attention can preserve useful information from a long context without reading every cached token.
Stratified Sparse Attention combines deterministic retention of high-scoring tokens with sampling of the residual tokens guided by proxy scores, then corrects sampled-token weights for their chances of inclusion.
On RULER-HARD, SSA scored above a practical hybrid that uses the same proxy scores and attention budget but samples residual tokens uniformly.
The comparison supports the authors’ interpretation that score-guided sampling helps in these conditions; an exact-score top-k method nevertheless scored higher in one reported condition.
These results cover selected models, tests and context lengths. On Ministral-3-8B at 128K context, Oracle-top-k scored higher than SSA on RULER-HARD; the practical GPU implementation also uses a fixed high-scoring-token fraction instead of the formal algorithm’s adaptive allocation, and clamps inclusion probabilities in a way the authors say adds a small bias.
Analysis
Deep
- Original paper authors
- Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez
- Submitted to arXiv
- License
- CC BY 4.0