引言

本文將探討KV-Cache如何通過在內存使用和計算時間之間進行巧妙的權衡，使像ChatGPT和DeepSeek這樣的語言模型在生成文本時更快。

總結11篇最近的研究論文，歸納三大類：token選擇、后處理壓縮技術和架構重新設計。包括DeepSeek的多頭潛在注意力（MLA），這些論文在這一基本思想的基礎上，進一步提高了大型語言模型（LLM）推理的時間效率。

一、思考

為什么文本生成如此緩慢

讓我們從一個簡單的類比開始。想象你在寫一個故事，每寫一個新詞，你都需要重新閱讀到目前為止的整個故事以保持一致性。故事越長，重新閱讀的時間就越長。這正是大型語言模型在文本生成時所面臨的問題。

自注意力的基本構建塊

現代語言模型的核心是一種稱為自注意力的機制。對于一個由n個標記（大致對應單詞）組成的序列，每個標記都需要“查看”或“關注”所有其他標記以理解上下文。

這種查看一切的過程的計算成本隨著序列長度的增長而增長：

• 對于n個標記，每個標記都需要查看所有nn個標記
• 這意味著成本與n x n = n²成正比
• 用數學符號表示，我們將其寫為O(n²)的復雜度

真正的問題：一次生成一個標記

當語言模型生成文本時，它一次生成一個標記，這就是事情變得計算密集的地方：

• 第一個標記：查看1個標記（成本：O(1²)）
• 第二個標記：查看2個標記（成本：O(2²)）
• 第三個標記：查看3個標記（成本：O(3²)）
• 以此類推，直到第n個標記：查看n個標記（成本：O(n²)）

如果我們將生成長度為的序列的所有這些成本加起來，我們得到：

這種O(n³)的成本意味著隨著文本的增長，生成時間會極其迅速地增長。例如，生成兩倍長的序列大約需要八倍的時間！顯然，我們需要一個更好的方法。

解決方案：鍵值（KV）緩存

KV 緩存背后的關鍵是，我們正在做大量冗余工作。在生成每個新標記時，我們會重新計算之前已經處理過的所有先前標記。讓我們看看如何解決這個問題。

什么是鍵值緩存？

可以將 KV 緩存想象成一個智能記事本，我們會在第一次看到每個 token 時記下有關它的重要信息。對于每個 token，我們計算并存儲兩件事：

• 鍵（k）：可以將其視為一種尋址機制——它有助于確定此標記與未來標記的相關性
• 值（v）：可以將其視為當此標記被發現相關時實際使用的信息

從數學上，我們計算這些為：

• 鍵：k = xW_k（其中是x標記，W_k是一個學習到的變換）
• 值：v = xW_v（其中W_v是另一個學習到的變換）

在生成一個新標記時，我們使用它的查詢（計算方式類似于鍵）通過將其與所有存儲的鍵進行比較來在我們的緩存中找到相關信息。然后使用匹配的值來幫助生成標記。

KV緩存如何加速

有了KV緩存，處理過程變得更加高效：

1. 當我們遇到一個新token時，只需要計算它的key和value一次
2. 對于所有后續的token，我們可以直接從緩存中查找這些預計算的值
3. 這意味著每個新token只需要做少量新的計算，而不是重新做所有之前的計算

顯然有一個權衡：

• 我們需要更多的內存來存儲所有的keys和values。對于一個具有：

     L層

     H注意力頭

     序列長度n

     key/value維度d_k，總的內存開銷為L x H x n x d_k x 2值（這個2是因為需要存儲keys和values）。這會隨著序列長度n線性增長(O(n))，但對于大模型來說，常數因子可能非常大。

• 但作為回報，我們將計算成本從O(n³)降低到O(n²)。

要理解為什么是O(n²)，讓我們看一下每一步的成本：

1. 第一步：處理一個token ->成本O(1)
2. 第二步：處理一個新token + 查找1個緩存的token -> 成本O(2)
3. 第三步：處理一個新token + 查找2個緩存的token -> 成本O(3)
4. 依此類推...

將這些加起來：

   O(1 + 2 + 3 + ... + n) = O(n²)

這相比O(n³)是一個顯著的改進！雖然我們仍然需要做查看所有前面的tokens的基礎工作<O(n²)>，但我們避免了每一步都進行昂貴的重新計算。

內存挑戰：為什么我們需要更好的解決方案

雖然KV緩存是一個強大的優化手段，但它伴隨著顯著的內存開銷。讓我們通過一個具體的例子來看看，使用像Llama3 70B這樣的現代大語言模型：

• L = 80層
• H = 64注意力頭
• B = 8批量大小為8個序列
• d_k= 128key/value維度
• 16位精度

處理一個批量（8個序列，每個序列1000個token）所需的內存為：

L x H x B x n x d_k x 2 x 2字節=80 x 64 x 8 x 1000 x 128 x 2 x 2字節=20.97GB

這種巨大的內存使用帶來了幾個挑戰：

1. 隨著序列長度線性增長
2. 與批量大小成倍增長，支持并行處理
3. 限制了我們可以處理的最大上下文長度
4. 限制了在內存受限設備上的部署

這些挑戰激發了研究界的一波創新，導致了各種優化KV緩存使用的技術。接下來，將探討這些前沿的解決方案。

二、如何改善傳統的KV緩存？

以下論文代表了KV緩存優化的關鍵創新。我們將通過三大主要方法來探索它們：token選擇、后處理壓縮技術和架構重設計。

2.1 Token 選擇和修剪方法（Token Selection and Pruning Approaches）

1) Heavy-Hitter Oracle (H2O)

• ??https://arxiv.org/abs/2306.14048??

H2O 引入了在KV緩存中識別和保留重要token的概念：

• 重型Token（Heavy-Hitter Tokens）：H2O 識別在生成過程中具有最高累計注意力分數的token，這些token遵循冪律分布。這些token對于模型的功能至關重要，因此在緩存中優先處理。
• 動態次模撤銷（Dynamic Submodular Eviction）：該方法將緩存管理問題框架化為一個優化問題，目標函數為次模函數F(S)，用于量化token集合的S重要性：

確保每次最多只移除一個token。這個貪心算法在計算上高效，并在次模約束下保證接近最優的性能。

• 結果：通過該方法，KV緩存大小減少了5倍，幾乎沒有精度損失，并且吞吐量提升了高達29倍。

2) StreamLLM

• ??https://arxiv.org/abs/2309.17453??

• 作者觀察到注意力匯聚（Attention Sinks）現象：解碼過程中，初始token充當自然的注意力錨點。
• 如果沒有這些注意力匯聚的token，傳統窗口注意力方法的性能會下降。
• 基于這一觀察，他們引入了滾動緩存（Rolling Cache），它保留了初始token，并處理最近的上下文，從而實現了無限長度序列的處理。
• 他們還展示了這些匯聚token可以通過訓練獲得，作為專用的注意力錨點，從而減少對多個初始token的依賴。

3) Value-Aware Token Pruning (VATP)

• ??https://arxiv.org/abs/2406.12335??

VATP 擴展了 H2O 的 token 重要性概念，考慮了注意力模式和價值向量的屬性：

性能與效率：

• 在16個 LongBench 任務中，VATP 在12-14個任務中超越了 H2O 和 Scissorhands 等基準。
• 在保持最小性能損失的情況下，實現了50%的有效壓縮。
• 引入的計算開銷幾乎可以忽略不計，并且與 Scissorhands 集成時兼容 FlashAttention。

2.2 后處理壓縮技術（Post-hoc Compression Techniques）

這些方法壓縮或優化KV緩存，同時保持標準的Transformer架構。

4) Adaptive KV Compression (FastGen)

• ??https://arxiv.org/pdf/2310.01801??

FastGen 通過觀察運行時的注意力模式引入了自適應壓縮：

自適應壓縮策略：

5) 動態內存壓縮（DMC）

• ??https://arxiv.org/pdf/2403.09636??

DMC 引入了自適應的 token 合并：

6)L₂范數基礎的壓縮

• ??https://arxiv.org/pdf/2406.11430??

本文提出了一個令人驚訝的觀察：緩存 KV 對的L₂范數與注意力分數之間存在明確的相關性，低L₂范數的鍵嵌入通常會導致解碼時的高注意力分數。因此，提出了一個簡單但有效的壓縮目標：

2.3 體系結構重設計

這些方法改變了 Transformer 架構，以更高效地處理 KV 緩存，通常將壓縮直接集成到架構中。

7) 多查詢注意力（MQA）

• ??https://arxiv.org/pdf/2305.13245??

• 核心思想：MQA 通過共享單個鍵值頭跨所有查詢頭來減少 KV 緩存大小，替代傳統的多頭注意力（MHA）：

K = XW_K ， V = XW_V

其中K和V是共享的鍵和值投影。

• 優點：將 KV 緩存大小減少了H（注意力頭的數量），顯著降低了內存帶寬開銷。
• 權衡：雖然 MQA 更快，但在需要多樣化注意力模式的任務中，通常會遭遇質量下降。

8) 分組查詢注意力（GQA）

• ??https://arxiv.org/abs/2305.13245??

• 核心思想：GQA 在完全多頭注意力和 MQA 之間進行插值，提供了推理速度和模型質量之間的可擴展權衡。它將查詢頭分為G組，每組共享一個單獨的鍵值頭：

9) 多頭潛在注意力（MLA）

• ??https://arxiv.org/abs/2405.04434??

DeepSeek的多頭潛在注意力（MLA）采用了一種新穎的方法來減少KV緩存開銷。雖然MQA和GQA通過頭共享來實現這一目標，MLA則采用低秩潛在壓縮技術，在保持多頭注意力的優點的同時，減少了KV緩存的大小。

• MLA通過將鍵（keys）和值（values）壓縮成低維度的潛在向量，來減少KV緩存的大小。
• 它將鍵值嵌入（key-value embeddings）降投到一個壓縮的潛在空間：

10) SnapKV

• ??https://arxiv.org/pdf/2404.14469??

11) 只緩存一次（YOCO）

• ??https://arxiv.org/pdf/2405.05254??

YOCO修改了Transformer架構以優化緩存：

• 全局緩存：使用解碼器-解碼器設計，只有一個共享的KV緩存。
• 復雜度減少：將內存從O(N x L)減少到O(N + L)，其中N是序列長度，L是層數。
• 高效注意力：自解碼器采用滑動窗口注意力或門控保留機制，使內存使用保持恒定（O(C)，其中C是小窗口大?。?。

結論

KV-Cache技術是將Transformer模型擴展和優化到實際應用中的核心。像動態逐出、壓縮和結構化近似等創新，持續推動著在長上下文或資源受限的場景中實現更高效的技術。KV-Cache仍然是一個活躍的研究領域，既提供了理論上的見解，也帶來了實際的改進。

公眾號大模型自然語言處理  作者：余俊暉

原文鏈接：??https://mp.weixin.qq.com/s/7j9sVIlJQDPnji9bas09ig???