المفاهيم الأساسية
大型語言模型 (LLM) 的效能提升需要大量的記憶體資源,尤其是在處理長文本時。文章提出了一種基於注意力頭重要性的 KV 快取壓縮方法 (HeadKV),透過評估每個注意力頭對上下文問答任務中檢索和推理能力的貢獻,將有限的 KV 快取資源動態分配給更重要的注意力頭,從而在不影響效能的前提下顯著減少記憶體使用並提升運算效率。
研究背景
隨著大型語言模型(LLM)的發展,模型能夠處理越來越長的輸入文本。然而,隨著輸入長度的增加,自注意力機制所需的記憶體使用量和延遲也顯著增加,其中關鍵值(KV)快取的記憶體佔用量尤為突出。為了解決這個問題,現有的研究提出了各種 KV 快取壓縮技術,但大多數方法依賴於層級分配策略,無法充分優化 KV 快取分配以應對下游任務。
研究目標
本文旨在提出一個更有效的 KV 快取壓縮方法,透過在注意力頭級別進行動態分配,將有限的 KV 快取資源集中於對下游任務貢獻更大的注意力頭,從而在不影響效能的前提下顯著減少記憶體使用並提升運算效率。
研究方法
本文提出的 HeadKV 方法包含三個關鍵組成部分:
注意力頭級重要性評分估計:
借鑒 Needle-in-a-Haystack 測試方法,但不同於以往僅關注檢索能力,本文設計了包含顯式上下文推理步驟的評估範例,以評估每個注意力頭對檢索和推理能力的貢獻。
修改了重要性評分計算方式,不再僅依賴於注意力分數最高的 token,而是考慮整個正確答案,並使用注意力權重進行調整,以獲得更準確的注意力頭重要性分佈。
注意力頭級 KV 快取分配:
根據每個注意力頭的重要性評分,動態分配 KV 快取預算。重要性評分越高的注意力頭,分配到的 KV 快取預算越大。
具體而言,每個注意力頭會被分配一個固定的基礎預算,並從所有注意力頭中提取一部分預算形成共享預算池。共享預算池中的預算會根據每個注意力頭的重要性評分進行分配,最終形成每個注意力頭的 KV 快取分配結果。
KV 快取選擇:
在確定每個注意力頭保留的 KV 快取條目數量後,採用基於注意力的選擇策略,保留與最後 α 個指令 token 最相關的條目。
實驗結果
在 LongBench 和 LooGLE 等長文本理解基準測試中,HeadKV 方法在幾乎所有任務上都優於現有方法,尤其是在資源受限的情況下(KV 大小為 64 和 128)。實驗結果表明,HeadKV 方法能夠在保持與完整 KV 快取相當的效能的同時,顯著減少記憶體使用量和解碼延遲。
研究結論
HeadKV 方法透過在注意力頭級別進行動態 KV 快取分配,並結合考慮檢索和推理能力的重要性評分估計方法,有效提升了 LLM 在長文本處理任務中的效率。
未來研究方向
進一步探索不同類型注意力頭的作用,例如參與上下文學習和真實性判斷的注意力頭,以進一步提升模型效能和準確性。
開發通用的任務特定評分估計演算法,例如利用特定任務的梯度信息更有效地分配 KV 快取預算,以針對不同應用場景進行定制化壓縮。
الإحصائيات
使用 64 個 KV 快取大小僅保留了 0.7% 的總 tokens。
在 Llama-3-8B-Instruct 模型和 1024 KV 快取大小的實驗中,HeadKV-R2 方法的平均準確率達到 32.95,超過了完整 KV 快取的平均準確率(32.90)。