36氪 - 科技频道 19小时前
10% KV Cache实现无损数学推理,这个开源方法解决推理大模型「记忆过载」难题
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

R-KV是一种高效的KV缓存压缩方法,专为解决大模型推理时的冗余问题而设计。它通过实时对token进行排序,保留信息丰富且多样的token,从而显著降低显存占用,提高吞吐量,并保持甚至提升准确率。R-KV无需改动模型结构,即插即用,适用于边端设备长链推理、多轮Agent以及加速强化学习采样等场景。实验结果表明,R-KV在数学基准测试中表现出色,性能优于现有KV压缩方法。

💡 R-KV的核心在于边生成边压缩(Decoding-Time Compression),在token写入KV缓存前就进行“去留”判断,有效阻止显存膨胀。

📊 R-KV通过重要性打分和冗余打分,综合评估每个token对后续答案的贡献度以及与其他token的相似度,从而筛选出关键且非冗余的token。

🚀 R-KV采用联合淘汰机制,优先保留高重要性、低冗余度的token,动态调整KV配额。实验表明,当λ≈0.1时效果最佳。

📈 性能测试显示,R-KV在数学基准测试中准确率不降反升,并显著提升了端到端吞吐量,尤其是在处理更大批处理大小时。

📱 R-KV适用于多种场景,包括边端设备长链推理、多轮Agent和加速强化学习采样,使得大模型推理在资源受限的设备上成为可能。

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

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

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)标记,仅保留重要且非冗余的标记:

Token还没写进KV,就先判断“去留”,彻底阻断显存膨胀。

多头注意力综合评估,每个Token对后续答案的贡献度。

计算Key向量余弦相似度,找出“复读机”式内容。

按「高重要+低冗余」优先级实时调度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也能跑

反思-重写-自评等复杂流程不再受显存限制。 

直接用于加速强化学习的采样过程

论文PDF:https://arxiv.org/pdf/2505.24133.pdf

项目主页:https://zefan-cai.github.io/R-KV.page/

代码仓库:https://github.com/Zefan-Cai/R-KV

本文来自微信公众号“量子位”,作者:R-KV团队,36氪经授权发布。

Fish AI Reader

Fish AI Reader

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

FishAI

FishAI

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

联系邮箱 441953276@qq.com

相关标签

R-KV 大模型 推理加速 显存优化
相关文章