通过 Gist Tokens 简化稀疏注意力

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