R-KV團(tuán)隊(duì) 投稿
量子位 | 公眾號(hào) QbitAI
推理大模型雖好,但一個(gè)簡(jiǎn)單的算數(shù)問(wèn)題能推理整整三頁(yè),還都是重復(fù)的“廢話”,找不到重點(diǎn)……
一種可以把大模型的“碎碎念”轉(zhuǎn)化為可控記憶條目的高效壓縮方法,出現(xiàn)了!
R-KV開(kāi)源登場(chǎng):顯存↓90%、吞吐×6.6、準(zhǔn)確率=100%。
它可以通過(guò)實(shí)時(shí)對(duì)token進(jìn)行排序,兼顧重要性和非冗余性,僅保留信息豐富且多樣化的token,從而解決大模型推理時(shí)的冗余問(wèn)題。
讓“長(zhǎng)時(shí)間推理”不再是奢侈品。
項(xiàng)目詳情可見(jiàn)文末鏈接。
R-KV三步走:冗余識(shí)別+重要性評(píng)估+動(dòng)態(tài)淘汰
鏈?zhǔn)剿伎迹–hain-of-Thought,CoT)讓LLM解題思路清晰可見(jiàn),卻也讓推理長(zhǎng)度指數(shù)級(jí)膨脹。
以DeepSeek-R1-Llama-8B為例,一道AIME數(shù)學(xué)題就能寫出3.2萬(wàn)個(gè)Token:模型權(quán)重15.5GB,KV緩存再吃4.1GB——顯存瞬間見(jiàn)底。
現(xiàn)有KV壓縮方法(SnapKV、StreamingLLM、H2O等)主要針對(duì)長(zhǎng)輸入設(shè)計(jì),可一旦模型在輸出端開(kāi)始“碎碎念”,相似句子之間互相打高分注意力,反而讓“按注意力刪低分”策略失靈:
造成關(guān)鍵步驟被誤刪、重復(fù)內(nèi)容卻被保留、準(zhǔn)確率斷崖式下跌等問(wèn)題。
而R-KV通過(guò)以下步驟,在模型解碼時(shí)實(shí)時(shí)壓縮KV緩存來(lái)處理冗余的鍵/值(KV)標(biāo)記,僅保留重要且非冗余的標(biāo)記:
- 邊生成邊壓縮(Decoding-Time Compression)
- Token還沒(méi)寫進(jìn)KV,就先判斷“去留”,徹底阻斷顯存膨脹。
- 重要性打分(Importance)
- 多頭注意力綜合評(píng)估,每個(gè)Token對(duì)后續(xù)答案的貢獻(xiàn)度。
- 冗余打分(Redundancy)
- 計(jì)算Key向量余弦相似度,找出“復(fù)讀機(jī)”式內(nèi)容。
- 聯(lián)合淘汰(Joint Eviction)
- 按「高重要+低冗余」優(yōu)先級(jí)實(shí)時(shí)調(diào)度KV配額,λ≈0.1時(shí)效果最佳。
整個(gè)流程訓(xùn)練-free、模型-agnostic,無(wú)需改動(dòng)模型結(jié)構(gòu),直接“即插即用”。因此可以直接被用到強(qiáng)化學(xué)習(xí)的采樣過(guò)程中,非常靈活。
可視化:R-KV vs. SnapKV
上圖展示了R-KV和純注意力基線SnapKV在相同解碼步驟中選擇了哪些token?;疑轿催x;由淺到深紅=被越多注意力頭選中。
可以看到,SnapKV關(guān)注點(diǎn)集中在離當(dāng)前Query最近的局部片段,甚至重復(fù)保留多次「3 students are leaving early…」等無(wú)用自述。
而R-KV選出的Token橫跨整段推理:題目關(guān)鍵詞30 students,關(guān)鍵中間值24,12及最終答案全部被保留,此外語(yǔ)義覆蓋面更廣。
通過(guò)結(jié)合注意力強(qiáng)度與冗余過(guò)濾,R-KV保留了重要上下文并去除噪聲,成功完成任務(wù);而SnapKV誤刪關(guān)鍵信息導(dǎo)致答案錯(cuò)誤。
得到結(jié)果:R-KV有更廣泛的覆蓋范圍、更高的信息多樣性和更顯著的去冗余能力。
性能測(cè)試:準(zhǔn)確率不降反升
可以看到,R-KV在具有挑戰(zhàn)性的數(shù)學(xué)基準(zhǔn)測(cè)試中大幅超越了基線,甚至超過(guò)了完整的KV。
在計(jì)算開(kāi)銷上,R-KV引入了重要性評(píng)分和冗余評(píng)分的額外計(jì)算,但總體開(kāi)銷適中,通常會(huì)被壓縮KV緩存帶來(lái)的注意力成本降低所抵消。隨著序列長(zhǎng)度的增加,這種權(quán)衡變得越來(lái)越有利。
對(duì)內(nèi)存節(jié)省和端到端吞吐量提升進(jìn)行實(shí)時(shí)分析,可以看到,當(dāng)批處理大小為1時(shí),R-KV在吞吐量上略優(yōu)于FullKV。這表明R-KV通過(guò)減少注意力計(jì)算所實(shí)現(xiàn)的加速效果超過(guò)了R-KV自身的計(jì)算開(kāi)銷。
然而,這種直接的速度提升僅占整體收益的一小部分,R-KV帶來(lái)的主要吞吐量提升來(lái)自于KV緩存壓縮,使模型能夠支持顯著更大的推理批處理大小。
對(duì)基于比例和固定KV緩存預(yù)算的端到端吞吐量進(jìn)行評(píng)估,發(fā)現(xiàn)R-KV始終能夠?qū)崿F(xiàn)比FullKV大得多的批處理大小和更高的吞吐量,同時(shí)不損失任務(wù)性能。
R-KV的適用場(chǎng)景如下:
- 邊端設(shè)備長(zhǎng)鏈推理
- 顯存斷崖縮減,讓消費(fèi)級(jí)GPU甚至手機(jī)NPU也能跑
- 多輪Agent
- 反思-重寫-自評(píng)等復(fù)雜流程不再受顯存限制。
- 直接用于加速?gòu)?qiáng)化學(xué)習(xí)的采樣過(guò)程
- training-free的方法即插即用。
論文PDF:https://arxiv.org/pdf/2505.24133.pdf
項(xiàng)目主頁(yè):https://zefan-cai.github.io/R-KV.page/
代碼倉(cāng)庫(kù):https://github.com/Zefan-Cai/R-KV
特別聲明:以上內(nèi)容(如有圖片或視頻亦包括在內(nèi))為自媒體平臺(tái)“網(wǎng)易號(hào)”用戶上傳并發(fā)布,本平臺(tái)僅提供信息存儲(chǔ)服務(wù)。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.