大型語言模型(LLM)在長上下文推理時面臨關鍵值(KV)緩存內存線性增長的嚴重瓶頸。隨着上下文長度增加,KV緩存佔用的顯存迅速膨脹,成為限制模型支持更長序列的主要障礙。現有的KV緩存壓縮方法大多依賴啓發式策略:一方面,啓發式預算分配基於統計先驗而非任務目標,導致不同注意力頭之間的資源錯配;另一方面,啓發式令牌選擇依賴於查詢-鍵交互或靜態歸納偏好(如注意力匯聚),這些方法無法動態適應具體任務的需求。這些問題使得現有方法在高壓縮率下性能顯著下降。
針對這一侷限,研究團隊提出了LKV(Learned KV Eviction)方法,將KV緩存壓縮重新表述為一個端到端可微優化問題。LKV包含兩個核心組件:LKV-H,用於學習任務優化的全局預算;以及LKV-T,用於推導令牌內在的重要性,而無需實例化注意力矩陣。這種設計完全繞過了手工設計的啓發式規則,使壓縮過程嚴格與任務目標對齊。LKV通過反向傳播直接優化壓縮效果,能夠根據具體任務自適應地調整各注意力頭的緩存容量和需要保留的令牌。
在廣泛的實驗評估中,LKV在LongBench和RULER兩個主流長上下文基準測試上均取得了最先進的性能。特別是,在LongBench上,LKV僅保留15%的KV緩存即可實現接近無損的性能,顯著優於此前的方法。這一結果驗證了LKV在高壓縮率下的有效性。進一步的分析表明,學習到的預算分配是保持壓縮後模型保真度的主要因素,證明數據驅動的分配對於克服手工啓發式的侷限性至關重要。
LKV的提出為LLM的高效推理提供了新的方向。它無需修改模型架構即可直接應用,並且能夠與現有的推理加速技術相結合。論文作者包括Enshuai Zhou等10位研究人員,論文於2026年4月22日提交至arXiv,並被歸類為機器學習(cs.LG)和計算與語言(cs.CL)兩個領域。該工作有望降低長上下文LLM的部署成本,推動其在長文檔分析、代碼生成等場景中的實際應用。