無問芯穹提出混合稀疏注意力方案MoA，實現最高8倍吞吐率提升

機器之心發佈

機器之心編輯部

隨着大語言模型在長文本場景下的需求不斷涌現，其核心的注意力機制（Attention Mechanism）也獲得了非常多的關注。

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

然而，現有稀疏注意力通常採用統一的稀疏注意力模式，即對不同的注意力頭和輸入長度應用相同的稀疏模式。這種統一的方法難以捕捉到大語言模型中多樣的注意力模式，導致不同注意力頭的不同的精度 - 代價權衡被忽略。

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

本工作現已開源，歡迎交流討論。

總覽

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

爲應對這一挑戰，研究團隊提出了混合稀疏注意力（Mixture of Sparse Attention, MoA）方法，它能夠爲不同的頭和層定製獨特的稀疏注意力配置。MoA 構建了一個包含多種注意力模式及其擴展規則的搜索空間。通過分析模型，評估潛在配置，MoA 可以爲每個注意力頭找到最優的稀疏注意力模式和擴展規則。

實驗結果顯示，無需任何訓練，MoA 就可以在保持平均注意力跨度不變的情況下，將有效上下文長度提升約3.9倍。模型效果上，在 Vicuna-7B、Vicuna-13B 和 Llama3-8B 模型上，MoA 將長文本信息檢索準確率提高了1.5-7.1倍，優於統一注意力基線方法 StreamingLLM。

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

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

背景介紹

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

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

現有方法

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

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

然而，這種統一跨度的滑動窗口方法並未考慮到模型本身的特性，導致大型模型在處理長文本時的有效上下文長度受到限制，進而影響了其在長文本場景下的表現。

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

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

下圖揭示了這個現象的一個可能解釋：雖然部分注意力頭專注於局部上下文，但另一些注意力頭則關注了整個輸入序列。因此，採用統一跨度的方法限制了那些關注全局上下文的頭的注意力跨度，並且給關注局部上下文的頭分配了過多的計算和內存資源。

此外，隨着輸入長度的增加，某些注意力頭需要比其他頭更快地擴展其注意力跨度，以防止性能顯著下降。遺憾的是，統一跨度的方法沒有考慮到這種異質性，未能針對不同注意力頭的需求分別調整其注意力範圍。另外，現有的模型壓縮技術在制定壓縮策略時通常基於通用語料庫和人工撰寫的輸出結果，這並不能精確地反映稀疏化對處理長上下文任務的具體影響。

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

本文方法

本文提出了一種名爲混合注意力（MoA）的方法，它是一種無需訓練或微調的異質稀疏注意力機制。如下圖所示，MoA 建立了一套異質彈性規則，作爲注意力跨度的搜索空間。對於每個注意力頭，MoA 將自動分析不同注意力跨度的效果，並最優化不同注意力頭的跨度。同時，MoA 精心設計了校準數據集，以確保它能精確地反映出稀疏注意力對處理長上下文任務的具體影響。

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

異質彈性規則的搜索空間

自動注意力跨度規則搜索

通過分析每個注意力值對預測結果的貢獻，我們可以計算出在當前輸入下，不同跨度滑窗對最終預測結果的影響，也就是該窗口所移除的所有注意力值影響的總和。

在實際操作中，本工作推導了注意力影響的形式化表達，並利用深度學習框架的反向傳播機制高效地計算所需的偏導數值。在校準數據集的不同輸入長度上，MoA 對滑動窗口的影響分別取平均值，以體現同一種異質彈性規則在不同長度輸入下的影響。

在分析階段完成後，MoA 能夠明確每種異質彈性規則在精度和效率之間的平衡。基於此，MoA 可以將較長的注意力跨度分配給那些對壓縮更爲敏感的注意力頭，而將較短的注意力跨度分配給那些對此不太敏感的注意力頭。

優化

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

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

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

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

校準數據集的設計與選擇

同時，MoA 也強調了數據工程在大語言模型壓縮中的重要性。本工作發現，使用具有長距離依賴性的數據集並參考原始大語言模型的響應對於準確分析壓縮的影響至關重要。

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

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

因此，本工作構建長距離依賴並通過與原始模型對齊來增強校準數據集。通過下表可以發現，這種數據集構建方式可以準確反映注意力影響，顯著提高壓縮後的模型的性能。

實驗結果

精度

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

MoA 超越基線稀疏注意力方法，並在 50% 的平均注意力密度下達到了與原始稠密模型相當的性能。我們計算了稀疏模型相對於原始模型的性能下降。

對於長上下文檢索檢索任務，MoA 最大相對準確性下降爲 8%，遠小於 StreamingLLM、InfLLM 和 H2O 的 87%、58% 和 44%。平均來說，MoA 的相對準確性下降在 1% 以下，而 StreamingLLM、InfLLM 和 H2O 的 51%、41% 和 20%。

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

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

對於長上下文理解任務，在使用 LV-Eval 和 LongBench 的基準測試中，MoA 只顯示了最大 5% 和 3% 的相對分數下降，而 StreamingLLM 則分別最大下降了 36% 和 27%；InfLLM 最大下降了 17% 和 5%；H2O 最大下降了 9% 和 4%

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

長上下文泛化。通過在 12k 長度內進行壓縮，MoA 能夠有效泛化到 32k-256k 的長度。如下表所示，在範化的長度上，MoA 的檢索精度比 InfLLM 和 StreamingLLM 高 1.9-3.3 倍，LV-Eval 評分高 1.2-1.4 倍，展現出與原始稠密模型相當的性能。如下表所示，

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

在 25% 注意力密度下對搜索空間進行消融實驗。

效率

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

下表比較了 MoA 相對於各種注意力機制和 LLM 框架的運行效率，並對 MoA 的每個設計帶來的效率提升進行了消融分析。在 50% 注意力密度下，MoA 相比 FlashAttention2 將解碼吞吐量提升了 6.6-8.2 倍。相比 H2O 和 InfLLM，解碼吞吐量提升了 1.2-4.0 倍。與包含高度系統級優化的 vLLM 框架 [11] 相比，MoA 仍實現了 1.7-1.9 倍的吞吐量提升。MoA 還將 GPU 總內存減少 1.2 到 1.4 倍。

這些吞吐量的提升來源於四個主要因素：生成過程中的靜態 KV-Cache（約 3.0 倍）；由於稀疏性減少了注意力計算（約 1.5 倍）；較小的 KV-Cache 內存支持了更大的批大小（約 1.4 倍）；以及我們針對 MoA 異構注意力所實現的 CUDA GPU 算子優化（約 1.2 倍）。

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

作者介紹

本論文的共同一作是清華大學電子工程系 NICS-EFC 實驗室的傅天予、黃浩峰和寧雪妃，他們來自 NICS-EFC 實驗室的 EffAlg 團隊和無問芯穹（Infinigence AI）。NICS-EFC 實驗室由汪玉教授帶領，實驗室的高效算法團隊（Efficient Algorithm Team，EffAlg）由寧雪妃助理研究員帶領。EffAlg 團隊的主要研究方向爲高效深度學習技術，團隊網站爲 https://nics-effalg.com/

引用

[1] Chen, Shouyuan, et al. "Extending Context Window of Large Language Models via Positional Interpolation." ArXiv, 2023, abs/2306.15595, https://api.semanticscholar.org/CorpusID:259262376.

[2] Tworkowski, Szymon, et al. "Focused Transformer: Contrastive Training for Context Scaling." ArXiv, 2023, abs/2307.03170, https://api.semanticscholar.org/CorpusID:259360592.

[3] Vaswani, Ashish, et al. "Attention is all you need." Advances in Neural Information Processing Systems, vol. 30, 2017.

[4] Xiao, Guangxuan, et al. "Efficient Streaming Language Models with Attention Sinks." The Twelfth International Conference on Learning Representations, 2024.

[5] Han, Chi, et al. "Lm-infinite: Simple on-the-fly length generalization for large language models." arXiv preprint arXiv:2308.16137, 2023.

[6] Zaheer, Manzil, et al. "Big bird: Transformers for longer sequences." Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 17283-17297.

[7] Li, Dacheng, et al. "How Long Can Open-Source LLMs Truly Promise on Context Length?" lmsys.org, June 2023, https://lmsys.org/blog/2023-06-29-longchat.

[8] Fu, Yao, et al. "Data Engineering for Scaling Language Models to 128K Context." ArXiv, 2024, abs/2402.10171, https://api.semanticscholar.org/CorpusID:267682361.

[9] Zhang, Zhenyu (Allen), et al. "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models." ArXiv, 2023, abs/2306.14048, https://api.semanticscholar.org/CorpusID:259263947.

[10] Dao, Tri, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." Advances in Neural Information Processing Systems, 2022.

[11] Kwon, Woosuk, et al. "Efficient Memory Management for Large Language Model Serving with PagedAttention." Proceedings of the 29th Symposium on Operating Systems Principles, 2023, https://api.semanticscholar.org/CorpusID:261697361.

[12] Xiao, Chaojun et al. “InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory.” (2024)