論文原文：https://openreview.net/pdf?id=tkiZQlL04w
最近學(xué)習(xí)華為的RazorAttention，水平有限，根據(jù)論文做了初步了解。

1 背景：KV緩存成為部署模型的主要瓶頸

長上下文大型語言模型（LLM）在不同任務(wù)的自然語言處理方面具有顯著的先進(jìn)性。在LLM 模型的應(yīng)用場景中，KV（Key-Value）緩存需要保存所有詞元的Key與Value，以便節(jié)省后續(xù)解碼時(shí)的計(jì)算復(fù)雜度。然而，隨著輸入序列的增長，KV Cache 所占用的空間也會不斷增加，總體達(dá)到 O(seq_len)的空間復(fù)雜度，導(dǎo)致顯存消耗迅速增加（極端情況下 KV Cache 與權(quán)重的顯存占比可達(dá)到 9:1），進(jìn)而影響模型推理性能、限制其部署。

1.1 當(dāng)前主流方式：

• 量化（Quantization）：量化方法通過降低存儲每個(gè)數(shù)據(jù)點(diǎn)的精度來減少內(nèi)存使用。使用較低的精度（如4位或8位）來表示數(shù)據(jù)。
• 窗口（Window）：這種方法通過限制緩存中存儲的數(shù)據(jù)量來優(yōu)化顯存使用。例如，只緩存最近的或最重要的數(shù)據(jù)。
• 稀疏化（Sparse）：這種方法通過僅存儲非零或重要的數(shù)據(jù)點(diǎn)來減少內(nèi)存占用。它利用數(shù)據(jù)的稀疏性來優(yōu)化存儲。

目前業(yè)界流行的 KV Cache 壓縮算法均為實(shí)時(shí)動態(tài)壓縮，即實(shí)時(shí)地計(jì)算一些指標(biāo)：注意力分?jǐn)?shù)（與 FlashAttention 不兼容）或者 Topk 判斷稀疏模式，對推理速度均有非常高的負(fù)面影響，在實(shí)際部署中不適用。

2 Razor Attention

基于Attention可解釋性，提出了RazorAttention靜態(tài)KV Cache壓縮算法，是目前業(yè)界唯一的靜態(tài)方法。RazorAttention的動機(jī)來源于 Anthropic 2022 年的論文 In-context Learning and Induction Heads，該文章介紹了 Decoder-Only架構(gòu) 的語言大模型獨(dú)有的上下文學(xué)習(xí)能力（In-Context Learning），并發(fā)現(xiàn)有一部分注意力頭和模型的上下文學(xué)習(xí)能力強(qiáng)相關(guān)，文中稱之為具有推斷能力的 Induction Heads。在此基礎(chǔ)上，本工作聯(lián)想到大模型的長序列能力也是上下文學(xué)習(xí)能力的一個(gè)子集，并圍繞此展開了研究。實(shí)驗(yàn)表明，RazorAttention 壓縮了 70%的KV Cache，并保持模型的長序列能力幾近語義無損。

image.png

2.1 算法原理

業(yè)界最新的長序列KV Cache壓縮算法大多數(shù)為在線動態(tài)壓縮方法（實(shí)時(shí)計(jì)算?些指標(biāo)），存在如下問題：

• 例如在線計(jì)算Attention得分，存在Overhead，且與FlashAttention不兼容
• 最新技術(shù)在往輕量化重要性計(jì)算發(fā)展，但是依然存在TopK計(jì)算；隨著序列長度增加，Top-K計(jì)算的overhead將顯著增加
• MInfer等算法基于稀疏模式，要求在線執(zhí)行稀疏計(jì)算
• 依賴當(dāng)前輸入的丟棄token無法還原，影響多輪對話等場景應(yīng)用

文中先以基于 ALiBi 編碼的模型為例，從理論上證明每個(gè)注意力頭擁有不同的視野范圍，因此可以通過動態(tài)調(diào)整每個(gè)頭的 KV Cache大小來實(shí)現(xiàn)壓縮。具體來說，需要計(jì)算出每個(gè)頭的有效視野Lh，并只保留最近的 Lh 個(gè) token 在 KV Cache中，而將距離較遠(yuǎn)的 token 丟棄掉。這樣可以保證在不損失準(zhǔn)確率的情況下，大幅減少存儲空間的需求。
在此基礎(chǔ)上，“不同注意力頭擁有不同視野”的結(jié)論可以擴(kuò)展到任意位置編碼（如RoPE）的語言大模型。本?作在Induction Heads基礎(chǔ)上發(fā)現(xiàn)了Echo Heads同樣對長文本任務(wù)起到關(guān)鍵作?，并統(tǒng)稱為檢索頭（Retrieval Heads）。檢索頭通常能有效利用長距離信息，對輸入擁有全局視野，其 KV Cache也因此非常重要：

• Echo Head：關(guān)注前文中與當(dāng)前token相同的token
• Induction Head：關(guān)注前文中與當(dāng)前token相同的下一個(gè)token

在壓縮算法設(shè)計(jì)上，RazorAttention對檢索頭的KV Cache進(jìn)行保護(hù)，確保上下文當(dāng)中重要且主要信息不丟失，并直接減少了非檢索頭的KV Cache大小，以達(dá)到壓縮的目的。在Non-retrieval Head上進(jìn)行壓縮，需要保留Attention Sink+Local Attention。此外，本文提出了Compensation token（補(bǔ)償token），對于丟棄掉的KV Cache以平均值形式保留在非檢索頭中，達(dá)到保護(hù)非檢索頭的局部視野的效果?；赗oPE位置編碼的KV Cache壓縮原理如下：

image.png

image.png

2.2 Alibi模型的attention

在傳統(tǒng)的Transformer模型中，位置嵌入通常是預(yù)先計(jì)算并添加到輸入序列的每個(gè)token上的。然而，Alibi不采用這種方法。它是在計(jì)算注意力分?jǐn)?shù)的過程中動態(tài)的考慮信息位置的。
通俗來說，Alibi類型的模型是使用相對位置來計(jì)算注意力分?jǐn)?shù)的。論文中發(fā)現(xiàn)，在這種情況的“不同的注意力擁有不同的視野”這個(gè)重要結(jié)論：基于Alibi模型，從理論上證明每個(gè)注意力頭擁有不同的視野范圍，因此可以通過動態(tài)調(diào)整每個(gè)頭的KVcache大小來實(shí)現(xiàn)壓縮。具體來說，需要計(jì)算出每個(gè)頭的有效視野 L.h，并只保留最近的Lh個(gè)token在KVcache中，而將距離較遠(yuǎn)的token丟掉。有點(diǎn)窗口化的意思，但是通過計(jì)算，在保證不損失準(zhǔn)確率的情況下，大幅減少存儲空間的需求。
論文中計(jì)算注意力頭視野的公式為：

image.png

   可以看到，它設(shè)置了一個(gè)超參數(shù)，當(dāng)前的詞對較遠(yuǎn)的注意力分?jǐn)?shù)低于這個(gè)設(shè)定的超參數(shù)的時(shí)候，就會將它舍棄，被認(rèn)為是不需要的信息。只需要讓每個(gè)注意力頭加起來的“視野”能夠覆蓋所有信息。保證語義信息不丟失。

2.3 RoPE模型attention

在論文中表示，“不同注意力頭擁有不同視野”的結(jié)論可以擴(kuò)展到任意位置編碼（如ROPE）的語言大模型。本工作在Induction Heads基礎(chǔ)上發(fā)現(xiàn)了EchoHeads同樣對長文本任務(wù)起到關(guān)鍵作用，并統(tǒng)稱為檢索頭（Retrieval Heads）。
檢索頭通常能有效利用長距離信息，對輸入擁有全局視野，其KV Cache也因此非常重要：

• Echo Head：關(guān)注前文中與當(dāng)前token相同的token
• Induction Head： 關(guān)注前文中與當(dāng)前token相同的下一個(gè)token

在壓縮算法設(shè)計(jì)上，RazorAttention對檢索頭的KVCache進(jìn)行保護(hù)，確保上下文當(dāng)中重要且主要信息不丟失，并直接減少了非檢索頭的KV Cache大小，以達(dá)到壓縮的目的。在Non-retrievalHead上進(jìn)行壓縮的目的，需要保留Attention Sink+Local Attention。
此外，本文提出了Compensation token（補(bǔ)償token），對于丟棄掉的KV Cache以平均值形式保留在非檢索頭中，達(dá)到保護(hù)非檢索頭的局部視野的效果?；赗oPE位置編碼的KV Cache壓縮原理如上圖，將部分信息壓縮。
可以簡單理解為，檢索頭作為小隊(duì)長，負(fù)責(zé)統(tǒng)領(lǐng)大局；而非檢索頭作為專家，專注自己的這部分工作內(nèi)容。從而做到性能幾乎沒有損失，而減少了70%的存儲。

3 總結(jié)

文中提出的RazorAttention是目前業(yè)界唯一的靜態(tài)KV Cache壓縮算法，一方面沒有在線動態(tài)計(jì)算開銷（既不依賴Attention score，也不涉及在線Topk計(jì)算），類似PTQ僅需少量輸入樣本進(jìn)行離線校準(zhǔn)，便可輕量快捷的決定Attention稀疏模式；另一方面能與FlashAttention等主流融合算子兼容，無明顯overhead。

參考：
https://blog.csdn.net/weixin_57667062/article/details/143898047