10% KV Cache實現無損數學推理!這個開源方法解決推理大模型「記憶過載」難題
推理大模型雖好,但一個簡單的算數問題能推理整整三頁,還都是重復的“廢話”,找不到重點……
一種可以把大模型的“碎碎念”轉化為可控記憶條目的高效壓縮方法,出現了!
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