MemDecay:區域感知的KV快取淘汰策略,提升LLM代理推理效率
研究提出MemDecay,一種訓練無關的區域感知KV快取淘汰策略,透過為不同語義區域的令牌分配不同優先順序和衰減率,在固定快取預算下保留關鍵資訊。實驗表明,系統令牌的半衰期遠長於暫存區令牌,且固定策略能在全部設定中保持完美準確率,而現有基線最多僅保留13/24。
近日,一篇來自arXiv的論文提出了MemDecay,一種無需訓練的區域感知KV快取淘汰策略,旨在解決大型語言模型(LLM)代理在推理過程中的記憶體瓶頸問題。LLM代理在複雜任務中需要累積系統指令、計劃、使用者輪次、檢索文件、工具輸出和中間推理等多種異構上下文,這些上下文的鍵值(KV)快取可能佔據大量視訊記憶體。傳統的快取淘汰策略通常對所有令牌一視同仁,基於注意力或最近性進行淘汰,但忽略了代理編排器已經擁有的語義結構資訊。
MemDecay的創新之處在於為不同語義區域的令牌分配不同的基礎優先順序和衰減率,並在令牌受到注意力時重新整理其保留分數。在固定的快取預算下,MemDecay會淘汰得分最低的頁面,同時允許關鍵區域(如系統指令)被固定(pinning),從而確保重要資訊不被丟失。此外,論文還提供了一種從測量的注意力生命週期中校準衰減率的方法,使得策略可以自適應地調整。
研究人員使用Qwen2.5-1.5B和3B模型,在約450和1,700個令牌的上下文中進行了評估。實驗結果表明,不同區域的注意力生命週期相差一個數量級:系統令牌的半衰期範圍為148到189個解碼步,而暫存區令牌僅為14到16步。透過固定系統區域,在所有設定下系統區域事實的準確率均達到了完全快取時的水平,而現有的基線方法最多隻能保留24個事實中的13個。隨著上下文增長,區域感知的保留策略仍然有效,而基於最近性的保留策略則迅速崩潰。不過,在非固定內容上,累積注意力保留策略表現更好,消融實驗表明注意力分數歸一化是當前公式的主要限制。這些結果清楚地表明,語義提示結構是KV快取管理的一個穩健訊號,同時也闡明瞭如何將其與基於注意力的重要性相結合。
MemDecay的提出為高效LLM代理推理提供了新的思路,有望在模型部署和推理成本最佳化中發揮重要作用。未來工作可以進一步探索如何將區域感知與其他注意力機制結合,以提升非固定內容的保留效果。