圖源：https://x.com/rohanpaul_ai/status/1966820906916815156

鍵值緩存（KV cache）已經(jīng)成為大模型快速運(yùn)行的核心技術(shù)，它就像一個(gè)「記憶庫」，能夠保存之前計(jì)算過的結(jié)果并重復(fù)使用，這樣就不用每次都重新計(jì)算同樣的內(nèi)容。

但是，這個(gè)記憶庫有個(gè)問題：輸入的文本越長，需要的存儲(chǔ)空間就越大，而且模型處理長文本時(shí)會(huì)變得非常慢。

為了應(yīng)對(duì)這些挑戰(zhàn)，現(xiàn)有的 KV cache 壓縮方法主要依賴基于規(guī)則的啟發(fā)式方法。當(dāng)前的方法可以歸類為三種范式：

跨所有層的固定位置信息保留

基于注意力權(quán)重的均勻?qū)蛹?jí)分配淘汰機(jī)制

具有預(yù)定義深度衰減的金字塔策略

雖然這些方法在降低內(nèi)存占用方面有效，但它們未能考慮兩個(gè)關(guān)鍵問題：

• transformer 層在信息處理中的不同功能角色
• 緩存與任務(wù)性能之間的動(dòng)態(tài)關(guān)系

僅依賴基于規(guī)則的 KV cache 預(yù)算分層分配，可能導(dǎo)致任務(wù)相關(guān)信息無法被最優(yōu)地保留。

針對(duì)這些限制，來自中國科學(xué)院大學(xué)、中國科學(xué)院自動(dòng)化研究所的 Bohan Yu 和蘇黎世聯(lián)邦理工學(xué)院的 Yekun Chai 受到（Chai 等，2022）的啟發(fā)，采用進(jìn)化算法直接基于任務(wù)性能搜索最優(yōu)的 KV cache 分配。

• 論文標(biāo)題：EvolKV: Evolutionary KV Cache Compression for LLM Inference
• 論文鏈接：https://arxiv.org/pdf/2509.08315

他們引入了 EvolKV，這是一個(gè)進(jìn)化框架，能夠自適應(yīng)地在 transformer 層之間分配 KV cache 預(yù)算，如圖 1 所示。它將每層 KV cache 預(yù)算制定為優(yōu)化變量，將其分為組，并采用進(jìn)化算法迭代搜索能夠直接最大化下游任務(wù)適應(yīng)度得分的組別配置。通過將任務(wù)驅(qū)動(dòng)優(yōu)化與層特定緩存剪枝相結(jié)合，EvolKV 實(shí)現(xiàn)了與不同層的不同貢獻(xiàn)相匹配的細(xì)粒度、性能感知分配。

與剛性啟發(fā)式方法相比，EvolKV 為以下游任務(wù)目標(biāo)為導(dǎo)向的逐層 KV cache 預(yù)算分配提供了一個(gè)靈活而有效的機(jī)制。首先，它將層 / 組級(jí)緩存預(yù)算制定為可學(xué)習(xí)參數(shù)，其中，作者將層分組為優(yōu)化單元以實(shí)現(xiàn)高效搜索。然后，它們使用黑盒進(jìn)化優(yōu)化方法直接最大化下游任務(wù)的性能。

通過這種方式，他們的方法能夠?qū)崿F(xiàn)任務(wù)感知的細(xì)粒度緩存分配，自動(dòng)適應(yīng)每個(gè)組或?qū)拥墓δ茇暙I(xiàn)。具體而言，它能夠適應(yīng)多樣化的評(píng)估標(biāo)準(zhǔn)，如準(zhǔn)確率和 F1 分?jǐn)?shù)，并在沒有預(yù)定義假設(shè)的情況下發(fā)現(xiàn)非均勻分布（即偏離啟發(fā)式固定長度或金字塔模式的模式）。

作者在 Mistral 7B-Instruct 和 Llama-3-8B-Instruct 上進(jìn)行了全面實(shí)驗(yàn)，在四個(gè)不同的基準(zhǔn)測試（十一項(xiàng)任務(wù)）上評(píng)估 EvolKV，涵蓋長上下文檢索、長上下文推理和數(shù)學(xué)任務(wù)。結(jié)果表明，任務(wù)優(yōu)化的 KV cache 分配產(chǎn)生了一致的改進(jìn)：

• 在 Needle-in-a-Haystack 基準(zhǔn)測試中，EvolKV 比最佳基線提高了多達(dá) 13%
• 在 RULER 基準(zhǔn)測試中，EvolKV 比最強(qiáng)基線提升了多達(dá) 3.6%
• 在 LongBench 評(píng)估中，它在廣泛的目標(biāo) KV cache 預(yù)算范圍內(nèi)（從 128 到 2048）始終優(yōu)于所有基線方法，并且在僅使用完整模型 1.5% 的 KV cache 預(yù)算的情況下，其性能顯著超過了完整模型。
• 對(duì)于 GSM8K，在 128 KV cache 預(yù)算下，EvolKV 比最強(qiáng)基線在準(zhǔn)確率上提高了多達(dá) 7 個(gè)百分點(diǎn)，保持了完整模型性能的 95.7%，而最強(qiáng)基線在 512 KV cache 預(yù)算下僅保持了 84.5%。

EvolKV 詳解

EvolKV 是一個(gè)動(dòng)態(tài)的、任務(wù)驅(qū)動(dòng)的進(jìn)化框架，通過利用下游任務(wù)的性能反饋來自適應(yīng)地為每一層分配 KV cache 預(yù)算。圖 2a 展示了 EvolKV 與其他方法之間預(yù)算分配的比較。

進(jìn)化壓縮的優(yōu)化目標(biāo)

進(jìn)化算法的工作原理是產(chǎn)生候選解決方案并評(píng)估它們的適應(yīng)度，然后根據(jù)適應(yīng)度反饋不斷改進(jìn)搜索策略，逐步引導(dǎo)整個(gè)群體朝著更好的解決方案發(fā)展。在本文中，EvolKV 將下游任務(wù)的性能反饋當(dāng)作適應(yīng)度分?jǐn)?shù)，并利用進(jìn)化算法來指導(dǎo)每一層的 KV cache 壓縮。

具體來說，在一個(gè)有 L 個(gè) transformer 層的語言模型中，作者用 k_i ∈ N 來表示第 i 層的 KV cache 預(yù)算，其中?i ∈ {1, . . . , L}。給定進(jìn)化算法為下游任務(wù) f (?) 產(chǎn)生的一組候選壓縮方案，他們的目標(biāo)是找到最優(yōu)方案 S*，這個(gè)方案既要最大化任務(wù)性能，又要盡量接近目標(biāo)平均 KV cache 預(yù)算 c：

其中 f (S) 是使用壓縮方案 S ∈ 時(shí)獲得的下游任務(wù)性能，超參數(shù) λ > 0 用來平衡原始性能和緩存效率。由于下游性能指標(biāo)種類繁多且數(shù)值范圍不同（比如準(zhǔn)確率、F1 分?jǐn)?shù)、ROUGE 分?jǐn)?shù)），作者采用了一個(gè)直接與任務(wù)性能進(jìn)行權(quán)衡的緩存效率項(xiàng)，以確保可比性。

緩存效率項(xiàng) CACHESCORE (S, c) ∈ [0, 1] 的工作機(jī)制是：如果某個(gè)方案的平均每層緩存預(yù)算超過了目標(biāo)預(yù)算 c，就給它較低的分?jǐn)?shù)；對(duì)于那些保持在目標(biāo)范圍內(nèi)的方案，則應(yīng)用平滑折扣：

其中 γ ∈ (0, 1] 是平滑因子。因此，這個(gè)目標(biāo)函數(shù)偏好那些提供強(qiáng)大任務(wù)性能且將平均 KV cache 預(yù)算保持在接近或低于期望預(yù)算的壓縮方案。

KV Cache 預(yù)算的分組

為了提高優(yōu)化效率，作者引入了組大小參數(shù) n_g，將 KV cache 預(yù)算 K = {k_1, k_2, . . . , k_L} 劃分為 J = ?L/n_g?個(gè)組，記為 G = {g_1, g_2, . . . , g_J}。每個(gè)組 g_j 包含連續(xù)的緩存預(yù)算子集，定義為 g_j =  ?_j ∈ {1, 2, . . . , J}。

為簡化起見，作者假設(shè)層總數(shù) L 能被組大小 n_g 整除，即 L = J?n_g。在這種表述下，候選壓縮方案在組級(jí)別應(yīng)用，記為?；谙掠稳蝿?wù)性能為每個(gè)組選擇的最優(yōu)方案記為。這種分組表述顯著減少了搜索空間，并在進(jìn)化搜索過程中促進(jìn)了更穩(wěn)定的優(yōu)化動(dòng)態(tài)。

進(jìn)化壓縮的迭代過程

本文中的 KV cache 預(yù)算優(yōu)化以分組方式進(jìn)行，如算法 1 所示，從底層到頂層依次進(jìn)行。在優(yōu)化每個(gè)組時(shí)，先前優(yōu)化組的 KV cache 預(yù)算被固定為各自的最優(yōu)方案，而其余組保持其初始值。如果候選方案 S_g 獲得比當(dāng)前最佳方案更高的適應(yīng)度分?jǐn)?shù) r，則相應(yīng)地更新當(dāng)前組的 KV cache 預(yù)算。這個(gè)過程反復(fù)進(jìn)行，直到所有組都被優(yōu)化。

KV 緩存預(yù)算補(bǔ)全

為了確保評(píng)估的公平性，作者對(duì)總大小偏離目標(biāo)的 KV cache 預(yù)算優(yōu)化結(jié)果進(jìn)行補(bǔ)全。具體而言，他們首先計(jì)算實(shí)際總 KV cache 預(yù)算與目標(biāo)總預(yù)算 T = c?L 之間的差異，記為?_cache = T ? A。然后根據(jù)各層在 A 中的原始占比，將這個(gè)差異按比例重新分配到各層。補(bǔ)全后的 KV cache 預(yù)算為 B = {b_1, b_2, . . . , b_L}，其中，i ∈ {1, 2, . . . , L}。

實(shí)驗(yàn)結(jié)果

在 LongBench 上的結(jié)果

表 1 報(bào)告了在 Mistral-7B-Instruct 上使用 16 個(gè) LongBench 子數(shù)據(jù)集的評(píng)估結(jié)果，所有訓(xùn)練樣本已被移除。在所有評(píng)估的 KV cache 預(yù)算中，EvolKV 始終獲得最高的平均性能，優(yōu)于所有基于規(guī)則的基線方法。此外，在包括 MultiFieldQA-en、2WikiMultihopQA、MuSiQue、TriviaQA 和 PassageRetrieval-en 在內(nèi)的幾個(gè)子數(shù)據(jù)集上，EvolKV 不僅與未壓縮的完整模型保持競爭力，甚至在某些 KV cache 預(yù)算下超越了完整模型。

表 2 展示了 Llama-3-8B-Instruct 上的類似結(jié)果，同樣排除了訓(xùn)練樣本。EvolKV 在所有 KV cache 預(yù)算下都表現(xiàn)出了優(yōu)異的性能。值得注意的是，在緩存預(yù)算為 128 時(shí)，EvolKV 在 TREC 子集上比最強(qiáng)基線高出 7.69 個(gè)百分點(diǎn)，突出了其對(duì)多樣化下游任務(wù)的強(qiáng)適應(yīng)性。

在 GSM8K 上的結(jié)果

圖 3b 展示了 EvolKV 為 Llama-3-8B-Instruct 優(yōu)化的 KV cache 預(yù)算分配。表 3 報(bào)告了 Llama-3-8B-Instruct 和 Mistral-7B-Instruct 兩個(gè)模型對(duì)應(yīng)的測試集準(zhǔn)確率。在所有配置中，EvolKV 在兩個(gè)模型上都始終優(yōu)于基線方法。具體而言，在 Llama-3-8B-Instruct 上，它相比最強(qiáng)競爭對(duì)手取得了顯著改進(jìn)，在 KV 緩存預(yù)算為 128、256 和 512 時(shí)，準(zhǔn)確率分別至少提升了 7.28、2.05 和 7.58 個(gè)百分點(diǎn)。值得注意的是，EvolKV 使用減少的緩存預(yù)算（c = 512）就達(dá)到了完整模型性能的 95.7%，顯著優(yōu)于所有基線方法，其中最佳基線結(jié)果僅達(dá)到 84.5%。

在 NIAH 和 RULER 上的結(jié)果

作者在 NIAH 上評(píng)估了 EvolKV 以及所有基線模型的長上下文檢索能力。圖 7 展示了這部分評(píng)估的結(jié)果：與基線方法相比，EvolKV 在 Llama3-8B-Instruct 上取得了超過 4 個(gè)百分點(diǎn)的改進(jìn)，在 Mistral-7B-Instruct 上取得了超過 13 個(gè)百分點(diǎn)的顯著提升。這些結(jié)果表明，EvolKV 有效地探索并利用了模型在長上下文檢索中的潛在層級(jí) KV cache 分配。

作者在 RULER 基準(zhǔn)測試上進(jìn)一步評(píng)估了 NIAH 中優(yōu)化的 KV 緩存分配。如表 4 所示，EvolKV 在平均得分上始終優(yōu)于所有基線方法，在 Mistral-7B-Instruct 上提升了多達(dá) 0.99 分，在 Llama-3-8B-Instruct 上提升了 3.6 分。這些結(jié)果進(jìn)一步證明了 EvolKV 強(qiáng)大的泛化能力、長上下文檢索和推理能力，因?yàn)閮?yōu)化的 KV 預(yù)算可以有效地遷移到其他基準(zhǔn)評(píng)估中，這表明 EvolKV 揭示了潛在的層級(jí)分配策略。

更多細(xì)節(jié)請參見原論文。