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

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

然而，現(xiàn)有稀疏注意力通常采用統(tǒng)一的稀疏注意力模式，即對(duì)不同的注意力頭和輸入長(zhǎng)度應(yīng)用相同的稀疏模式。這種統(tǒng)一的方法難以捕捉到大語言模型中多樣的注意力模式，導(dǎo)致不同注意力頭的不同的精度 - 代價(jià)權(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)注全局信息，有的則聚焦局部；有的注意力范圍隨輸入長(zhǎng)度增加而擴(kuò)展，有的則保持不變。然而，現(xiàn)有的統(tǒng)一稀疏注意力機(jī)制破壞了這些固有的特性。

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

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

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

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

背景介紹

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

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

現(xiàn)有方法

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

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

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

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

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

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

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

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

本文方法

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

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

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

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

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

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

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

分析

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

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

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

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

優(yōu)化

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

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

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

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

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

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

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

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

因此，本工作構(gòu)建長(zhǎng)距離依賴并通過與原始模型對(duì)齊來增強(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)測(cè)試（長(zhǎng)上下文檢索，長(zhǎng)上下文理解）上和之前的靜態(tài)和動(dòng)態(tài)稀疏注意力方法（StreamingLLM [4],H2O [9] 和 InfLLM [12]）進(jìn)行了比較。

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

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

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

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

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

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

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

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

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

效率

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

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

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

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

作者介紹

本論文的共同一作是清華大學(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/