無問芯穹提出混合稀疏注意力方案MoA，加速長文字生成，實現最高8倍吞吐率提升

机器之心發表於2024-11-08
隨著大語言模型在長文字場景下的需求不斷湧現，其核心的注意力機制（Attention Mechanism）也獲得了非常多的關注。

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

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

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

本工作現已開源，歡迎交流討論。
圖片
  • 程式碼：https://github.com/thu-nics/MoA
  • 主頁：https://nics-effalg.com/MoA
  • arXiv：https://arxiv.org/abs/2406.14909

總覽

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

為應對這一挑戰，研究團隊提出了混合稀疏注意力（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 稀疏注意力模式的搜尋空間時，本工作充分考慮了大語言模型注意力模式的天然異質性和彈性特點。正如圖 (a) 展示的那樣，本工作選用了對硬體友好的異質跨度滑窗作為本工作稀疏注意力的掩膜。同之前研究 [4][5] 類似，本工作對最初的幾個文字施加了全域性注意力。

為了捕捉注意力跨度隨輸入長度變化的不同表現，本工作將注意力頭 h 的注意力跨度 S 定義為輸入長度 N 的函式，即圖片。其中，圖片和 圖片是控制基本跨度和其隨輸入長度變化的擴充套件比例的超引數。每個注意力頭的 圖片 和 圖片 都可以從一系列離散選項中挑選

自動注意力跨度規則搜尋

對於有很多層的大語言模型，異質彈性規則的搜尋空間可能會變得非常大。因此，本工作設計了自動注意力跨度規則搜尋方法，來對大語言模型的每一個注意力頭確定最優的 圖片 和 圖片。這包括分析和最佳化兩步。

分析

給定大語言模型，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).

相關文章