用信息价值判断是否扩大注意力块选择,提升长文本理解准确率。
Uncertainty-gated selection for block-sparse attention
- 根据查询的置信度动态扩展保留的注意力块,避免关键信息丢失。
- 在215个上下文长度下,召回率提升28个百分点,接近稠密注意力表现。
- 适用于多种模型架构,显著提升长序列推理效率与精度。
块稀疏注意力通过每查询在键块中选择top-k来替代O(N²)的softmax,实现长上下文语言模型的可扩展性。但该方法存在短视问题:当第k与第(k+1)个块得分接近时,选择器无法追加预算,导致携带答案线索的块被丢弃且不可恢复。本文提出一种信息价值路由机制,衡量每个查询的top-k决策确定性,并对置信度最低的查询加倍保留块数;该规则不依赖骨干网络,可与现有块评分方法(如Quest)叠加使用。在LongBench-v2 medium数据集上(n=215),router-on-Quest的配对召回率达0.75,相比SSA式基线(top-k)提升28个百分点(McNemar p<0.01),且在相同上下文中距离稠密注意力仅差2个百分点。该提升在三个架构的四种模型(Qwen2.5、Mistral-Nemo、Qwen3.6)上复现。在128K上下文长度下,该方法在Qwen2.5-7B-1M和Qwen3.6上分别保持稠密精度的0.81和0.89,而基线仅0.09;同时融合选择与内核的流水线运行时间分别为稠密的0.62x和0.80x。
原文摘要 · Abstract (English)
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) -- and lands within 2 pp of dense on RULER NIAH multikey at the same context. The lift reproduces on four models from three architectures (Qwen2.5, Mistral-Nemo, Qwen3.6). At 128K, the router preserves 0.81 and 0.89 of dense accuracy on Qwen2.5-7B-1M and Qwen3.6 (vs. SSA-style top-k at 0.09 on the former) while the fused selection-plus-kernel pipeline runs at 0.62x and 0.80x dense wall time.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。