智源社区 17小时前
10% KV Cache实现无损数学推理!这个开源方法解决推理大模型「记忆过载」难题
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

R-KV是一种高效的KV缓存压缩方法,旨在解决大模型推理过程中产生的冗余问题。该技术通过实时对token进行排序,保留信息丰富且多样化的token,从而显著降低显存占用、提升吞吐量,并保持准确率。R-KV的核心在于“边生成边压缩”,通过重要性评估和冗余打分,联合淘汰冗余标记。实验结果表明,R-KV在数学基准测试中表现优异,并能支持更大的推理批处理大小,适用于边端设备、长链推理和强化学习等场景。

💡 R-KV的核心在于“边生成边压缩”,在模型解码时实时压缩KV缓存,仅保留重要且非冗余的标记,彻底阻断显存膨胀。

📊 R-KV通过三个关键步骤实现高效压缩:首先进行重要性打分,评估每个Token对后续答案的贡献度;其次进行冗余打分,计算Key向量余弦相似度,找出重复内容;最后联合淘汰,按「高重要+低冗余」优先级实时调度KV配额。

🚀 R-KV性能卓越:在数学基准测试中准确率不降反升,大幅超越基线;显存占用降低90%,吞吐量提升6.6倍。R-KV能够支持更大的推理批处理大小,适用于边端设备、长链推理和强化学习等场景。

推理大模型虽好,但一个简单的算数问题能推理整整三页,还都是重复的“废话”,找不到重点……

一种可以把大模型的“碎碎念”转化为可控记忆条目的高效压缩方法,出现了!

R-KV开源登场:显存↓90%、吞吐×6.6、准确率=100%

它可以通过实时对token进行排序,兼顾重要性和非冗余性,仅保留信息丰富且多样化的token,从而解决大模型推理时的冗余问题。

让“长时间推理”不再是奢侈品。

项目详情可见文末链接。

R-KV三步走:冗余识别+重要性评估+动态淘汰

链式思考(Chain-of-Thought,CoT)让LLM解题思路清晰可见,却也让推理长度指数级膨胀。

以DeepSeek-R1-Llama-8B为例,一道AIME数学题就能写出3.2万个Token:模型权重15.5GB,KV缓存再吃4.1GB——显存瞬间见底。

现有KV压缩方法(SnapKV、StreamingLLM、H2O等)主要针对长输入设计,可一旦模型在输出端开始“碎碎念”,相似句子之间互相打高分注意力,反而让“按注意力删低分”策略失灵:

造成关键步骤被误删、重复内容却被保留、准确率断崖式下跌等问题。

而R-KV通过以下步骤,在模型解码时实时压缩KV缓存来处理冗余的键/值(KV)标记,仅保留重要且非冗余的标记:

    边生成边压缩(Decoding-Time Compression)
    Token还没写进KV,就先判断“去留”,彻底阻断显存膨胀。
    重要性打分(Importance)
    多头注意力综合评估,每个Token对后续答案的贡献度。
    冗余打分(Redundancy)
    计算Key向量余弦相似度,找出“复读机”式内容。
    联合淘汰(Joint Eviction)
    按「高重要+低冗余」优先级实时调度KV配额,λ≈0.1时效果最佳。

整个流程训练-free模型-agnostic,无需改动模型结构,直接“即插即用”。因此可以直接被用到强化学习的采样过程中,非常灵活。

可视化:R-KV vs. SnapKV

上图展示了R-KV和纯注意力基线SnapKV在相同解码步骤中选择了哪些token。灰色=未选;由浅到深红=被越多注意力头选中。

可以看到,SnapKV关注点集中在离当前Query最近的局部片段,甚至重复保留多次「3 students are leaving early…」等无用自述。

而R-KV选出的Token横跨整段推理:题目关键词30 students,关键中间值24,12及最终答案全部被保留,此外语义覆盖面更广。

通过结合注意力强度与冗余过滤,R-KV保留了重要上下文并去除噪声,成功完成任务;而SnapKV误删关键信息导致答案错误。

得到结果:R-KV有更广泛的覆盖范围、更高的信息多样性和更显著的去冗余能力。

性能测试:准确率不降反升


可以看到,R-KV在具有挑战性的数学基准测试中大幅超越了基线,甚至超过了完整的KV。

在计算开销上,R-KV引入了重要性评分和冗余评分的额外计算,但总体开销适中,通常会被压缩KV缓存带来的注意力成本降低所抵消。随着序列长度的增加,这种权衡变得越来越有利。

对内存节省和端到端吞吐量提升进行实时分析,可以看到,当批处理大小为1时,R-KV在吞吐量上略优于FullKV。这表明R-KV通过减少注意力计算所实现的加速效果超过了R-KV自身的计算开销。

然而,这种直接的速度提升仅占整体收益的一小部分,R-KV带来的主要吞吐量提升来自于KV缓存压缩,使模型能够支持显著更大的推理批处理大小。

对基于比例和固定KV缓存预算的端到端吞吐量进行评估,发现R-KV始终能够实现比FullKV大得多的批处理大小和更高的吞吐量,同时不损失任务性能。

R-KV的适用场景如下:

    边端设备长链推理
    显存断崖缩减,让消费级GPU甚至手机NPU也能跑
    多轮Agent
    反思-重写-自评等复杂流程不再受显存限制。
    直接用于加速强化学习的采样过程
    training-free的方法即插即用。

论文PDF:https://arxiv.org/pdf/2505.24133.pdf
项目主页:https://zefan-cai.github.io/R-KV.page/
代码仓库:https://github.com/Zefan-Cai/R-KV

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法!

—  —

🌟 点亮星标 🌟

科技前沿进展每日见

内容中包含的图片若涉及版权问题,请及时与我们联系删除

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

R-KV 大模型 KV缓存压缩 推理优化
相关文章