arXiv:2604.20920cs.LG2026-04被引 3

用简化的概括令牌实现高效长文本注意力,推理速度更快且更准。

Simplified Sparse Attention via Gist Tokens

  • 通过插入概括令牌并训练模型压缩关键信息,无需改动架构。
  • 在LongBench上比压缩和稀疏注意力基线表现更好,检索生成任务超全注意力5.7分。
  • 支持层级结构,32倍压缩下仍保持高精度,适合长文本生成场景。

稀疏注意力可降低长上下文推理成本,但多数方法引入新架构组件。本文提出简化稀疏注意力(SSA),无需任何架构改动。具体地,先在交错插入概括令牌的序列上进行持续预训练,优化标准下一个词预测损失,同时通过注意力掩码限制模型对上下文的访问范围,使模型学会将每段关键信息编码至概括令牌中。推理时,仅用查询与少量概括令牌计算注意力得分,选择性展开最相关的前k段,并重新引入原始标记。由于查询仅与概括令牌交互,避免了直接与完整键值缓存计算带来的内存带宽开销,无需额外的辅助键值缓存。在LongBench上,SSA在相同压缩比下持续优于压缩和推理时稀疏注意力基线。更显著的是,在检索增强生成任务中,经持续预训练后,SSA性能甚至超过全注意力模型达5.7分以上,归因于其选择性展开机制能聚焦查询相关片段并有效过滤噪声。进一步提出层次化概要-概要变体(H-SSA),在高达32倍压缩比下实现对数线性解码复杂度,同时保持或提升准确率。代码已开源。

原文摘要 · Abstract (English)

Sparse attention can reduce the cost of long-context inference, but most variants introduce new architectural components. We introduce Simplified Sparse Attention (SSA), a simpler approach to sparse attention that requires no architectural changes. Concretely, we first perform continued pretraining on sequences interleaved with gist tokens. We optimize the standard next-token loss as usual, but the gist tokens use an attention mask to restrict what parts of the context the language model can attend to; this teaches the model to pack each chunk's important information into the gist tokens. At inference time, SSA scores chunks via attention between the current query and the small set of gist tokens, selectively unfolding the top-k chunks by reintroducing their corresponding raw tokens. Since the query is scored only against the gist tokens, we avoid the memory-bandwidth cost associated with naive scoring against the full KV cache, without requiring the auxiliary KV cache approach used by sparse attention methods. On LongBench, SSA consistently outperforms compression and inference-time sparse-attention baselines under the same compression ratio. More strikingly, in retrieval-augmented generation, SSA can even outperform full attention after continued pretraining by over 5.7 points. We attribute this to the ability of SSA's selective unfolding, which concentrates attention on the query-relevant chunks and effectively filters out noise. SSA further extends to a hierarchical gist-of-gist variant (H-SSA) that achieves log-linear decoding complexity while maintaining or improving accuracy at high compression ratios up to 32x. The code is available at https://github.com/yuzhenmao/simplified-sparse-attention/.

稀疏注意力长文本生成高效推理概括令牌

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。