通过Gist Tokens的简化稀疏注意力

Hugging Face Daily Papers 论文

摘要

本文介绍了简化稀疏注意力(SSA),一种在持续预训练中使用Gist令牌的方法,能够在推理时无需架构更改即可实现高效的分块选择,取得了高压缩比,并在LongBench和检索增强生成等长上下文任务上优于基线。

稀疏注意力可以降低长上下文推理的成本,但大多数变体引入了新的架构组件。我们提出了简化稀疏注意力(SSA),这是一种更简单的稀疏注意力方法,无需任何架构更改。具体来说,我们首先在交错插入Gist令牌的序列上进行持续预训练。我们照常优化标准的下一个令牌损失,但Gist令牌使用注意力掩码来限制语言模型可以关注的上下文部分;这教会了模型将每个块的重要信息打包到Gist令牌中。在推理时,SSA通过当前查询与一小部分Gist令牌之间的注意力来对块进行评分,并通过重新引入对应的原始令牌来选择性展开top-k块。由于查询仅针对Gist令牌进行评分,我们避免了与完整KV缓存进行朴素评分相关的内存带宽成本,而无需像其他稀疏注意力方法那样使用辅助KV缓存方法。在LongBench上,在相同压缩比下,SSA持续优于压缩和推理时的稀疏注意力基线。更引人注目的是,在检索增强生成中,经过持续预训练后,SSA甚至能比全注意力高出5.7个百分点。我们将此归因于SSA选择性展开的能力,它将注意力集中在与查询相关的块上,并有效过滤掉噪声。SSA进一步扩展到分层Gist-of-Gist变体(H-SSA),该变体实现了对数线性解码复杂度,同时在高压缩比(高达32倍)下保持或提高准确性。代码可在 https://github.com/yuzhenmao/simplified-sparse-attention/ 获取。
查看原文
查看缓存全文

缓存时间: 2026/06/29 22:05

论文页面 - 基于要点令牌的简化稀疏注意力

来源:https://huggingface.co/papers/2604.20920

摘要

简化稀疏注意力(SSA)通过在预训练期间使用基于要点令牌的注意力掩码来降低长上下文推理成本,从而在推理时无需架构修改即可实现高效的分块选择。

稀疏注意力(https://huggingface.co/papers?q=Sparse%20attention)可以降低长上下文推理的成本,但大多数变体都会引入新的架构组件。我们提出了简化稀疏注意力(https://huggingface.co/papers?q=Sparse%20Attention)(SSA),这是一种更简单的稀疏注意力(https://huggingface.co/papers?q=sparse%20attention)实现方法,无需任何架构更改。具体来说,我们首先在穿插有要点令牌(https://huggingface.co/papers?q=gist%20tokens)的序列上进行持续预训练(https://huggingface.co/papers?q=continued%20pretraining)。我们照常优化标准的下一令牌损失(https://huggingface.co/papers?q=next-token%20loss),但要点令牌(https://huggingface.co/papers?q=gist%20tokens)使用注意力掩码(https://huggingface.co/papers?q=attention%20mask)来限制语言模型可以关注的上下文部分;这教会模型将每个分块的重要信息打包到要点令牌(https://huggingface.co/papers?q=gist%20tokens)中。在推理时,SSA通过当前查询与少量要点令牌(https://huggingface.co/papers?q=gist%20tokens)之间的注意力对分块进行评分,并通过重新引入相应的原始令牌来选择性地展开top-k分块。由于查询仅与要点令牌(https://huggingface.co/papers?q=gist%20tokens)评分,我们避免了与对整个KV缓存(https://huggingface.co/papers?q=KV%20cache)进行朴素评分相关的内存带宽成本,并且无需稀疏注意力(https://huggingface.co/papers?q=sparse%20attention)方法所使用的辅助KV缓存(https://huggingface.co/papers?q=KV%20cache)方法。在LongBench(https://huggingface.co/papers?q=LongBench)上,SSA在相同压缩比下持续优于压缩和推理时稀疏注意力基线。更引人注目的是,在检索增强生成(https://huggingface.co/papers?q=retrieval-augmented%20generation)中,经过持续预训练(https://huggingface.co/papers?q=continued%20pretraining)后,SSA甚至可以比完整注意力高出超过5.7个点。我们将此归因于SSA的选择性展开(https://huggingface.co/papers?q=selective%20unfolding)能力,它将注意力集中在与查询相关的分块上,并有效过滤掉噪声。SSA进一步扩展到分层要点之要点(https://huggingface.co/papers?q=hierarchical%20gist-of-gist)变体(H-SSA),该变体实现了对数线性解码复杂度(https://huggingface.co/papers?q=log-linear%20decoding%20complexity),同时在高压缩比(高达32倍)下保持或提高了准确性。代码可在https://github.com/yuzhenmao/simplified-sparse-attention/获取。

查看arXiv页面(https://arxiv.org/abs/2604.20920)查看PDF(https://arxiv.org/pdf/2604.20920)添加到收藏(https://huggingface.co/login?next=%2Fpapers%2F2604.20920)

在你的代理中获取此论文:

hf papers read 2604.20920

没有最新的CLI?curl -LsSf https://hf.co/cli/install.sh | bash

引用此论文的模型0

没有模型链接此论文

请在你的模型README.md中引用arxiv.org/abs/2604.20920以从此页面链接。

引用此论文的数据集0

没有数据集链接此论文

请在你的数据集README.md中引用arxiv.org/abs/2604.20920以从此页面链接。

引用此论文的Space0

没有Space链接此论文

请在你的Space的README.md中引用arxiv.org/abs/2604.20920以从此页面链接。

包含此论文的收藏集0

没有收藏集包含此论文

请将此论文添加到一个收藏集(https://huggingface.co/new-collection)以从此页面链接。

相似文章

MiniMax 稀疏注意力

Hugging Face Daily Papers

MiniMax 稀疏注意力 引入了一种分块稀疏注意力机制,针对超长上下文的大语言模型实现了显著的加速。在1M上下文长度下,每个token的注意力计算减少28.4倍,在H800 GPU上预填充阶段实际速度提升14.2倍,解码阶段提升7.6倍。该方法附带了一个开源推理内核以及一个公开发布的多模态模型。