R-KV: Redundancy-Aware KV Cache Compression for Reasoning Models

发表时间: 2025-12 · arXiv:2505.24133 (NeurIPS 2025)

原文: https://arxiv.org/abs/2505.24133

Zefan Cai, Wen Xiao, Hanshi Sun, Cheng Luo, Yikai Zhang, Ke Wan, Yucheng Li, Yeyang Zhou, Li-Wen Chang, Jiuxiang Gu, Zhen Dong, Anima Anandkumar, Abedelkadir Asi, Junjie Hu
(University of Wisconsin - Madison, Microsoft, Carnegie Mellon University, California Institute of Technology, University of California - San Diego, University of Surrey, University of California - Berkeley)

速读

一句话结论
本文提出了一种专为推理模型设计的冗余感知 KV Cache 压缩方法 R-KV，通过联合评估 Token 的注意力和语义冗余度，在仅保留 10% 到 16% 缓存的情况下实现了无损甚至超越全量缓存的推理准确率，并大幅提升了系统吞吐量。

要解决什么问题
大语言模型在处理复杂推理任务时（如使用 DeepSeek-R1 等推理模型），通常会生成极长的思维链（CoT）和多次反思步骤。这种自回归生成范式会导致输出长度远超输入提示，例如解决复杂数学题可能生成 3.2 万个 Token，单次请求除消耗 15.5GB 显存加载权重外，还需 4.1GB 存储键值缓存（KV Cache）。随着生成长度的激增，KV Cache 的显存占用会急剧膨胀，导致严重的显存瓶颈和极低的并发处理能力。现有的 KV Cache 压缩方法主要针对长输入提示（Prefilling 阶段）设计，且通常仅依赖注意力分数来评估 Token 的重要性。然而，推理模型的输出具有极高的冗余性，包含大量重复的自我验证、中间计算和啰嗦的自我对话。在标准的基于注意力的压缩机制下，这些重复片段会因为与之前生成的相似文本高度相关而获得极高的注意力分数。如果单纯保留“高注意力”的 Token，系统会过度保留这些毫无新意的冗余反思，同时可能会错误地裁剪掉那些注意力分数不高但对推理至关重要的零散信息。这种机制上的错位导致现有压缩方法在部署推理模型时，要么显存降不下来，要么严重破坏模型的推理准确率。因此，如何精准识别并剔除高注意力但高冗余的 Token，保留真正重要且语义多样的上下文，是当前推理模型长文本生成面临的核心卡点。

怎么做的
为绕开单纯依赖注意力分数导致的冗余陷阱，本文提出 R-KV 方法，核心思路是在解码阶段动态进行 KV Cache 压缩，将 Token 的重要性与非冗余性结合进行联合筛选。该方法由三个关键部件构成：重要性评分机制、冗余度估计机制以及联合驱逐机制。首先，重要性评分机制负责找出对后续生成有关键作用的 Token。方法利用最后 $\alpha$ 个观察 Token 对历史候选 Token 的注意力权重来衡量重要性。为避免极端异常值干扰，R-KV 在局部滑动窗口内对注意力分数进行最大池化操作，得到平滑后的重要性分数。其次，冗余度估计机制负责揪出那些语义重复的废话 Token。该机制通过计算不同 Key 向量之间的余弦相似度矩阵来衡量语义重合度。为防止误删刚生成、仍具短期上下文价值的重复 Token，设计中强制将每个 Token 与其最近的 $\beta$ 个高相似度 Token 的得分清零。随后，计算每个 Token 与其他所有 Token 的平均相似度，并通过 Softmax 归一化得到冗余度分数 $R_i^h$。分数越高，说明该 Token 的内容越是被其他 Token 共享，冗余度越大。最后，联合选择策略将上述两个维度结合，为每个注意力头 $h$ 下的每个 Token $i$ 计算最终的保留分数，其定义性公式为：

$$Z_i^h = \lambda I_i^h - (1 - \lambda) R_i^h$$

其中 $\lambda$ 是平衡重要性与冗余度的超参数，重要性分数 $I_i^h$ 越高越倾向于保留，冗余度分数 $R_i^h$ 越高越倾向于剔除。在具体执行时，R-KV 为保留的 KV Token 分配固定大小的预算缓存，并为新生成的 Token 分配固定大小的缓冲区。每当缓冲区填满一段固定长度的文本后，R-KV 会将预算缓存中的 Token 与缓冲区中除最后 $\alpha$ 个观察 Token 之外的候选 Token 拼接在一起，计算所有候选者的 $Z_i^h$ 分数，并仅将得分最高的 Token 存回预算缓存中。这种设计使得系统能够在不破坏推理链条完整性的前提下，精准剔除无意义的重复思考，从而将显存占用维持在恒定水平。

效果如何
实验在数学推理数据集 MATH-500 和 AIME 2024 上进行，使用了 DeepSeek-R1-Distill-Llama-8B 和 DeepSeek-R1-Distill-Qwen-14B 模型，最大生成长度设定为 1.6 万到 3.2 万个 Token，并采用非零温度采样（Temperature 为 0.6）计算 Pass@1 准确率。对比基线包括代表纯注意力路线的最先进压缩方法 SnapKV，以及不进行任何压缩的 FullKV（代表无损生成黄金标准）。量化结果显示，R-KV 在压缩效率和准确率上取得了突破性表现。在 DeepSeek-R1-Distill-Llama-8B 模型上，R-KV 仅需保留 10.34% 的原始 KV Cache 即可在 AIME 2024 数据集上达到与未压缩模型相当的准确率，而同等预算下 SnapKV 的性能仅能达到 FullKV 的 60%。当保留约 16% 的 KV Cache 时，R-KV 准确率甚至达到 FullKV 基线的 105%，表明剔除冗余信息反而有助于模型避开重复死循环，提升推理质量。在吞吐量方面，由于 R-KV 将显存占用降低了 90%，使得系统能够支持更大的批处理大小。在 16K 上下文长度设置下，R-KV 相比 FullKV 支持的批处理大小提升了 9 倍，端到端吞吐量提升了 6.6 倍；在固定缓存预算下，吞吐量甚至可提升 9.2 倍。作者也坦诚了局限性：R-KV 目前与 Paged Attention 等高级内存管理机制的兼容性仍是挑战，且需要推理服务框架提供专门的压缩接口，否则频繁重新分配和释放显存的工程开销，可能会抵消掉很大一部分加速收益。

主要贡献

核心问题：近期的大型语言模型（LLMs）在复杂推理和自我反思方面展现出卓越能力，但推理模型（如DeepSeek-R1）在部署时面临一个关键挑战：它们倾向于生成过长且冗余的推理轨迹，导致自回归生成过程中键值（KV）缓存快速增长，从而引发难以承受的内存需求。现有的KV缓存压缩方法主要针对长输入提示（Prompt）设计，缺乏对长生成输出的深入探索。同时，标准的基于注意力的KV缓存压缩方法在处理推理模型时往往会失效，因为重复的冗余内容会为其自身生成很高的注意力信号，导致重要但分散的推理信息被丢弃，而冗余的自我反思被过度保留。

研究目标：提出一种专为推理模型中冗余Token设计的KV缓存压缩策略，在解码过程中选择性地保留“重要且非重复的上下文”，从而在解决实际内存限制的同时，保持模型关键的推理能力。

创新点：
1. 提出了R-KV（Redundancy-aware KV Cache Compression for Reasoning models），这是一种免训练且与模型无关的解码期KV缓存压缩方法，专门针对推理模型的冗余生成问题。
2. 引入了混合评分机制，包含基于注意力权重的关键度评分机制以保留关键Token，以及基于键向量（Key Vectors）实时语义相似度分析的动态冗余度评分机制以识别重复Token。
3. 设计了联合驱逐机制，平衡冗余度与关键度，优化缓存效率。实验表明，R-KV仅需使用10%的KV缓存即可保持近100%的完整KV缓存性能，在16%的缓存预算下甚至达到了完整性能的105%，并带来了90%的内存节省和6.6倍的吞吐量提升。

R-KV架构图
R-KV架构图

关键观察与设计原则

推理模型中的冗余现象：推理模型通常会生成详细的思维链（CoT）和多个反思步骤，导致响应长度远超标准模型。在MATH-500和AIME 2024数据集上，DeepSeek-R1的蒸馏模型（8B、7B、14B）生成的输出长度比Ground Truth长出8到14倍以上。然而，并非所有新增的Token都能提供有意义的内容，解码上下文很大程度上被重复内容主导。数据显示，推理模型生成输出中1-gram和2-gram的平均频率始终高于Ground Truth（高出约5.7倍），这表明其生成内容存在高度重复。
推理模型与Ground Truth的生成长度及n-gram频率对比

现有KV压缩方法在处理冗余时的失效：大多数现有的KV缓存压缩方法主要基于Token的上下文重要性（通常通过Key和Query之间的注意力分数衡量，如【3, Snapkv: Llm knows what you are looking for before generation+2024】）来选择Token。这种方法虽然能保留关键上下文，但未能考虑冗余问题。在推理模型中，重复的内容往往会获得不成比例的高注意力分数，因为它们与之前生成的重复文本高度相似。这导致冗余Token被过度保留，无谓地膨胀了KV缓存大小，却未提供新的有效信息。可视化结果表明，基于注意力的SnapKV方法选择了大量与自我反思和最终答案结论相关的重复Token。
SnapKV选择的KV缓存可视化

方法细节

解码期压缩机制：与关注预填充（Prefilling）阶段以管理长上下文输入的现有方法（如SnapKV、PyramidKV等）不同，R-KV专注于推理模型的解码（Decoding）阶段，这是一个生成输出远长于输入的特殊场景。具体而言，R-KV为两部分分配内存：一个是大小为$B_{budget}$的缓存（Cache），用于存储保留的KV Tokens；另一个是大小为$B_{buffer}$的缓冲区（Buffer），用于存放新生成的文本Tokens。总内存需求为$B_{total} = B_{budget} + B_{buffer}$。在模型于缓冲区生成每个固定长度的文本段后，R-KV会执行KV缓存压缩。在每个文本段结束时，遵循先前工作【3, Snapkv: Llm knows what you are looking for before generation+2024】的做法，最后$\alpha$个Tokens始终作为观察Tokens（Observation Tokens）保留在缓存中。接着，将缓存中现有的$B_{budget}$个Tokens与缓冲区中的前$B_{buffer} - \alpha$个Tokens拼接，形成$n = B_{budget} + B_{buffer} - \alpha$个候选KV Tokens。每个候选Token会被分配一个选择分数，随后选出得分最高的$k = B_{budget} - \alpha$个Tokens填入剩余的缓存预算中（外加$\alpha$个观察Tokens）。此过程在保留关键上下文的同时压缩了KV缓存，实现了自回归解码期间的高效内存利用。

基于注意力权重的关键度评分：R-KV利用注意力权重来估计Token的重要性，其直觉在于获得更高注意力的Tokens对解码贡献更大。具体来说，计算每个Key Token从解码期间最后$\alpha$个观察Tokens接收到的注意力分数。除了标准的多头注意力（MHA），R-KV还支持分组查询注意力（GQA）的评分估计。
- 多头注意力（MHA）：给定最后$\alpha$个观察Tokens作为查询$Q^h \in \mathbb{R}^{\alpha \times d}$，以及每个注意力头$h$的$n$个键状态$K^h \in \mathbb{R}^{n \times d}$，注意力分数$A^h \in \mathbb{R}^{\alpha \times n}$的计算方式为：

$$A^h = \operatorname{softmax}(Q^h \cdot (K^h)^\top / \sqrt{d})$$


- 分组查询注意力（GQA）：在GQA中，每个键/值头$h$由一组$G$个不同的查询头（索引为$g \in [0, G)$）共享。共享的键/值状态记为$\pmb{K}^h, \pmb{V}^h \in \mathbb{R}^{n \times d}$，组内的$G$个查询状态记为$\pmb{Q}^{h,0}, \ldots, \pmb{Q}^{h,G-1} \in \mathbb{R}^{\alpha \times d}$。组内每个查询头的注意力分数计算为：

$$\pmb{A}_{\mathrm{group}}^{h,g} = \pmb{Q}^{h,g} \cdot (\pmb{K}^h)^\top / \sqrt{d} \in \mathbb{R}^{\alpha \times n}, \quad \mathrm{for } g = 0, \ldots, G-1$$
随后，使用最大池化（Max-pooling）操作在组维度上将这$G$个矩阵聚合为头组$h$的单一矩阵$A_{\mathrm{group}}^h$。最后，沿Key Token维度重新归一化得到最终的注意力权重$A^h$：
$$A_{\mathrm{group}}^h = \operatorname*{max pool} \left( \left[ A_{\mathrm{group}}^{h,0}, \ldots, A_{\mathrm{group}}^{h,G-1} \right] \right) \in \mathbb{R}^{\alpha \times n}, \quad A^h = \mathrm{softmax} \left( A_{\mathrm{group}}^h \right) \in \mathbb{R}^{\alpha \times n}$$
- 分数稳定化与关键度估计：由于从$A^h$得出的单Token重要性分数可能包含异常高值，导致估计不稳定，R-KV采用滑动窗口（大小为$2W$）对最近Tokens的单Token重要性分数应用最大池化操作。设$A_{j,i}^h$为第$j$个Query对第$i$个Key的注意力分数，计算稳定后的分数$\tilde{A}_{j,i}^h$，最终得到保留第$i$个Token的关键度分数$I_i^h$：
$$\tilde{A}_{j,i}^h = \operatorname*{max} \left( A_{j,i-W}^h, \ldots, A_{j,i}^h, \ldots, A_{j,i+W-1}^h \right), \quad I_i^h = \frac{1}{\alpha} \sum_{j=0}^{\alpha-1} \tilde{A}_{j,i}^h \in \mathbb{R}$$

基于语义相似度的冗余度估计：为了识别冗余Tokens，R-KV使用余弦相似度测量键状态（Key States）之间的语义相似度。与其他Tokens相似度高的Tokens被视为潜在冗余。
- Key Tokens间的余弦相似度：给定头$h$的键Tokens $K^h \in \mathbb{R}^{n \times d}$，首先将每个键向量$K_i^h$归一化为$\overline{\mathbf{K}}_i^h$，然后计算余弦相似度矩阵$S^h$：

$$\overline{\mathbf{K}}_i^h = \frac{\mathbf{K}_i^h}{\| \mathbf{K}_i^h \|_2 + \epsilon} \in \mathbb{R}^d, \quad S^h = \overline{\mathbf{K}}^h (\overline{\mathbf{K}}^h)^\top \in \mathbb{R}^{n \times n}, \quad S_{i,i}^h = 0, \forall i \in [0, n)$$


为了防止Token与自身被标记为冗余，将对角线元素$S_{i,i}^h$置为0。
- 强制保留近期Tokens：由于直接移除所有冗余Tokens可能会损害模型性能，R-KV在具有高相似度的Tokens中，仅保留最近生成的$\beta$个Tokens。具体操作是，对于每个Token $i$，找出高度相似的Token索引集合$\mathcal{T}_i^h = \{ j \mid S_{j,i}^h > T, j \in [0, n) \}$（$T$为相似度阈值）。从中提取包含最多$\beta$个最大索引（即最近的$\beta$个相似Tokens）的子集$\mathcal{T}_{i,\beta}^h$。然后将矩阵$S^h$中这些近期Tokens对应的相似度分数清零，即$S_{j,i}^h \gets 0, \forall j \in \mathcal{T}_{i,\beta}^h$。
- 冗余度评分计算：对于每个头$h$中的每个键Token $i$，计算其平均相似度分数$\bar{S}_i^h = \frac{1}{n} \sum_{j=0}^{n-1} S_{j,i}^h$。高平均值表明该Token的内容很大程度上与其他Tokens共享。最后，通过Softmax操作对$\bar{S}_i^h$进行归一化，得到限定数值范围的单Token冗余度分数$R_i^h$：

$$R_i^h = \left( \mathrm{softmax} \left( [ \bar{S}_0^h, \ldots, \bar{S}_{n-1}^h ] \right) \right)_i$$

KV缓存保留的联合选择策略：为了在有效管理缓存的同时保留基本上下文，R-KV整合了关键度分数和冗余度分数。给定每个注意力头的预算$B_{budget}$，目标是保留最大化信息多样性且最小化冗余的Tokens。每个Token $i$在头$h$中的最终选择分数$Z_i^h$计算如下：

$$Z_i^h = \lambda I_i^h - (1 - \lambda) R_i^h$$


其中，$I_i^h$越高表示Token越重要，$R_i^h$越高表示冗余度越大。超参数$\lambda$控制了优先保留重要Token与减少冗余Token之间的权衡。

实验环境

实验结果

补充细节

超参数$\lambda$的选择分析：
对顶层（$N_{layer}=31$）头$h=0$的关键度分数$\mathbf{I}^h$和冗余度估计$\mathbf{R}^h$分布分析表明，$\mathbf{I}^h$是稀疏的且由少数异常值主导，而相似度分布（决定$\mathbf{R}^h$）相对密集。当$\lambda=0$（纯冗余策略）时，无法保证保留初始的四个Tokens，这会严重损害LLM的生成能力（如先前工作指出），因此$\lambda$至少应为0.01。当$\lambda > 0.1$时，选择指标被注意力分数主导。实验（图6）证实，$\lambda=0.1$时在MATH-500上取得了最高准确率，而纯冗余（$\lambda=0$）或纯注意力（$\lambda=1$）策略性能最差，证明了两种指标的互补性。
不同λ下R-KV与纯注意力/纯冗余度指标的KV选择分数对比
与图5相同方法的性能对比

基于注意力的方法无法捕获冗余的原因：
通过对比R-KV与SnapKV选择的Tokens发现，SnapKV选择的Tokens覆盖范围有限，倾向于选择靠近Query的Tokens（局部注意力集中），且会选中远离Query但属于高度冗余和不重要的片段（如反复出现的“3 students are leaving early”）。相反，R-KV选择的Tokens更加多样化，分布更均匀，能够捕获更丰富的上下文表示。

效率与吞吐量分析：

结论

本文提出了R-KV，一种专为LLMs复杂推理挑战量身定制的解码期KV缓存压缩方法。针对推理模型生成过长且冗余输出导致的内存负担，R-KV通过联合评估Token的重要性和冗余度，保留了核心推理内容并丢弃了重复信息。该方法仅需10%-34%的原始KV缓存即可保持近乎完整的模型性能，大幅超越现有压缩方法。同时，R-KV在长序列生成场景下实现了最高13倍的Batch Size扩展和9倍的加速，作为一种免训练、模型无关的解决方案，为推理LLMs的部署（特别是强化学习工作流的Rollout阶段）提供了高可扩展性的支持。

附录

算法实现细节：
- GQA的最大池化：最新的开源LLMs（如Llama3、Qwen2）广泛采用分组查询注意力（GQA）。在KV缓存驱逐策略中，需要将注意力分数从Query头维度降维至KV头维度。先前工作（如SnapKV）主要采用平均池化（Mean Pooling），而R-KV假设最大池化（Max Pooling）能更好地为每个Query头保留最关键的Tokens。经验结果表明，最大池化能带来更好的性能，因此在所有主要实验中均采用此方法。

内存与计算复杂度分析：
- 内存节省公式：生成过程中，R-KV所需总内存为$M_{total} = M_{\theta} + M_{budget} + M_{buffer} + M_{\alpha}$（其中$M_{\theta}$为模型权重，$M_{\alpha}$为存储最后$\alpha$个Query状态的缓存）。相比之下，FullKV需要$M_{full}$来保留所有$B_{full}$个KV Tokens。因此，R-KV节省的内存为$M_{saving} = M_{full} - M_{budget} - M_{buffer} - M_{\alpha}$。
- 计算开销：关键度评分的计算复杂度为$O(\alpha B_{budget})$，冗余度估计的复杂度为$O(B_{budget}^2)$。每个生成段的总额外开销为$O(\alpha B_{budget} + B_{budget}^2)$。无压缩的生成复杂度为$O(B_{full} B_{buffer})$，而R-KV的复杂度降为$O((B_{budget} + B_{buffer}) B_{buffer})$。对于推理模型而言，$B_{full}$通常极大，使用相对较小的$B_{budget}$能有效降低计算成本，且减少KV缓存带来的注意力计算加速远超计算压缩分数的开销。