SelKV：基于逐标记合并或丢弃及注意力补偿的选择性 KV 缓存合并

# 选择性KV缓存合并：基于逐Token合并或丢弃与注意力补偿机制

**来源：** https://arxiv.org/html/2607.16213

**Soumia Bouyahiaoui**
阿尔及利亚国家人工智能学院（ENSIA），西迪阿卜杜拉校区，阿尔及尔，阿尔及利亚
[email protected]

**Manel Kara laouar**
阿尔及利亚国家人工智能学院（ENSIA），西迪阿卜杜拉校区，阿尔及尔，阿尔及利亚
[email protected]

**Aicha Boutorh**
阿尔及利亚国家人工智能学院（ENSIA），西迪阿卜杜拉校区，阿尔及尔，阿尔及利亚
[email protected]

**Mohamed Hadj Ameur**
阿尔及利亚国家人工智能学院（ENSIA），西迪阿卜杜拉校区，阿尔及尔，阿尔及利亚
[email protected]

###### 摘要

大型语言模型（LLMs）以自回归方式生成文本，依赖一个键值（KV）缓存，其内存占用随上下文长度线性增长，成为主要瓶颈。最近的压缩方法通过token合并来缓解这一问题；然而，这些方法通常依赖无差别的聚合，这会损害表示质量并引入*注意力塌陷*，即合并后的token虽然编码了多个输入，却在softmax中获得了与单个token相同的质量。我们提出一个无需训练的双组件框架，用于KV缓存压缩，以解决这些局限。首先，一个*软余弦门控*根据值向量相似性自适应地调节合并决策，抑制或丢弃不相似的token以保持语义保真度。其次，我们引入一个*注意力比率*补偿机制，该机制在解码时应用一个基于预填充注意力统计信息的logit偏置，以纠正合并引起的softmax不平衡。在LongBench（16个英文数据集）上评估时，仅保留25%的KV缓存，我们的框架在代表性的单次（one-shot）基线中取得了强劲的压缩性能。它在评估的分组查询注意力（GQA）模型上尤其鲁棒，保持了近乎无损的生成质量。此外，在复杂的多文档QA任务中，该方法超过了完整缓存基线，并在100k token下实现了3.3倍解码加速。

*关键词* KV缓存压缩，token合并，注意力补偿，大型语言模型，推理效率

## 1 引言

大型语言模型（LLMs）在各种任务中表现出色，但其自回归推理受到键值（KV）缓存的瓶颈限制，该缓存存储了过去的关键和值向量用于注意力计算。对于一个具有L层、H个注意力头（维度d）的模型，KV缓存随序列长度n增长为O(L·H·d·n)，即使在中等上下文长度下也会消耗数个GB。随着应用对更长上下文的需求增加，KV缓存内存已成为批次大小、吞吐量和部署成本的主要约束。

目前有两种主要的KV缓存压缩方法。*逐出*方法[19（https://arxiv.org/html/2607.16213#bib.bib6），7（https://arxiv.org/html/2607.16213#bib.bib5），15（https://arxiv.org/html/2607.16213#bib.bib4），3（https://arxiv.org/html/2607.16213#bib.bib8）]永久丢弃低重要性token，减少了缓存大小，但不可逆地丢失了它们携带的信息。*合并*方法如KVMerger[14（https://arxiv.org/html/2607.16213#bib.bib9）]、WeightedKV[17（https://arxiv.org/html/2607.16213#bib.bib10）]、D2O[12（https://arxiv.org/html/2607.16213#bib.bib11）]和LOOK-M[13（https://arxiv.org/html/2607.16213#bib.bib12）]则将被逐出的token合并到保留的缓存条目中，当信息分散在多个token中时，这可以保留更多的语义内容。我们关注当前合并方法的两个局限：

#### 问题1：统一的合并或丢弃决策

大多数方法对所有*被逐出的*token应用相同的策略：要么全部合并（KVMerger、WeightedKV、LOOK-M），要么全部丢弃[19（https://arxiv.org/html/2607.16213#bib.bib6），7（https://arxiv.org/html/2607.16213#bib.bib5），15（https://arxiv.org/html/2607.16213#bib.bib4）]。D2O[12（https://arxiv.org/html/2607.16213#bib.bib11）]根据注意力模式将token划分为丢弃集和合并集，但这仍然是一个二元决策，对分配给合并集的每个token以完整权重进行合并。没有任何现有方法使用表示相似性在连续、逐token的基础上调节*合并多少*。当一个被逐出的token的值向量与其合并目标不相似时，加权平均可能会破坏保留的表示。我们的实验说明了这种权衡：统一合并在HotpotQA上将多文档QA提升了+0.92点，但在TREC上将少样本分类降低了-4.0点。EMS[18（https://arxiv.org/html/2607.16213#bib.bib14）]同样指出*“过度合并和过度逐出都会导致次优性能，”*但没有提供逐token机制。

#### 问题2：注意力塌陷

当多个token合并到单个缓存条目中时，该条目仍然大致获得与未合并token相同的softmax注意力，尽管它代表多个原始token。KeepKV[11（https://arxiv.org/html/2607.16213#bib.bib15）]将这一现象形式化为*注意力塌陷*（定理3.2）：一个代表m个合并token的位置被系统性地低估了约∼m倍。因此，许多现有的合并方法可能未充分利用它们保留的信息[14（https://arxiv.org/html/2607.16213#bib.bib9），17（https://arxiv.org/html/2607.16213#bib.bib10），12（https://arxiv.org/html/2607.16213#bib.bib11），8（https://arxiv.org/html/2607.16213#bib.bib13）]。

我们通过SelKV（选择性KV缓存合并）解决这两个问题，这是一个无需训练的框架，包含两个可以集成到现有合并流程中的互补机制：

1.  **软余弦门控**（问题1）：对于每个被逐出的token，我们计算其值向量与合并目标之间的余弦相似性。门控g = max(cos_sim, 0) 然后连续地调节合并强度：相似token完全合并（g≈1），正交token被丢弃（g=0），中间情况部分合并。无需学习参数或阈值调整。
2.  **注意力补偿**（问题2）：我们在解码期间添加一个注意力比率logit偏置，使用预填充注意力权重来估计每个合并位置应该获得多少注意力质量。与基于计数的朴素补偿（log(1+Mi)）不同，这种公式是自然校准的。

自适应逐token路由在视觉Transformer中已有先例，DiffRate[4（https://arxiv.org/html/2607.16213#bib.bib16）]表明学习的合并或剪枝决策可以优于统一策略。我们的软余弦门控是KV缓存的无需训练对应物。

我们提出SelKV（选择性KV缓存合并），一个无需训练的框架，包含一个*软余弦门控*，自适应调节逐token合并强度，以及*注意力比率补偿*，重新校准softmax注意力朝向合并位置。我们在LongBench[1（https://arxiv.org/html/2607.16213#bib.bib3）]（16个英文数据集）上评估了三个模型，涵盖MHA和GQA架构，KV保留预算为25%。SelKV在*分组查询注意力*（GQA）上尤其有效，最接近完整缓存性能。它还在多个多文档QA设置中超过了完整缓存基线，表明选择性合并可以充当隐式注意力过滤器。在100k token下，压缩缓存比完整缓存解码快3.3×。

本文的其余部分组织如下：第2节（https://arxiv.org/html/2607.16213#S2）回顾相关工作，第3节（https://arxiv.org/html/2607.16213#S3）介绍提出的方法；第4节（https://arxiv.org/html/2607.16213#S4）报告实验评估并与SOTA模型进行基准比较；第5节（https://arxiv.org/html/2607.16213#S5）总结并展望未来方向。

## 2 相关工作

KV缓存压缩方法主要分为三类：token逐出、token合并和量化。量化方法如KIVI[9（https://arxiv.org/html/2607.16213#bib.bib21）]和ZipCache[6（https://arxiv.org/html/2607.16213#bib.bib22）]降低精度而非token数量，并且与我们的方法基本互补，因此我们在这里重点讨论逐出和合并。

### 2.1 Token逐出

Token逐出方法通过永久丢弃被认为不重要的条目来减少KV缓存，只保留一个子集用于后续解码。

**逐步骤逐出：** H2O[19（https://arxiv.org/html/2607.16213#bib.bib6）]将缓存逐出形式化为一个动态子模问题，识别累积了不成比例注意力质量的*重要头*（Heavy Hitter）token，并在每个解码步骤丢弃其余token，同时保留一个最近token窗口。VATP[5（https://arxiv.org/html/2607.16213#bib.bib7）]扩展了这一思想，指出注意力权重本身是不完整的重要性代理：token对输出的实际贡献是其注意力权重与值幅度的乘积。虽然有效，但逐步骤方法在*每个*生成步骤重新评估重要性并重新组织缓存，带来了显著的解码开销。

**一次性逐出：** 另一种方法在预填充后一次性压缩缓存，然后使用固定压缩缓存进行生成。StreamingLLM[15（https://arxiv.org/html/2607.16213#bib.bib4）]只保留前几个*注意力汇聚*（attention sink）token和一个滑动最近窗口，无需基于注意力的评分，但牺牲了所有中间上下文信息。SnapKV[7（https://arxiv.org/html/2607.16213#bib.bib5）]使用一个观察窗口覆盖最后w个查询位置，识别每个注意力头的重要KV位置，并通过平均池化核进行平滑以保留token簇。PyramidKV[3（https://arxiv.org/html/2607.16213#bib.bib8）]应用了逐层自适应预算：较低层保留更多token，而较高层更依赖最近上下文。更近期的一次性方法包括RocketKV[2（https://arxiv.org/html/2607.16213#bib.bib23）]、EvolKV[16（https://arxiv.org/html/2607.16213#bib.bib24）]和LAVa[10（https://arxiv.org/html/2607.16213#bib.bib26）]。所有逐出方法都永久销毁了被丢弃token携带的信息。当重要证据分布在许多位置时，如在多文档QA或长距离摘要中，这尤其代价高昂。

### 2.2 Token合并

合并方法并非完全丢弃被逐出的token，而是将其信息整合到保留的缓存条目中。KVMerger[14（https://arxiv.org/html/2607.16213#bib.bib9）]识别具有相似关键表示的token序列，并通过高斯核加权合并方案进行融合。D2O[12（https://arxiv.org/html/2607.16213#bib.bib11）]根据注意力模式将token划分为应该丢弃的和应该合并的，但在每个分区内对*所有*token应用*相同*策略。LOOK-M[13（https://arxiv.org/html/2607.16213#bib.bib12）]在局部滑动窗口内合并token以保持空间局部性。WeightedKV[17（https://arxiv.org/html/2607.16213#bib.bib10）]在合并期间应用注意力加权凸组合。ChunkKV[8（https://arxiv.org/html/2607.16213#bib.bib13）]将固定大小的语义块视为压缩单元而非单个token。

大多数现有方法应用一个统一的合并或丢弃决策：要么*所有*被逐出的token以完整权重合并（KVMerger、WeightedKV、LOOK-M），要么*所有*被丢弃（H2O、SnapKV、StreamingLLM）。D2O使用注意力模式将token划分为丢弃组和合并组，但分配给合并组的token仍然无条件合并。EMS[18（https://arxiv.org/html/2607.16213#bib.bib14）]明确观察到*“过度合并和过度逐出都会导致次优性能，”*但没有提出连续逐token的门控机制。我们的软余弦门控通过基于值向量相似性在[0,1]范围内调节合并强度来填补这一空白。

另一个问题影响所有合并方法。当多个token合并到单个缓存条目中时，合并后的条目获得的softmax注意力大约与单个未合并token相同。KeepKV[11（https://arxiv.org/html/2607.16213#bib.bib15）]将之形式化为*注意力塌陷*（定理3.2）：合并位置被系统性地低估了约∼m倍。KeepKV提出基于原始合并计数log(1+Mi)的logit偏置作为修正。然而，原始合并计数可能极端（在75%压缩下高达∼200），使得基于计数的补偿不稳定。此外，大多数合并方法如KVMerger、WeightedKV、D2O、ChunkKV根本没有包含任何补偿。

## 3 方法

我们提出**SelKV**，一个可插入的选择性KV缓存压缩框架，用于解决第1节中识别的两个问题。给定一个具有L层、H个注意力头和头维度d的预训练Transformer，SelKV在预填充后*一次性*压缩KV缓存，然后使用压缩缓存进行解码。该框架引入两个机制：(1) 一个*软余弦门控*，逐token决定合并多少，以及 (2) *注意力补偿*，纠正合并位置的注意力不足。两者均无需训练，可以添加到现有合并方法中。图1（https://arxiv.org/html/2607.16213#S3.F1）展示了完整流程，在预填充后经过六个阶段：重要性评分（§3.2（https://arxiv.org/html/2607.16213#S3.SS2））、token选择（§3.3（https://arxiv.org/html/2607.16213#S3.SS3））、合并目标路由（§3.4（https://arxiv.org/html/2607.16213#S3.SS4））、带软余弦门控的选择性合并（§3.5（https://arxiv.org/html/2607.16213#S3.SS5））、注意力补偿（§3.6（https://arxiv.org/html/2607.16213#S3.SS6））和RoPE重定位（§3.7（https://arxiv.org/html/2607.16213#S3.SS7））。token选择支持*每头*模式（每个注意力头独立选择自己的token）和*共享*模式（所有头保留相同token集），注意力补偿可以跨层平均应用或通过前向钩子*逐层*应用。

![图1](https://arxiv.org/html/2607.16213/images/overview.png)
*图1：选择性KV缓存压缩流程概览。预填充后，该方法对token重要性进行评分，选择保留的token，将被逐出的token路由到合并目标，应用软余弦门控，对合并位置进行注意力补偿，并在解码前重新定位RoPE索引。*

### 3.1 概述与符号

设 x = (x1, ..., xn) 表示长度为 n 的输入序列。我们需要 (i) 完整 KV 缓存 { (K^(ℓ), V^(ℓ)) }_{ℓ=1}^L，其中 K^(ℓ), V^(ℓ) ∈ ℝ^{H×n×d}；以及 (ii) 在长度为 W 的观察窗口中的查询位置的注意力权重，用于重要性评分和路由。朴素的 eager 预填充会实现 O(n^2) 注意力，这在长上下文中不可行。相反，我们采用*两遍*策略：(1) 运行 SDPA 预填充以构建完整 KV 缓存，不具体化注意力矩阵；然后 (2) 仅对最后 W 个 token 使用 eager 注意力对缓存的关键进行重新运行，产生每层大小为 H×W×n 的注意力矩阵。这降低了预填充内存从 O(n^2) 到 O(W·n)，同时产生与 SnapKV 风格观察窗口相同的注意力权重。

给定目标保留比例 ρ ∈ (0,1)，我们设置目标缓存长度为 m = ⌊ρ·n⌋。