隨著大語言模型在長文本場景下的需求不斷涌現(xiàn)，其核心的注意力機(jī)制（Attention Mechanism）也獲得了非常多的關(guān)注。

注意力機(jī)制會計(jì)算一定跨度內(nèi)輸入文本（令牌，Token）之間的交互，從而實(shí)現(xiàn)對上下文的理解。隨著應(yīng)用的發(fā)展，高效處理更長輸入的需求也隨之增長 [1][2]，這帶來了計(jì)算代價的挑戰(zhàn)：注意力高昂的計(jì)算成本和不斷增長的鍵值緩存（KV-Cache）代價。稀疏注意力機(jī)制可以有效緩解內(nèi)存和吞吐量的挑戰(zhàn)。

然而，現(xiàn)有稀疏注意力通常采用統(tǒng)一的稀疏注意力模式，即對不同的注意力頭和輸入長度應(yīng)用相同的稀疏模式。這種統(tǒng)一的方法難以捕捉到大語言模型中多樣的注意力模式，導(dǎo)致不同注意力頭的不同的精度 - 代價權(quán)衡被忽略。

最近，來自清華大學(xué)、無問芯穹和上海交通大學(xué)的研究團(tuán)隊(duì)發(fā)表了《MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression》，提出通過混合不同稀疏度的注意力頭，使用 25% 的注意力稠密度，就可以記憶幾乎 100% 的上下文。

本工作現(xiàn)已開源，歡迎交流討論。

• 代碼：https://github.com/thu-nics/MoA
• 主頁：https://nics-effalg.com/MoA
• arXiv：https://arxiv.org/abs/2406.14909

總覽

在大語言模型中，不同的注意力頭表現(xiàn)出各異的注意力模式和擴(kuò)展規(guī)則：有的關(guān)注全局信息，有的則聚焦局部；有的注意力范圍隨輸入長度增加而擴(kuò)展，有的則保持不變。然而，現(xiàn)有的統(tǒng)一稀疏注意力機(jī)制破壞了這些固有的特性。

為應(yīng)對這一挑戰(zhàn)，研究團(tuán)隊(duì)提出了混合稀疏注意力（Mixture of Sparse Attention, MoA）方法，它能夠?yàn)椴煌念^和層定制獨(dú)特的稀疏注意力配置。MoA 構(gòu)建了一個包含多種注意力模式及其擴(kuò)展規(guī)則的搜索空間。通過分析模型，評估潛在配置，MoA 可以為每個注意力頭找到最優(yōu)的稀疏注意力模式和擴(kuò)展規(guī)則。

實(shí)驗(yàn)結(jié)果顯示，無需任何訓(xùn)練，MoA 就可以在保持平均注意力跨度不變的情況下，將有效上下文長度提升約 3.9 倍。模型效果上，在 Vicuna-7B、Vicuna-13B 和 Llama3-8B 模型上，MoA 將長文本信息檢索準(zhǔn)確率提高了 1.5-7.1 倍，優(yōu)于統(tǒng)一注意力基線方法 StreamingLLM。

此外，MoA 縮小了稀疏與稠密模型之間的能力差距，在 50% 平均注意力跨度下，長上下文理解基準(zhǔn)測試集的最大相對性能下降從基線方法的 9%-36% 降低至 5% 以內(nèi)。

在運(yùn)行效率上，MoA 的稀疏注意力使得生成過程中 KV-Cache 長度不擴(kuò)大便于內(nèi)存管理，減少了注意力計(jì)算量，降低了存儲量從而可增大批大小。結(jié)合 CUDA GPU 算子優(yōu)化，MoA 相比于 FlashAttention2 和 vLLM 將 7B 和 13B 稠密模型的生成吞吐量分別可提升 6.6-8.2 以及 1.7-1.9 倍。方法在 12K 以內(nèi)的輸入長度上搜索壓縮方案，壓縮后模型可以在長達(dá) 256K 的輸入長度上高效且精準(zhǔn)地檢索信息。

背景介紹

多頭自注意力（MHA）機(jī)制是大型語言模型的核心功能之一 [3]。該機(jī)制從輸入序列出發(fā)，通過線性變換將其轉(zhuǎn)換成查詢（Q）、鍵（K）和值（V）三個矩陣。這些矩陣與之前序列的鍵值緩存（KV-Cache）相結(jié)合，共同計(jì)算出注意力矩陣（A）。為了保持自回歸特性，這一計(jì)算過程會通過因果掩膜（M）進(jìn)行調(diào)整，最終得到輸出（O）。具體公式如下：

在大語言模型的自回歸推理過程中，分為兩個階段：預(yù)填充和解碼。在預(yù)填充階段，模型會處理整個輸入序列，以生成初始的響應(yīng)令牌。隨后進(jìn)入解碼階段，模型利用新產(chǎn)生的令牌以及之前緩存的 K 和 V 矩陣，逐步生成后續(xù)令牌，直至完成整個序列的生成。雖然這種迭代方法效果顯著，但隨著 KV-Cache 的不斷擴(kuò)展，它也帶來了內(nèi)存和計(jì)算資源的需求增加。

現(xiàn)有方法

之前的研究通過引入稀疏注意力方法來應(yīng)對大型語言模型處理長上下文時的效率挑戰(zhàn)。對于生成式的大型語言模型，主流的稀疏模式是采用統(tǒng)一跨度滑窗：即不論注意力頭還是輸入長度如何，都使用固定、均勻跨度的滑動窗口掩膜，這樣每個文本僅關(guān)注其鄰近的上下文區(qū)域。

此外，還會對最初的幾個文本施加全局注意力，以便它們能夠關(guān)注到所有其他文本。這種局部注意模式通過丟棄當(dāng)前注意跨度之外的 KV-Cache，顯著降低了長序列場景下的內(nèi)存需求 [4][5]。原則上，盡管單個模型層的注意力是局部的，但通過多層模型的逐步傳遞，每個詞最終都能獲取全局信息，從而在理論上可以實(shí)現(xiàn)比平均注意力跨度更長的有效上下文長度 [6]。

然而，這種統(tǒng)一跨度的滑動窗口方法并未考慮到模型本身的特性，導(dǎo)致大型模型在處理長文本時的有效上下文長度受到限制，進(jìn)而影響了其在長文本場景下的表現(xiàn)。

根據(jù)之前的研究定義，本工作將有效上下文長度定義為在內(nèi)容檢索任務(wù)上能夠達(dá)到 90% 以上精度的最大輸入長度 [1][2]。研究表明，像 StreamingLLM [4] 這樣的統(tǒng)一跨度滑窗方法，其有效上下文長度往往難以超出平均注意力跨度。如下圖所示，當(dāng)使用輸入長度 50% 的跨度進(jìn)行稀疏注意力時，統(tǒng)一跨度滑窗無法有效地從窗口外的內(nèi)容中檢索信息，而且這一問題隨著輸入長度的增加而愈發(fā)嚴(yán)重。

在LongEval數(shù)據(jù)集上[7]，使用不同注意力方法的Vicuna-7B模型在不同輸入長度和檢索位置的檢索精度。大海撈針實(shí)驗(yàn)[8]將大量的鍵值對作為輸入，并在不同位置測試給鍵檢索值時的精度。(a) 原始稠密注意力模型；(b)統(tǒng)一跨度滑窗StreamingLLM，注意力跨度減半，超出跨度時檢索效果降低；(c) MoA，平均注意力跨度減半，超出跨度時檢索效果依然優(yōu)秀。

下圖揭示了這個現(xiàn)象的一個可能解釋：雖然部分注意力頭專注于局部上下文，但另一些注意力頭則關(guān)注了整個輸入序列。因此，采用統(tǒng)一跨度的方法限制了那些關(guān)注全局上下文的頭的注意力跨度，并且給關(guān)注局部上下文的頭分配了過多的計(jì)算和內(nèi)存資源。

此外，隨著輸入長度的增加，某些注意力頭需要比其他頭更快地?cái)U(kuò)展其注意力跨度，以防止性能顯著下降。遺憾的是，統(tǒng)一跨度的方法沒有考慮到這種異質(zhì)性，未能針對不同注意力頭的需求分別調(diào)整其注意力范圍。另外，現(xiàn)有的模型壓縮技術(shù)在制定壓縮策略時通?；?strong>通用語料庫和人工撰寫的輸出結(jié)果，這并不能精確地反映稀疏化對處理長上下文任務(wù)的具體影響。

來自Vicuna-7B模型不同注意力頭的注意力矩陣示例。每個注意力矩陣是從LongEval數(shù)據(jù)集的256個輸入上取平均得到的。

本文方法

本文提出了一種名為混合注意力（MoA）的方法，它是一種無需訓(xùn)練或微調(diào)的異質(zhì)稀疏注意力機(jī)制。如下圖所示，MoA 建立了一套異質(zhì)彈性規(guī)則，作為注意力跨度的搜索空間。對于每個注意力頭，MoA 將自動分析不同注意力跨度的效果，并最優(yōu)化不同注意力頭的跨度。同時，MoA 精心設(shè)計(jì)了校準(zhǔn)數(shù)據(jù)集，以確保它能精確地反映出稀疏注意力對處理長上下文任務(wù)的具體影響。

MoA的概覽：(a) 稀疏注意力的搜索空間涵蓋了滑窗跨度的異構(gòu)彈性規(guī)則；(b) 自動壓縮過程始于精心設(shè)計(jì)的校準(zhǔn)數(shù)據(jù)集。MoA通過分析這個數(shù)據(jù)集中每個注意力值對模型預(yù)測的影響，揭示了不同候選彈性規(guī)則在不同輸入長度下的準(zhǔn)確性損失。在優(yōu)化階段，MoA為每個注意力頭挑選出最合適的彈性規(guī)則，使得其在滿足平均跨度約束的同時，盡可能減少模型的性能損失。

異質(zhì)彈性規(guī)則的搜索空間

在構(gòu)建 MoA 稀疏注意力模式的搜索空間時，本工作充分考慮了大語言模型注意力模式的天然異質(zhì)性和彈性特點(diǎn)。正如圖 (a) 展示的那樣，本工作選用了對硬件友好的異質(zhì)跨度滑窗作為本工作稀疏注意力的掩膜。同之前研究 [4][5] 類似，本工作對最初的幾個文本施加了全局注意力。

為了捕捉注意力跨度隨輸入長度變化的不同表現(xiàn)，本工作將注意力頭 h 的注意力跨度 S 定義為輸入長度 N 的函數(shù)，即。其中，和 是控制基本跨度和其隨輸入長度變化的擴(kuò)展比例的超參數(shù)。每個注意力頭的  和  都可以從一系列離散選項(xiàng)中挑選

自動注意力跨度規(guī)則搜索

對于有很多層的大語言模型，異質(zhì)彈性規(guī)則的搜索空間可能會變得非常大。因此，本工作設(shè)計(jì)了自動注意力跨度規(guī)則搜索方法，來對大語言模型的每一個注意力頭確定最優(yōu)的  和 。這包括分析和優(yōu)化兩步。

分析

給定大語言模型，MoA 首先會在校準(zhǔn)數(shù)據(jù)集的一系列文本上進(jìn)行分析，以評估移除每個注意力值對模型最終預(yù)測結(jié)果的具體影響。具體而言，本工作采用了一階泰勒展開的方法進(jìn)行影響評估：

通過分析每個注意力值對預(yù)測結(jié)果的貢獻(xiàn)，我們可以計(jì)算出在當(dāng)前輸入下，不同跨度滑窗對最終預(yù)測結(jié)果的影響，也就是該窗口所移除的所有注意力值影響的總和。

在實(shí)際操作中，本工作推導(dǎo)了注意力影響的形式化表達(dá)，并利用深度學(xué)習(xí)框架的反向傳播機(jī)制高效地計(jì)算所需的偏導(dǎo)數(shù)值。在校準(zhǔn)數(shù)據(jù)集的不同輸入長度上，MoA 對滑動窗口的影響分別取平均值，以體現(xiàn)同一種異質(zhì)彈性規(guī)則在不同長度輸入下的影響。

在分析階段完成后，MoA 能夠明確每種異質(zhì)彈性規(guī)則在精度和效率之間的平衡?；诖?，MoA 可以將較長的注意力跨度分配給那些對壓縮更為敏感的注意力頭，而將較短的注意力跨度分配給那些對此不太敏感的注意力頭。

優(yōu)化

根據(jù)分析結(jié)果，MoA 會為每個注意力頭挑選出最佳彈性規(guī)則。在優(yōu)化過程中，用戶可以設(shè)定不同輸入長度下的注意力密度（即平均注意力跨度與輸入長度的比值）作為限制條件，而 MoA 則會在滿足這一條件的同時，力求最小化各長度下的預(yù)測誤差。

具體而言，MoA 構(gòu)建了一個多目標(biāo)優(yōu)化問題：最小化不同輸入長度下的預(yù)測損失，同時確保平均注意力密度不超過用戶設(shè)定的限制。MoA 采用混合整數(shù)優(yōu)化方法求解這一多目標(biāo)優(yōu)化問題。求解得到的所有壓縮方案均達(dá)到了帕累托最優(yōu)，意味著在沒有增加其他長度預(yù)測損失的前提下，無法進(jìn)一步降低任何特定長度的預(yù)測損失。

為了確保模型能夠有效泛化至未曾見過的輸入長度，MoA 會從所有帕累托最優(yōu)壓縮方案中，選擇在未見過的驗(yàn)證長度上損失最小的方案作為最終采用的壓縮策略。

通過自動注意力跨度規(guī)則搜索，MoA 在遵守用戶定義的密度約束的同時，找到合適的異質(zhì)彈性規(guī)則來最小化由注意力稀疏化引起的準(zhǔn)確性損失。

校準(zhǔn)數(shù)據(jù)集的設(shè)計(jì)與選擇

同時，MoA 也強(qiáng)調(diào)了數(shù)據(jù)工程在大語言模型壓縮中的重要性。本工作發(fā)現(xiàn)，使用具有長距離依賴性的數(shù)據(jù)集并參考原始大語言模型的響應(yīng)對于準(zhǔn)確分析壓縮的影響至關(guān)重要。

本工作指出了常用的通用語言建模數(shù)據(jù)集的主要問題。這類數(shù)據(jù)集，例如人類編寫的文本語料庫，通過在整個語料庫上進(jìn)行下一個詞預(yù)測作為監(jiān)督信號。但是這主要捕獲的是臨近上下文之間的注意力模式，而忽略了長期上下文依賴性，無法解決像長距離檢索這樣的全局注意力任務(wù)。

同時，模型響應(yīng)和人類編寫的監(jiān)督之間存在顯著的不對齊。例如，對于同一個問題，人類可能會回答 'Blue'，而模型可能會回答 'The blue color'。使用人類的答案進(jìn)行監(jiān)督，注意力影響是基于預(yù)測 'Blue' 的概率轉(zhuǎn)移量化的，這與最終目標(biāo)背道而馳，即難以保持原始模型預(yù)測 'The' 的關(guān)鍵注意力。

因此，本工作構(gòu)建長距離依賴并通過與原始模型對齊來增強(qiáng)校準(zhǔn)數(shù)據(jù)集。通過下表可以發(fā)現(xiàn)，這種數(shù)據(jù)集構(gòu)建方式可以準(zhǔn)確反映注意力影響，顯著提高壓縮后的模型的性能。

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

精度

MoA 實(shí)驗(yàn)在多種模型（Vicuna-{7B, 13B) 和 Llama-3-{8B, 70B}）和多種基準(zhǔn)測試（長上下文檢索，長上下文理解）上和之前的靜態(tài)和動態(tài)稀疏注意力方法（StreamingLLM [4],H2O [9] 和 InfLLM [12]）進(jìn)行了比較。

MoA 超越基線稀疏注意力方法，并在 50% 的平均注意力密度下達(dá)到了與原始稠密模型相當(dāng)?shù)男阅?。我們?jì)算了稀疏模型相對于原始模型的性能下降。

對于長上下文檢索檢索任務(wù)，MoA 最大相對準(zhǔn)確性下降為 8%，遠(yuǎn)小于 StreamingLLM、InfLLM 和  H2O 的 87%、58% 和 44%。平均來說，MoA 的相對準(zhǔn)確性下降在 1% 以下，而 StreamingLLM、InfLLM 和 H2O 的  51%、41% 和 20%。

如下圖 (a) 所示，MoA 將其有效上下文長度擴(kuò)展到注意力跨度的大約 3.9 倍。圖 (b) 顯示，在固定的 8k 輸入長度下，MoA 只需要 25% 注意力就可以達(dá)到 90% 以上的檢索準(zhǔn)確性。圖 (c) 顯示 MoA 在 12k 長度內(nèi)壓縮后，可以保持和原稠密模型一致的有效上下文長度。

在LongEval上進(jìn)行上下文檢索準(zhǔn)確性測試。使用Vicuna-7B模型：(a)改變注意力跨度，比較有效上下文長度，(b)將輸入長度設(shè)定為8k，比較不同注意力密度下的檢索準(zhǔn)確性。使用Llama3-8B模型：(c)將密度設(shè)定為50%，比較輸入長度增加時的檢索準(zhǔn)確性。

對于長上下文理解任務(wù)，在使用 LV-Eval 和 LongBench 的基準(zhǔn)測試中，MoA 只顯示了最大 5% 和 3% 的相對分?jǐn)?shù)下降，而 StreamingLLM 則分別最大下降了 36% 和 27%；InfLLM 最大下降了 17% 和 5%；H2O 最大下降了 9% 和 4%

不同注意力方法在 50% 密度下進(jìn)行 (a) LV-Eval和(b) LongBench 長上下文理解基準(zhǔn)測試。測試使用 Vicuna-7B和13B模型，以及 Llama3-70B 模型。分?jǐn)?shù)相對于原始稠密模型進(jìn)行歸一化。

長上下文泛化。通過在 12k 長度內(nèi)進(jìn)行壓縮，MoA 能夠有效泛化到 32k-256k 的長度。如下表所示，在范化的長度上，MoA 的檢索精度比 InfLLM 和 StreamingLLM 高 1.9-3.3 倍，LV-Eval 評分高 1.2-1.4 倍，展現(xiàn)出與原始稠密模型相當(dāng)?shù)男阅堋Ｈ缦卤硭荆?/span>

消融實(shí)驗(yàn)進(jìn)一步評估了不同程度的稀疏注意力混合對于最終性能的影響。從基本的統(tǒng)一跨度滑窗開始，通過依次引入不同程度的異質(zhì)性（層間，注意力頭間，輸入長度間），最終模型性能在不斷提升。

在 25% 注意力密度下對搜索空間進(jìn)行消融實(shí)驗(yàn)。

效率

我們使用 Huggingface 框架支持的 FlashAttention [10] 和 vLLM 框架作為基線，對比 MoA 的效率。

下表比較了 MoA 相對于各種注意力機(jī)制和 LLM 框架的運(yùn)行效率，并對 MoA 的每個設(shè)計(jì)帶來的效率提升進(jìn)行了消融分析。在 50% 注意力密度下，MoA 相比 FlashAttention2 將解碼吞吐量提升了 6.6-8.2 倍。相比 H2O 和 InfLLM，解碼吞吐量提升了 1.2-4.0 倍。與包含高度系統(tǒng)級優(yōu)化的 vLLM 框架 [11] 相比，MoA 仍實(shí)現(xiàn)了 1.7-1.9 倍的吞吐量提升。MoA 還將 GPU 總內(nèi)存減少 1.2 到 1.4 倍。

這些吞吐量的提升來源于四個主要因素：生成過程中的靜態(tài) KV-Cache（約 3.0 倍）；由于稀疏性減少了注意力計(jì)算（約 1.5 倍）；較小的 KV-Cache 內(nèi)存支持了更大的批大小（約 1.4 倍）；以及我們針對 MoA 異構(gòu)注意力所實(shí)現(xiàn)的 CUDA GPU 算子優(yōu)化（約 1.2 倍）。

不同框架在7B和13B模型上的效率分析。MoA 每個設(shè)計(jì)帶來的效率提升通過消融分析分為四個部分。所有稀疏注意力方法都使用50%的注意力密度。解碼吞吐量在A100-80GB GPU 顯存能容納的最大批大小下進(jìn)行評估。

作者介紹

本論文的共同一作是清華大學(xué)電子工程系 NICS-EFC 實(shí)驗(yàn)室的傅天予、黃浩峰和寧雪妃，他們來自 NICS-EFC 實(shí)驗(yàn)室的 EffAlg 團(tuán)隊(duì)和無問芯穹（Infinigence AI）。NICS-EFC 實(shí)驗(yàn)室由汪玉教授帶領(lǐng)，實(shí)驗(yàn)室的高效算法團(tuán)隊(duì)（Efficient Algorithm Team，EffAlg）由寧雪妃助理研究員帶領(lǐng)。EffAlg 團(tuán)隊(duì)的主要研究方向?yàn)楦咝疃葘W(xué)習(xí)技術(shù)，團(tuán)隊(duì)網(wǎng)站為 https://nics-effalg.com/