Top-K Retrieval with Fixed-Size Linear-Attention Completion: Backbone- and KV-Format-Preserving Attention for KV-Cache Read Reduction

📄 arXiv: 2604.05438v1 📥 PDF

作者: Yasuto Hoshi, Daisuke Miyashita, Jun Deguchi

分类: cs.LG, cs.CL

发布日期: 2026-04-07

💡 一句话要点

提出固定大小线性注意力补全的Top-K检索，减少KV缓存读取，提升长文本生成效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 长文本生成 KV缓存 注意力机制 Top-K检索 线性注意力 信息检索 Transformer模型

📋 核心要点

  1. 现有长文本生成方法受限于KV缓存读取，尤其是在长上下文场景下，效率瓶颈明显。
  2. 提出检索-补全注意力模块，通过固定大小特征图摘要估计未检索token的注意力贡献，减少KV读取。
  3. 实验表明，该方法在长文本基准测试中优于传统Top-K选择，尤其在高熵注意力头中提升显著。

📝 摘要（中文）

长文本生成日益受到解码时键值(KV)缓存流量的限制，尤其是在KV卸载到GPU内存之外时。查询感知的检索（例如，Top-K选择）通过仅加载KV对的子集来减少此流量，但对子集上的softmax进行重新归一化会在注意力权重分散在未检索的token上时引入偏差。我们提出了一种检索-补全注意力模块，该模块保持骨干网络权重和KV缓存格式不变。对于每个查询，我们计算sink/tail anchors和查询相关的检索到的Top-K token上的精确注意力，并使用预填充时计算的固定大小特征图摘要来估计剩余中间区域的分子和分母。我们将精确和估计的贡献在未归一化的域中相加，并应用单个归一化，从而恢复丢失的softmax权重，而无需额外的注意力端KV读取。在长文本基准测试中，所提出的方法在匹配的token等效读取预算下优于仅选择的Top-K，在高熵头中增益最大。

🔬 方法详解

问题定义：论文旨在解决长文本生成中，由于KV缓存读取量过大导致的效率瓶颈问题。现有的Top-K检索方法虽然能减少KV读取，但直接在检索到的子集上进行softmax归一化会导致偏差，忽略了未检索token的注意力贡献，影响生成质量。

核心思路：核心思想是在Top-K检索的基础上，通过一个固定大小的特征图来估计未检索token的注意力贡献，从而在减少KV读取的同时，尽可能地恢复完整的注意力分布。这样既能降低计算复杂度，又能减少因信息损失带来的偏差。

技术框架：该方法主要包含以下几个阶段：1) Top-K检索：根据查询向量检索最相关的K个KV对。2) 精确注意力计算：对sink/tail anchors和检索到的Top-K token计算精确的注意力权重。3) 特征图摘要：在预填充阶段计算一个固定大小的特征图，用于总结未检索token的注意力信息。4) 注意力补全：使用特征图估计未检索token的注意力贡献，并将其与精确计算的注意力权重相加。5) 归一化：对所有注意力权重进行一次归一化，得到最终的注意力分布。

关键创新：关键创新在于使用固定大小的特征图来估计未检索token的注意力贡献。这种方法避免了直接读取所有KV对，从而显著减少了KV缓存的读取量。同时，通过特征图捕获了未检索token的整体信息，减少了因信息损失带来的偏差。此外，该方法保持了骨干网络权重和KV缓存格式不变，易于集成到现有的Transformer模型中。

关键设计：特征图的大小是关键参数，需要在计算复杂度和信息损失之间进行权衡。论文中具体如何生成和使用这个特征图的细节（例如，特征图的维度、生成方式、如何与精确注意力权重融合等）需要参考原文。此外，损失函数的设计也至关重要，需要确保特征图能够准确地估计未检索token的注意力贡献。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在长文本基准测试中，在匹配的token等效读取预算下，优于仅选择的Top-K方法。尤其是在高熵注意力头中，性能提升更为显著。这意味着该方法能够更有效地利用有限的KV缓存资源，并更好地处理复杂的注意力模式。

🎯 应用场景

该研究成果可应用于各种需要处理长文本的场景，例如长文档摘要、对话生成、代码生成等。通过减少KV缓存读取，可以显著提高长文本生成的效率，降低计算成本，并支持更大规模的语言模型部署。该方法还有助于在资源受限的设备上运行大型语言模型，例如移动设备和边缘设备。

📄 摘要（原文）

Long-context generation is increasingly limited by decode-time key-value (KV) cache traffic, particularly when KV is offloaded beyond GPU memory. Query-aware retrieval (e.g., Top-K selection) reduces this traffic by loading only a subset of KV pairs, but renormalizing the softmax over the subset introduces bias when attention mass is spread over unretrieved tokens. We propose a retrieval-completion attention module that keeps backbone weights and the KV-cache format unchanged. For each query, we compute exact attention over sink/tail anchors and the query-dependent retrieved Top-K tokens, and estimate the remaining mid-region numerator and denominator using a fixed-size feature-map summary computed at prefill time. We add the exact and estimated contributions in the unnormalized domain and apply a single normalization, recovering the missing softmax mass without additional attention-side KV reads. Across long-context benchmarks, the proposed method improves over selection-only Top-K at matched token-equivalent read budgets, with the largest gains in high-entropy heads.