Uncertainty-gated selection for block-sparse attention
Block-sparse attention scales long-context language models by replacing the O(N^2) softmax with a per-query top-k selection over key blocks. This cutoff is myopic: when the k-th and (k+1)-th blocks are nearly tied in score, the selector commits without spending extra budget, and a dropped block carrying answer evidence is unrecoverable downstream. We propose a value-of-information router that measures, for each query, how decisively the top-k cut was made, and doubles the kept set for the queries where that gap is smallest; the rule is backbone-agnostic and stacks with existing block-scoring methods such as Quest. On LongBench-v2 medium at n=215 (the entire dataset subset), router-on-Quest reaches paired recall 0.75 vs. top-k 0.47 -- +28 pp over the SSA-style baseline (McNemar p<0.01) -- while attending to ~4% of the context; it lands within 2 pp of dense on RULER NIAH multikey at the same context. The lift reproduces on three models from two architecture families and at 128K context, where router-on-Quest holds paired recall 0.81 while the SSA-style baseline collapses to 0.09. The fused selection-plus-kernel pipeline runs at 0.57x dense wall time at 128K.
Authors
- Thomas Rossi (ORCID: https://orcid.org/0009-0000-6127-1048)
Institutions
- École des hautes études en sciences sociales (FR)
Publication Details
- Journal
- Zenodo (CERN European Organization for Nuclear Research)
- Published
- 2026-06-13
- DOI
- https://doi.org/10.5281/zenodo.20683079
- Primary Topic
- Natural Language Processing Techniques
- Type
- preprint