Score-guided sampling improves sparse attention
Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity · arXiv:2610.28441
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.
What the evidence shows
- The long-context trade-off
Stratified Sparse Attention (SSA) is a method for reducing the work of language-model attention over long contexts. As a model generates text, its cache of earlier keys and values grows, and attention repeatedly reads from it.
Sparse methods read only some tokens, but deterministic selection can miss relevant attention mass, while uniform sampling does not distinguish among the residual tokens.
- A combined selection method
SSA divides an attention budget between tokens retained for high proxy scores and sampled tokens from the residual context. It samples those remaining tokens according to their proxy scores, then calculates exact attention scores for selected tokens and corrects for the probability that each sampled token was included. This combines deterministic retention with score-guided coverage.
- A diagnostic for sampling
The authors first compare score-guided sampling with uniform sampling using relative L2 attention error, a measure of how far sparse attention differs from the reference. Score-guided sampling has lower error across the settings shown.
This diagnostic supports the sampling choice, though the cited figure passage does not name its model or dataset and the result is not itself a language-model accuracy test.
- Evidence on long contexts
On RULER-HARD, which tests performance on long-context tasks and reports accuracy averaged over task categories, SSA outscored vAttention(PQ) for Ministral-3-8B at both reported context lengths under 50x attention sparsity.
The methods use the same PQ scores and attention budget, but SSA guides residual sampling with scores while vAttention(PQ) samples uniformly. The authors interpret this comparison as evidence for score-guided sampling.
- Evidence and its limits
The result has a clear exception: at 128K on Ministral-3-8B, Oracle-top-k scored above SSA on RULER-HARD. A separate AIME 2025 result for Qwen3.5-9B favors SSA over the reported alternatives at 50x sparsity, extending the evidence to long-form generation.
These comparisons support SSA under particular conditions, not a general ranking across all models and tests.
Deep
- Original paper authors
- Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez
- Submitted to arXiv
- License
- CC BY 4.0