作者/机构: Jiayi Tian, Seyedarmin Azizi, Yequan Zhao, Erfan Baghaei Potraghloo, Sean McPherson, Sharath Nittur Sridhar, Zhengyang Wang, Zheng Zhang, Massoud Pedram, Souvik Kundu
本文旨在解决大型推理模型(LRMs)在进行思维链(CoT)推理时,因生成冗长轨迹而导致的键值(KV)缓存开销巨大的问题。这一开销构成了内存和吞吐量的瓶颈,限制了模型的有效部署。
核心问题与研究目标:
现有KV缓存压缩方法在处理长CoT任务时效果不佳。例如,H2O、SnapKV等方法主要针对长上下文预填充(prefill)场景,在Co-T推理任务中准确率下降明显。近期提出的R-KV虽关注到CoT中的冗余token,但在多批次(multi-batch)推理时准确率会显著下降,且其逐token的驱逐策略破坏了语义结构,常导致模型生成更长的推理路径。此外,多批次推理中的填充(padding)token会挤占有效的KV预算,导致性能恶化。
创新点与主要贡献:
为了解决上述问题,本文提出了SkipKV,一个无需训练的、句子感知的KV压缩框架,通过在粗粒度的句子层面进行选择性驱逐和生成,实现高效的CoT推理。SkipKV不仅保持了推理的连贯性,还在固定的KV缓存预算下,实现了准确率和生成长度之间的更优权衡。其核心贡献包括:
成果概述:
在多个推理基准(AIME-24, LiveCodeBench, MATH-500, GSM8K)上的广泛评估表明,SkipKV效果显著。与现有方法相比,在相似的压缩预算下,SkipKV的准确率最高提升了26.7%。与SOTA方法相比,SkipKV生成的序列长度最多减少了1.6倍,吞吐量最多提升了1.7倍。如下图1所示,SkipKV在较小的KV预算下,既保持了高准确率,又缩短了生成长度。
下图2直观对比了不同的KV缓存驱逐策略。SkipKV通过分组句子内的token并驱逐句子冗余度高的区域,实现了高准确率和更短的生成长度。
本节分析了代表性的SOTA CoT token驱逐方法的关键局限性。
观察1:在多批次解码中,进行KV驱逐会导致推理准确率相较于单批次解码有所下降。
KV缓存驱逐方法的优势在于能够减少KV内存使用,从而支持更大的批次大小并提高生成吞吐量。为了分析现有SOTA KV驱逐方法(如【Cai et al., R-kv: Redundancy-aware kv cache compression for trainingfree reasoning models acceleration. 2025】和【Zhang et al., H2o: ´ heavy-hitter oracle for efficient generative inference of large language models. 2023】)在CoT推理任务上的准确率鲁棒性,我们在批次大小为1和10的情况下分别评估了准确率。具体而言,我们在MATH-500数据集上使用R1-Llama-8B模型,在不同KV预算下(以基线FullKV的KV内存为100%)进行了测试。如图3(左)所示,对于R-KV和H2O,多批次解码导致了显著的准确率下降,尤其是在较低的KV预算下。这一局限性可能归因于多批次场景下,填充(padding)token导致有效KV预算的减少。这些填充token是为了对齐批次中不同候选序列的可变长度prefill序列而添加的。当通过KV驱逐实施固定的内存预算时,填充token会消耗内存空间,导致有效KV预算减少,从而可能损害准确率。此外,这些token还可能扭曲注意力分布,导致不稳定的token级重要性估计,进一步加剧推理一致性的下降。如图3(中)所示,MATH-500中的prefill token长度在样本间差异很大,一个批次内的差异可能超过400个token。这种差异导致序列批处理时产生过多的填充,增加了KV缓存的浪费,并加剧了带驱逐的多批次解码中的性能下降。
观察2:在KV预算减少时,总生成长度通常会比没有任何KV压缩时增加。
图3(右)显示了在MATH-500上评估的Deepseek-R1-Qwen 7B(批次大小为10)的准确率和平均token长度(连续线)。除了影响准确率外,与没有压缩的模型相比,使用R-KV压缩的模型在所有KV预算下都持续生成更长的序列。这一观察表明,token驱逐可能导致模型延长其推理链并生成更长的输出。这可能归因于KV缓存中上下文信息的丢失:当部分有效上下文被移除时,模型可能会通过重新生成推理步骤或添加冗余片段来补偿缺失的信息。更长的生成导致了因前向传播次数增加而带来的计算成本增加,从而抵消了通过缓存大小限制实现的效率增益。
观察3:Token级别的驱逐常常导致单词的碎片化移除,从而引起大型推理模型(LRM)的过度思考。
为了进一步探究生成长度增加的原因,我们在图4中可视化了经过R-KV驱逐后保留的token。我们从MATH-500中选择一个样本,并展示了模型输出,其中保留的token用蓝色高亮,被驱逐的用灰色表示。这里的局限性可归纳为两方面。首先,纯粹基于token级冗余分数【Cai et al., R-kv: Redundancy-aware kv cache compression for trainingfree reasoning models acceleration. 2025】的驱逐常常会从关键的数学计算步骤中移除数字,从而扰乱推理流程。其次,token级别的驱逐可能常常保留正确答案的碎片化token,误导LRM重新验证部分结果并生成不必要地长或不确定的推理链。例如,在图4中,R-KV保留的token频繁地包括来自中间推理步骤的不相连的数字实体(例如,6 + 12i − 3i → +2i)或来自最终答案的片段(例如,(6, 9) → (, 9))。这种碎片化的保留导致模型反复重新推导或检查答案的某些部分,从而产生更长但冗余的推理轨迹。这些发现强调了需要一种在更高层次语义单元(例如句子或推理段落)上操作的连贯驱逐策略,以实现准确率和生成长度之间更好的权衡。
本节首先分析了推理轨迹的句子级属性,这些属性启发了SkipKV的设计(如图6所示)。
定义:成对句子相似度(Pairwise Sentence Similarity, PSS)
两个句子之间的PSS定义为两者向量嵌入之间的余弦相似度。设两个句子的向量嵌入分别为$v_i, v_j \in R^d$,PSS得分计算如下:
![][351b88f7802cbec6c4fa6b2f7bfe24836f9294a7cb25e4a07aad191e8a1bb10b]
为了从语义角度理解成功和失败的推理轨迹之间的差异,我们对模型使用FullKV生成的输出中的句子(思考)级属性进行了细粒度分析。我们关注两个方面:(i)生成句子之间的语义相似性,以及(ii)非执行性(不太重要)思考的比例。值得注意的是,遵循【Chen et al., Seal: Steerable reasoning calibration of large language models for free. 2025】的定义,我们将生成的token分为执行性(与实际答案相关的重要token)和非执行性(为通过反思或过渡来验证响应而生成的次要token)思考。
观察4:正确和错误的推理响应都会生成高度相似的句子,后者通常生成更高比例的相似句子。
我们在AIME24上使用R1-Qwen-7B和R1-Llama-8B进行了一项计算PSS的实验。具体来说,我们首先收集模型输出文本,并将其分割成由换行符和标点符号模式(例如,‘\n’, ‘.\n’)分隔的单个句子。然后使用基于BERT的句子转换器【Reimers & Gurevych, Sentence-bert: Sentence embeddings using siamese bert-networks. 2019】将每个句子编码为有意义的语义嵌入向量。接着,我们计算所有句子嵌入对的PSS,并将PSS得分≥0.95的对标记为高度相似,从而得到语义冗余的度量。有趣的是,如图5(上)所示,与正确响应相比,错误响应始终表现出高达1.7倍的更高比例的相似句子。
观察5:与正确响应相比,错误响应生成的非执行性思考比例要高得多。
如图5(下)所示,与正确响应相比,错误的生成始终表现出更高比例的非执行性思考。特别是在AIME24上分别使用R1-Qwen-7B和R1-Llama-8B进行评估时,我们观察到错误输出中的非执行性思考分别比正确输出高约2.6倍和1.8倍。这些模式表明,当模型失败时,它们倾向于产生重复或停滞的推理步骤,重访语义相似的内容或生成元级别的评论,而不是执行具体的问题解决行动。
句子相似度评分与token重要性/冗余度评分结合
我们首先介绍句子级的相似度评分,然后提出累积评分机制来决定要驱逐的句子和token。
句子相似度得分。我们使用此分数来计算句子级别的冗余度。然而,使用句子转换器来生成句子的潜在向量表示是不切实际的,因为它会在解码过程中带来显著的计算开销。因此,我们利用最后一层的隐藏状态,表示为$H \in R^{bs \times N \times d}$,作为句子片段的潜在上下文表示。具体来说,对于批次中的每个样本,我们通过将总序列(N)划分为由基于标点的分隔符分隔的小块来识别每个句子片段的开始和结束索引(图6中的步骤1-2)。对于批次ID为k的句子i,我们然后计算其向量嵌入$v_i \in R^{1 \times 1 \times d}$作为该句子中向量的均值,
![][556b950ab1bd497c4e307f0889dafd0a6eda83af75c7fd9445b17c1c4ca6f3c5]
这里,$b_i$和$e_i$是句子i的开始和结束索引。然后我们根据公式1计算PSS。接着,我们将冗余句子集定义为
![][461dadf14f4bbf37b20b91fb09ac8eebae585db23f189b1c9546ead2b9d89094]
这意味着如果一对(i, j)超过了预定义的阈值$\tau$(例如0.95),则较早的句子i被标记为冗余,而较晚的句子j被保留。
Token重要性得分。令$Q \in R^{bs \times H_q \times \alpha \times d}$表示最近$\alpha$个查询token的观察窗口【Li et al., Snapkv: Llm knows what you are looking for before generation. 2024】,$K, V \in R^{bs \times H_k \times N' \times d}$表示当前的键和值缓存。这里,$H_q$和$H_k$分别表示查询头和键值头的数量。大多数LRM使用分组查询注意力,其中对于每个查询头$i \in \{1, . . . , H_q\}$,其关联的键值头由$g(i) \in \{1, . . . , H_k\}$表示,其中$g(i) = \lfloor i/n \rfloor$。n是一个整数,定义了每个K/V被共享的查询头数量。批次中每个查询头i的注意力重要性表示为$A^i \in R^{bs \times 1 \times \alpha \times N'}$,其中
![][ae65b2d026a55c484e72432e3a003ca9cd634d8cf419a5cf98160f76f6210374]
这里,M表示用于多批次解码的注意力掩码。然后我们计算一个键头$h_k$的注意力重要性为$I^{h_k} = \text{softmax}(\text{maxpool}(A^{h_k \cdot n}: A^{h_k \cdot n+(n-1)}))$。我们对每个头的观察窗口$\alpha$上的注意力重要性矩阵进行归一化,得到token重要性矩阵$I^{h_k}_{\alpha} \in R^{bs \times N'}$。
Token冗余度得分。受R-KV的启发,我们现在总结token冗余度移除方法。对于每个键值头$h_k$,我们将键状态定义为$K^{h_k} \in R^{bs \times 1 \times N' \times d}$,token冗余度$R^{h_k} \in R^{bs \times N'}$计算为
![][19ec082b30261baf875428028572c0f234aae1d1ee8e529b57463ce26027acf0]
这里,我们也考虑注意力掩码M来减少填充token对冗余度得分的影响。
句子冗余驱动的累积得分。然后我们通过结合token重要性得分$I^{h_k}_\alpha$,token冗余度得分$R^{h_k}$,和句子相似度得分$\lambda_{i,j}$来制定KV压缩的总体驱逐得分(图6中的步骤3b):
![][7a8e0c15c6b69c3d94050f762e94aeecba6834a47165098ccfb5b961bdf1e64c]
$\sigma$控制了token重要性和冗余度之间的权衡。如前所述,$\lambda_{i,j}$是与句子对(i, j)相关的句子冗余度得分,由第i个句子的所有token共享。为了满足特定的token预算,我们按照$I_{final}$的升序驱逐token。重要的是,由于冗余句子的相似度得分(≥0.95)通常比token级别的得分(~0.1)高一个数量级,公式6确保了高度冗余的句子在token级别驱逐之前被移除。
KV缓存句子范围监控逻辑。为了确保冗余集P中的句子能从缓存空间中被一致地驱逐,我们引入了一个KV缓存范围监控机制(图6中的步骤3a和3c)。该机制由一个映射函数$\Phi$定义,它将原始生成空间(gs)中的第i个句子跨度转换为缓存空间(cs)中的跨度,即$\Phi(b^{(gs)}_i, e^{(gs)}_i) \rightarrow (b^{(cs)}_i, e^{(cs)}_i)$。为了正确地为P中的每个句子关联句子级得分,我们将token跨度ID$(b^{(gs)}_i, e^{(gs)}_i)$和相应的PSS记录在一个查找表T中,如下:
![][84649861b09c5b1546b6d0d77e10eb005432e900e3e5e40f4203d3f807c0afe6]
通过应用映射函数$\Phi$,我们获得了与缓存对齐的查找表:
![][a5730b3762c1a1a805b98479593a8d4f6afd9a4ff5dbb444d93a69b37e9bee25]
然后,在公式(6)中使用与缓存对齐的查找表$\Phi(T)$来确定在驱逐期间$i \in P$是否成立。如图7所示,句子范围监控逻辑需要在整个压缩过程中跟踪生成空间和缓存空间之间的映射。压缩步骤1后剩余的token成为压缩步骤2的cs的初始token,映射负责处理gs到cs的索引。
驱逐前更新缓存范围。由于我们每隔一定数量的解码迭代执行一次驱逐,在一次驱逐步骤中,我们可能会看到新生成的token——通常是完成或开始句子——被附加到gs token上。因此,它们的句子跨度必须被添加到KV缓存中,遵循两种情况:(1)在第一步初始化句子范围(图7,顶部),和(2)在后续步骤中附加新的范围(图7,底部)。在情况(1)中,缓存范围被简单地初始化为生成空间中的范围:
![][2a7ad97165b5ba14076303ecd3b52625d9fa6df9a631a595f938797d77cb469f]
在压缩步骤t,我们用$l^{(gs)}$和$l^{(cs)}$表示驱逐前的生成和缓存序列长度,而驱逐后的缓存长度受预算B的限制。在情况(2)中,对于当前压缩步骤中添加的每个新句子,我们按句子索引的递增顺序依次更新句子范围。对于第i个句子,它是:
![][0af3f8a324b4e5f97a52fc5adc9313c5aaac110f1d69d9fac96863ca24fb7913]
结束位置$e^{(cs)}_i$是通过从当前缓存长度$l^{(cs)}_t$中减去偏移量$\Delta$来确定的。这里的$\Delta$表示从第(i+1)个句子开始的gs中剩余token的长度。这个偏移量等于新附加token的长度,可以计算为总生成长度减去第i个句子在生成空间中的结束位置。
驱逐后更新缓存范围。如图7所示,句子跨度必须根据每个压缩步骤中幸存的token索引重新映射到新的缓存坐标空间。设幸存token索引的集合为$P = \{p_1, p_2, . . . , p_B\}$,其中$0 \le p_1 < p_2 < · · · < p_B < l^{(cs)}_t$,$p_k$表示第k个幸存token的索引。相应的缓存范围则更新为:
![][1b2f479b305c10201decc2982c16e3514d566d23ee546d84874f605e62194e54]
这里,$b^{(cs)}_i$被重新分配为不早于原始起点的最早幸存索引,而$e^{(cs)}_i$则被分配为不晚于原始终点的最晚幸存索引。如果在压缩后的缓存空间中没有token幸存,相应的句子范围将从驱逐后的缓存空间中丢弃。
自适应引导机制
除了在生成后移除冗余思考,我们进一步提出通过一种自适应引导机制在生成前跳过不必要的思考。为此,我们在某些LRM层的隐藏状态上添加一个引导向量,强制模型变得精确。然而,与早期使用固定强度($\alpha_0$)引导向量的引导工作【Chen et al., Seal: Steerable reasoning calibration of large language models for free. 2025】不同,我们提出了一种基于非执行性思考计数的强度自适应方法。下面介绍了自适应引导的细节。
动态调整引导强度
受SEAL【Chen et al., Seal: Steerable reasoning calibration of large language models for free. 2025】的启发,我们首先使用来自MATH训练集的500个样本构建引导向量,旨在将潜在表示转向执行式推理。引导向量计算为执行性和非执行性思考之间的平均潜在表示差异:$V = H_E − H_O$,其中$H_E$和$H_O$分别表示执行性和非执行性思考的平均隐藏状态。在推理步骤t中,引导向量被注入到所选层k的隐藏表示中,即$H_k \leftarrow H_k + \alpha_t \cdot V$,其中$\alpha_t$表示引导强度,控制着向执行导向方向的潜在调整程度。我们根据模型观察到的推理行为对$\alpha_t$进行动态调整。具体来说,如图6(左)所示,我们维持一个非执行性思考的运行计数器$N_o$,并更新$\alpha_t$为$\alpha_t \leftarrow \alpha_0 + \gamma \cdot N_o$,其中$\alpha_0$和$\gamma$分别是初始引导强度和预定义的增量因子,用于控制引导的激进程度。通过这种方式,每个样本的引导强度在整个推理过程中自适应调整,使模型能够通过选择性地抑制不必要的思考来缩短生成长度。
减少填充token以增加有效KV预算
对于批次中的每个样本,其有效的KV预算由$B' = B - \Delta_{pad}$给出,其中$\Delta_{pad}$是该样本的填充token数量。对于任何样本,$\Delta_{pad}$的计算公式为:
![][cdd0c2821ef36a41f8dd99350d1fb8c4bd11d3b2a68a3384774995a66d55bd9e]
其中$N_p$和$N^{max}_p$分别代表当前样本的预填充(prefill)大小和批次中最大预填充长度。显然,对于预填充长度差异很大的批次,$\Delta_{pad}$可能会很高。因此,为了减少浪费的填充数量,我们提出了批次分组策略。具体来说,我们首先按预填充长度的升序对所有样本进行排序和重新排列。然后,我们使用这个排序后的样本空间将它们分组为大小为bs的批次。如公式12所示,由于$N^{max}_p$和$N_p$之间的差异减小,填充token数量减少,从而增加了$B'$。图8展示了一个批次分组中的样本重排示例。在实验章节5.3中,我们通过经验验证了增加的$B'$在提高LRM多批次解码准确率方面的有效性。
模型:
数据集:
硬件配置:
软件与实验设置:
与驱逐方法的准确率比较:
如图9所示,在多个推理数据集和模型规模上,SkipKV与H2O、R-KV等先前的KV缓存驱逐方法进行了比较。与以往在多批次解码中准确率严重下降的token级驱逐方法不同,SkipKV在所有推理任务和模型上,都能以显著更低的KV内存保持高准确率。例如,在挑战性的AIME-24任务中,SkipKV在R1-Qwen-14B上使用比FullKV低6.7倍的KV缓存内存,达到了与之相同的准确率,并且在所有模型上都保持了超过4倍的KV内存节省。此外,在有限的KV预算下,SkipKV的性能甚至优于FullKV,例如在LiveCodeBench上,用2倍少的KV内存,准确率提升了5.2%(R1-Qwen-7B)。
与驱逐方法的Token长度比较:
如图10所示,SkipKV不仅准确率更高,还通过生成更少的token显著提升了生成效率。先前的token级驱逐方法在所有模型和任务上生成的token都比FullKV多,而SkipKV与FullKV相比,总token长度最多减少了28%。与R-KV相比,SkipKV在R1-Qwen-7B、R1-Qwen-14B和R1-Llama-8B上分别减少了高达32%、39%和48%的token,相当于生成延迟降低了1.5-2倍。例如,在MATH-500数据集上,SkipKV在R1-Qwen-7B上生成的token减少约15%,同时使用的KV内存比FullKV少3倍,实现了3.5倍的整体内存-延迟效益。
与其他高效推理方法的比较:
如图11所示,本文将SkipKV与专注于减少生成长度的SEAL方法进行了比较。虽然SEAL能缩短推理轨迹并保持准确率,但其输出token长度的减少(约10%)仅带来了有限的KV内存节省。相比之下,SkipKV通过联合跳过KV的生成和存储,与SEAL相比,准确率最多提升13.3%,KV内存消耗减少了6.6倍,有效地平衡了内存压缩和推理保真度。
吞吐量分析:
如表1所示,在GSM8K数据集上对不同方法的端到端吞吐量进行了评估。FullKV和SEAL支持的批次大小远小于基于驱逐的方法,导致吞吐量可扩展性有限。相比之下,R-KV和SkipKV通过固定大小的KV缓存有效缓解了内存压力,支持的批次大小最多增加了2.8倍。总体而言,R-KV比FullKV实现了高达7.6倍的加速,而SkipKV比FullKV进一步加速了9.6倍。在相同批次大小下,得益于更短的生成长度,SkipKV的性能比R-KV高出1.7倍。
表1. 在GSM8K上不同批次大小下的总延迟和吞吐量比较。在单张A100-40GB上评估。
SkipKV组件分析:
从表2中可以看出,在AIME24基准上,逐步集成SkipKV的三个组件——句子评分、自适应引导和批次分组——在推理效率和准确率上都带来了持续的改进。句子评分通过移除冗余句子提供了适度增益;加入自适应引导通过动态跳过非执行性思考显著减少了token总长度;最后,结合批次分组构成的完整SkipKV配置,与R-KV基线相比,准确率最高提升了20%,token总长度最多减少了30%,证明了SkipKV分层设计的有效性。
表2. 在AIME24上使用R1-Qwen-7B对SkipKV进行的消融实验。与R-KV基线相比,逐步引入句子评分、自适应引导和批次分组可以提高准确率并减少总token长度。
批次分组对有效KV预算的影响:
如图12所示,MATH-500数据集的prefill token长度在样本间差异很大,导致批次内变异性大。通过排序和分组后,批次内的prefill长度变异性显著降低,从而减少了所需的填充量。如表3所示,批次分组使得几乎全部KV预算都分配给了有效token,从而在较低的KV预算下有效保持了准确率。
表3. 在MATH-500上使用R1-Qwen-7B和SkipKV评估批次分组的效果。
本文介绍了SkipKV,一个面向句子的KV压缩框架,旨在提升推理效率的同时保持准确性。它通过选择性地跳过KV的生成和存储来降低内存占用。受LRM输出中句子级结构经验观察的启发,SkipKV引入了句子优先的KV驱逐策略和句子类型自适应的引导向量,以实现更连贯、更高效的生成。为了在多批次解码中进一步实现鲁棒的准确性,本文提出了一种批次分组策略来优化有效KV预算的分配。与SOTA替代方案R-KV相比,SkipKV在相似的压缩预算下,准确率提升高达26.7%,同时生成的token数量减少了1.6倍。此外,在多批次设置中,SkipKV与基线FullKV LRM相比,在相似的准确率下,吞吐量提升高达9.6倍。
我们在算法1中提供了SkipKV存储跳过机制的伪代码。假设总共有n个句子,我们将生成空间中的句子跨度集合及其在KV缓存中对应的范围定义为
![][68797e7355ebeb1e8c1867d7db5ae572ccad98989f3c1f322bc85a7c20e276e1]
其中每个$(b^{(gs)}_i, e^{(gs)}_i)$表示生成空间中第i个句子的token跨度,而$(b^{(cs)}_i, e^{(cs)}_i)$表示其在缓存中的相应范围。我们通过算法1获得压缩后的$K_{cache}$, $V_{cache}$以及更新后的缓存空间句子范围$M^{(cs)}$。
算法1 跳过KV缓存存储
需要: 驱逐前的KV缓存长度$l^{(cs)}_t$,KV缓存预算$B$,句子数量$n$,生成和缓存空间中的句子范围$M^{(gs)}, M^{(cs)}$,token跨度-PSS得分查找表$T$。
确保: 压缩后的缓存$K_{cache}, V_{cache}$,更新后的缓存范围$M^{(cs)}$。
1:
2: # 驱逐前更新
3: # 情况(1): 初始化空缓存范围。
4: if $M^{(cs)}$ 为空 then
5: for 每个句子范围$M^{(gs)}_i$ where $e^{(gs)}_i \le l^{(cs)}_t$ do
6: 使用公式(9)初始化$M^{(cs)}_i$。
7: end for
8: # 情况(2): 附加新句子范围。
9: else
10: for 每个句子范围$M^{(gs)}_i$ where $i \ge n$ do
11: 使用公式(10)更新$M^{(cs)}_i$。
12: if $e^{(cs)}_i \le l^{(cs)}_t$ then $M^{(cs)} \leftarrow M^{(cs)} \cup \{M^{(cs)}_i\}$。
13: end for
14: end if
15:
16: # KV驱逐
17: 通过公式(8)将冗余得分映射到缓存范围$\Phi(T)$;通过公式(6)以固定预算$B$驱逐冗余token并更新$K_{cache}, V_{cache}$。
18:
19: # 驱逐后更新
20: 对于$M^{(cs)}$中的每个$M^{(cs)}_i$,使用公式(11)进行更新。
21: 更新句子计数$n \leftarrow \text{len}(M^{(gs)})$。
完整的SkipKV过程在算法2中详细说明。为清晰起见,我们定义SkipKV使用的辅助token集如下。换行分隔符集定义为
![][f14533cfbf634d46934dd1b43ae8f330b9b224e4ec968999369ce61cca58b055]
非可执行关键字集定义为
![][7621f80387df3e415032e8bb026a08b424034423e70d70461f4db5a90707ac78]
以简化解码期间token级别的思考类型检测。这些集合在解码前一次性初始化,并用于记录带标签的输入句子范围。
算法2 SkipKV算法
需要: 大型推理模型$f(\theta, \cdot)$,输入内容$X$,KV缓存预算$B$,压缩步长间隔$\Delta t$,换行分隔符集$D$,非可执行关键字集$N$,引导层索引$L_s$,引导增量因子$\gamma$,初始引导强度$\alpha_0$,引导向量$V$,最大生成长度$N_{max}$。
确保: 生成的文本$Y$。
1: while $t < N_{max}$ do
2: # 1-2. 记录带标签的生成句子范围
3: $b^{(gs)}_0 \leftarrow 0$
4:
5: for $x_i$ in $X$ do
6: if $x_i \in D$ then $M^{(gs)}_i \leftarrow (b^{(gs)}_i, x_i)$
7: if $(\forall t \in [b^{(gs)}_i, x_i]) \in N$ then
8: 用'Others'标记$M^{(gs)}_i$
9: end if
10: $b^{(gs)}_{i+1} \leftarrow x_i + 1$
11: end for
12:
13: # 3. 跳过KV存储
14: $K_{cache}, V_{cache}, y \leftarrow f(\theta, X)$
15: for $k$ in 解码器层$L$ do
16: if $t \pmod{\Delta t} == 0$ then
17: 使用算法1更新$K_{cache}, V_{cache}$和$M^{(cs)}$
18: end if
19: # 跳过KV生成
20: if $k == L_s$ then $H_k \leftarrow H_k + \alpha_t \cdot V$
21: end for
22:
23: # 4-5. 计算PSS
24: 计算$M^{(gs)}$中句子的冗余分数,并使用公式(7)记录冗余集$T$。
25:
26: # 更新引导强度
27: $N_o \leftarrow \text{count}(M^{(gs)}_i\text{.key()} == \text{'Others'}, \forall i)$
28: $\alpha_t \leftarrow \alpha_0 + N_o \cdot \gamma$
29:
30: # 自回归更新
31: $y \rightarrow Y$; $y \rightarrow X$; $t \leftarrow t + 1$
32: end while
超参数。遵循R-KV,我们每128个解码步骤压缩和更新一次KV缓存,并将注意力-冗余度得分权衡因子设置为$\sigma = 0.1$。句子评分指标中的相似度阈值$\tau$选择在0.95到0.99的范围内。引导强度$\alpha$初始化为1或1.25,引导强度增量因子$\gamma$设置为0.02。我们在R1-Qwen-7B和R1-Llama-8B的第20层,以及R1-Qwen-14B的第35层插入引导向量。
基线。我们将我们的方法与基于驱逐的KV压缩方法H2O【Zhang et al., H2o: ´ heavy-hitter oracle for efficient generative inference of large language models. 2023】、RKV【Cai et al., R-kv: Redundancy-aware kv cache compression for trainingfree reasoning models acceleration. 2025】,以及基于引导的高效推理方法SEAL【Chen et al., Seal: Steerable reasoning calibration of large language models for free. 2025】进行比较。FullKV作为保留完整KV缓存的参考方法,为解码准确率提供了黄金标准。
图13显示了在MATH-500、AIME-24和LiveCodeBench数据集上,R1-Llama-8B模型的准确率和总token长度,比较了SkipKV与先前的驱逐策略(包括H2O和R-KV)以及FullKV基线。
图14比较了不同KV驱逐策略在使用R1-Qwen-7B模型于AIME-24、LiveCodeBench和MATH-500上的推理行为。对于R-KV和SkipKV,我们在每个数据集上的KV预算分别设置为2220 (20%)、2000 (28%)、1024 (30%),在此预算下SkipKV达到了与FullKV相同的准确率。图的上半部分展示了非执行性思考的比例,反映了不必要的重新验证步骤的频率;而下半部分报告了高相似度句子的比例,这表明了重复的推理模式。
如第4节观察4和5所讨论,正确和错误样本之间的对比显示,非执行性和冗余推理在错误的生成中积累得更严重。在两个面板中,SkipKV有效地缩小了正确和错误样本之间的差距,在两种情况下都减少了不必要或重复思考的生成。在数量上,与FullKV相比,SkipKV在正确和错误的生成中产生的非执行性思考分别减少了约4倍和8倍,并进一步分别超过SEAL 1.8倍和1.3倍。此外,得益于其样本级的自适应引导机制,SkipKV降低了非执行性思考比例在样本间的方差。尽管R-KV旨在移除冗余token,但它仅将错误答案中的相似句子减少了高达11%,而明确设计用于针对冗余句子的SkipKV始终实现了超过2倍的减少。总体而言,SkipKV在受限的KV预算下,在保持简洁、面向执行的推理方面表现出明显优势。
本节我们展示了来自MATH-500数据集的定性示例,以比较R-KV和SkipKV的生成行为(图15和16)。我们可视化了R1-Qwen-7B推理模型在KV缓存预算为1024下的生成输出。在所选示例中,两种方法都产生了正确的最终答案;然而,SkipKV生成的token比R-KV少了约20%。在两张图中,非执行性思考用黄色高亮,最终答案片段用蓝色框高亮。红色文本的深浅表示选择每个token的注意力头数量。从驱逐模式的可视化中可以看出,SkipKV主要从KV缓存中移除完整的句子,而R-KV则倾向于驱逐碎片化的token。R-KV基于冗余度的评分常常导致它移除位于关键数学推理路径上的数值token,从而破坏了逻辑一致性并导致更长的生成。值得注意的是,R-KV还驱逐了最终答案区域的token,这会迷惑模型并触发不必要的重新验证步骤。相比之下,SkipKV保留了关键的推理和答案片段,表现出更连贯和高效的生成行为。