DeepSeek的多頭潛在注意力（MLA）和及其11種KV-Cache技巧演進大總結 原創
引言
本文將探討KV-Cache如何通過在內存使用和計算時間之間進行巧妙的權衡，使像ChatGPT和DeepSeek這樣的語言模型在生成文本時更快。
總結11篇最近的研究論文，歸納三大類：token選擇、后處理壓縮技術和架構重新設計。包括DeepSeek的多頭潛在注意力（MLA），這些論文在這一基本思想的基礎上，進一步提高了大型語言模型（LLM）推理的時間效率。
一、思考
為什么文本生成如此緩慢
讓我們從一個簡單的類比開始。想象你在寫一個故事，每寫一個新詞，你都需要重新閱讀到目前為止的整個故事以保持一致性。故事越長，重新閱讀的時間就越長。這正是大型語言模型在文本生成時所面臨的問題。
自注意力的基本構建塊
現代語言模型的核心是一種稱為自注意力的機制。對于一個由n個標記（大致對應單詞）組成的序列，每個標記都需要“查看”或“關注”所有其他標記以理解上下文。
這種查看一切的過程的計算成本隨著序列長度的增長而增長：
- 對于n個標記，每個標記都需要查看所有nn個標記
- 這意味著成本與n x n = n2成正比
- 用數學符號表示，我們將其寫為O(n2)的復雜度
真正的問題：一次生成一個標記
當語言模型生成文本時，它一次生成一個標記，這就是事情變得計算密集的地方：
- 第一個標記：查看1個標記（成本：O(12)）
- 第二個標記：查看2個標記（成本：O(22)）
- 第三個標記：查看3個標記（成本：O(32)）
- 以此類推，直到第n個標記：查看n個標記（成本：O(n2)）
如果我們將生成長度為的序列的所有這些成本加起來，我們得到：
這種O(n3)的成本意味著隨著文本的增長，生成時間會極其迅速地增長。例如，生成兩倍長的序列大約需要八倍的時間！顯然，我們需要一個更好的方法。
解決方案：鍵值（KV）緩存
KV 緩存背后的關鍵是，我們正在做大量冗余工作。在生成每個新標記時，我們會重新計算之前已經處理過的所有先前標記。讓我們看看如何解決這個問題。
什么是鍵值緩存？
可以將 KV 緩存想象成一個智能記事本，我們會在第一次看到每個 token 時記下有關它的重要信息。對于每個 token，我們計算并存儲兩件事：
- 鍵（k）：可以將其視為一種尋址機制——它有助于確定此標記與未來標記的相關性
- 值（v）：可以將其視為當此標記被發現相關時實際使用的信息
從數學上，我們計算這些為：
- 鍵：k = xWk（其中是x標記，Wk是一個學習到的變換）
- 值：v = xWv（其中Wv是另一個學習到的變換）
在生成一個新標記時，我們使用它的查詢（計算方式類似于鍵）通過將其與所有存儲的鍵進行比較來在我們的緩存中找到相關信息。然后使用匹配的值來幫助生成標記。
KV緩存如何加速
有了KV緩存，處理過程變得更加高效：
- 當我們遇到一個新token時，只需要計算它的key和value一次
- 對于所有后續的token，我們可以直接從緩存中查找這些預計算的值
- 這意味著每個新token只需要做少量新的計算，而不是重新做所有之前的計算
顯然有一個權衡：
- 我們需要更多的內存來存儲所有的keys和values。對于一個具有：
L層
H注意力頭
序列長度n
key/value維度dk，總的內存開銷為L x H x n x dk x 2值（這個2是因為需要存儲keys和values）。這會隨著序列長度n線性增長(O(n))，但對于大模型來說，常數因子可能非常大。
- 但作為回報，我們將計算成本從O(n3)降低到O(n2)。
要理解為什么是O(n2)，讓我們看一下每一步的成本：
- 第一步：處理一個token ->成本O(1)
- 第二步：處理一個新token + 查找1個緩存的token -> 成本O(2)
- 第三步：處理一個新token + 查找2個緩存的token -> 成本O(3)
- 依此類推...
將這些加起來：
O(1 + 2 + 3 + ... + n) = O(n2)
這相比O(n3)是一個顯著的改進！雖然我們仍然需要做查看所有前面的tokens的基礎工作<O(n2)>，但我們避免了每一步都進行昂貴的重新計算。
內存挑戰：為什么我們需要更好的解決方案
雖然KV緩存是一個強大的優化手段，但它伴隨著顯著的內存開銷。讓我們通過一個具體的例子來看看，使用像Llama3 70B這樣的現代大語言模型：
- L = 80層
- H = 64注意力頭
- B = 8批量大小為8個序列
- dk= 128key/value維度
- 16位精度
處理一個批量（8個序列，每個序列1000個token）所需的內存為：
L x H x B x n x dk x 2 x 2字節=80 x 64 x 8 x 1000 x 128 x 2 x 2字節=20.97GB
這種巨大的內存使用帶來了幾個挑戰：
- 隨著序列長度線性增長
- 與批量大小成倍增長，支持并行處理
- 限制了我們可以處理的最大上下文長度
- 限制了在內存受限設備上的部署
這些挑戰激發了研究界的一波創新，導致了各種優化KV緩存使用的技術。接下來，將探討這些前沿的解決方案。
二、如何改善傳統的KV緩存？
以下論文代表了KV緩存優化的關鍵創新。我們將通過三大主要方法來探索它們：token選擇、后處理壓縮技術和架構重設計。
2.1 Token 選擇和修剪方法（Token Selection and Pruning Approaches）
1) Heavy-Hitter Oracle (H2O)
H2O 引入了在KV緩存中識別和保留重要token的概念：
- 重型Token（Heavy-Hitter Tokens）：H2O 識別在生成過程中具有最高累計注意力分數的token，這些token遵循冪律分布。這些token對于模型的功能至關重要，因此在緩存中優先處理。
- 動態次模撤銷（Dynamic Submodular Eviction）：該方法將緩存管理問題框架化為一個優化問題，目標函數為次模函數F(S)，用于量化token集合的S重要性：
確保每次最多只移除一個token。這個貪心算法在計算上高效，并在次模約束下保證接近最優的性能。
- 結果：通過該方法，KV緩存大小減少了5倍，幾乎沒有精度損失，并且吞吐量提升了高達29倍。
2) StreamLLM
- 作者觀察到注意力匯聚（Attention Sinks）現象：解碼過程中，初始token充當自然的注意力錨點。
- 如果沒有這些注意力匯聚的token，傳統窗口注意力方法的性能會下降。
- 基于這一觀察，他們引入了滾動緩存（Rolling Cache），它保留了初始token，并處理最近的上下文，從而實現了無限長度序列的處理。
- 他們還展示了這些匯聚token可以通過訓練獲得，作為專用的注意力錨點，從而減少對多個初始token的依賴。
3) Value-Aware Token Pruning (VATP)
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)
FastGen 通過觀察運行時的注意力模式引入了自適應壓縮：
自適應壓縮策略：
5) 動態內存壓縮（DMC）
DMC 引入了自適應的 token 合并：
6)L2范數基礎的壓縮
本文提出了一個令人驚訝的觀察：緩存 KV 對的L2范數與注意力分數之間存在明確的相關性，低L2范數的鍵嵌入通常會導致解碼時的高注意力分數。因此，提出了一個簡單但有效的壓縮目標：
2.3 體系結構重設計
這些方法改變了 Transformer 架構，以更高效地處理 KV 緩存，通常將壓縮直接集成到架構中。
7) 多查詢注意力（MQA）
- 核心思想：MQA 通過共享單個鍵值頭跨所有查詢頭來減少 KV 緩存大小，替代傳統的多頭注意力（MHA）：
K = XWK ， V = XWV
其中K和V是共享的鍵和值投影。
- 優點：將 KV 緩存大小減少了H（注意力頭的數量），顯著降低了內存帶寬開銷。
- 權衡：雖然 MQA 更快，但在需要多樣化注意力模式的任務中，通常會遭遇質量下降。
8) 分組查詢注意力（GQA）
- 核心思想：GQA 在完全多頭注意力和 MQA 之間進行插值，提供了推理速度和模型質量之間的可擴展權衡。它將查詢頭分為G組，每組共享一個單獨的鍵值頭：
9) 多頭潛在注意力（MLA）
DeepSeek的多頭潛在注意力（MLA）采用了一種新穎的方法來減少KV緩存開銷。雖然MQA和GQA通過頭共享來實現這一目標，MLA則采用低秩潛在壓縮技術，在保持多頭注意力的優點的同時，減少了KV緩存的大小。
- MLA通過將鍵（keys）和值（values）壓縮成低維度的潛在向量，來減少KV緩存的大小。
- 它將鍵值嵌入（key-value embeddings）降投到一個壓縮的潛在空間：
10) SnapKV
11) 只緩存一次（YOCO）
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???
