撞車DeepSeek NSA，Kimi楊植麟署名的新注意力架構MoBA釋出，程式碼也公開
机器之心發表於2025-02-19
昨天下午，DeepSeek 釋出了一篇新論文，提出了一種改進版的注意力機制 NSA；加上還有創始人兼 CEO 梁文鋒親自參與，一時之間吸引眼球無數，參閱報導《剛剛！DeepSeek 梁文鋒親自掛名，公開新注意力架構 NSA》。但其實就在同一天，月之暗面也釋出了一篇主題類似的論文，並且同樣巧合的是，月之暗面創始人兼 CEO 楊植麟也是該論文的署名作者之一。並且不同於 DeepSeek 只發布了一篇論文，月之暗面還發布了相關程式碼。且這些程式碼已有一年的實際部署驗證，有效性和穩健性都有保證。這篇論文提出了一種名為 MoBA 的注意力機制，即 Mixture of Block Attention，可以直譯為「塊注意力混合」。據介紹，MoBA 是「一種將混合專家（MoE）原理應用於注意力機制的創新方法。」該方法遵循「更少結構」原則，並不會引入預定義的偏見，而是讓模型自主決定關注哪些位置。- 論文標題：MoBA: Mixture of Block Attention for Long-Context LLMs
- 論文地址：https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf
- 專案地址：https://github.com/MoonshotAI/MoBA
與 DeepSeek 的 NSA 注意力機制新論文一樣，月之暗面這篇 MoBA 論文也收穫了諸多好評。最近一段時間，大語言模型（LLM）發展到了前所未有的規模，我們希望它們能夠處理模擬人類認知的複雜任務。具體來說，實現 AGI 的一個關鍵能力是處理、理解和生成長序列的能力，這對於從歷史資料分析到複雜推理和決策過程的各種應用至關重要。這種對擴充套件上下文處理能力的需求不僅體現在長輸入提示理解的普及上，還體現在最近 Kimi k1.5、DeepSeek-R1 和 OpenAI o1/o3 中展示的，對長鏈思維（CoT）輸出能力的探索中。然而，由於傳統注意力機制帶來的計算複雜度呈二次函式增長，擴充套件 LLM 的序列長度並非易事。這一挑戰催生了一波研究浪潮，其中一個重要方向是利用注意力分數的固有稀疏性。這種稀疏性既源於數學上的 softmax 操作，也源於生物學上的觀察，即人們已在與記憶儲存相關的大腦區域中觀察到了稀疏連線。現有方法通常利用預定義的結構約束（例如基於 sink 的注意力機制或滑動視窗注意力機制）來利用這種稀疏性。雖然這些方法可能有效，但它們往往高度依賴於特定任務，可能會限制模型的整體泛化能力。另一種方法是採用動態稀疏注意力機制，例如 Quest、Minference 和 RetrievalAttention，這些方法在推理時選擇部分 token 進行計算。儘管這些方法可以減少長序列的計算量，但它們並未顯著緩解長上下文模型的高昂訓練成本。最近，線性注意力模型作為一種有前景的替代方案出現，例如 Mamba、RWKV 和 RetNet。這些方法用線性近似替代了傳統的基於 softmax 的注意力機制，從而降低了長序列處理的計算開銷。然而，由於線性注意力與傳統注意力機制存在顯著差異，適配現有的 Transformer 模型通常需要高昂的轉換成本，或者需要從頭開始訓練全新的模型。更重要的是，它們在複雜推理任務中的有效性證據仍然有限。因此，一個關鍵的研究問題隨之而來：我們如何設計一種穩健且適應性強的注意力架構，既能保留原始 Transformer 框架，又能遵循「更少結構」的原則，使模型能夠在不依賴預定義偏差的情況下自主決定關注哪些部分？理想情況下，這種架構能夠在完全注意力和稀疏注意力模式之間無縫切換，從而最大限度地與現有預訓練模型相容，並在不犧牲效能的前提下實現高效推理和加速訓練。為此研究人員提出了塊注意力混合機制（Mixture of Block Attention, MoBA），這是一種全新架構，它基於混合專家系統（MoE）的創新原理，並將其應用於 Transformer 模型的注意力機制中。MoE 此前主要應用於 Transformer 的前饋網路（FFN）層，而 MoBA 首次將其引入長上下文注意力機制中，允許動態選擇與每個查詢 token 相關的歷史關鍵塊和值塊。這種方法不僅提高了 LLM 的效率，還使其能夠處理更長、更復雜的提示，而無需按比例增加資源消耗。相比之下，MoBA 透過將上下文劃分為塊並採用門控機制選擇性地將查詢 token 路由到最相關的塊，解決了傳統注意力機制的計算效率低下的問題。這種塊稀疏注意力顯著降低了計算成本，為更高效地處理長序列鋪平了道路。模型能夠動態選擇最具資訊量的關鍵塊，從而提高了效能和效率，這對於涉及大量上下文資訊的任務尤為有益。在該工作中，月之暗面詳細介紹了 MoBA 的架構，首先是其塊劃分和路由策略，其次是與傳統注意力機制相比的計算效率。他們也進行了實驗，證明了 MoBA 在處理長序列任務中的卓越效能。MoBA 透過動態選擇歷史片段（塊）進行注意力計算，擴充套件了 Transformer 模型的能力。MoBA 的靈感來源於混合專家（Mixture of Experts，簡稱 MoE）技術和稀疏注意力技術。前者主要應用於 Transformer 架構中的前饋網路（Feedforward Network，簡稱 FFN）層，而後者已被廣泛用於擴充套件 Transformer 以處理長文字上下文。本文創新點在於月之暗面將 MoE 原理應用於注意力機制本身，從而實現對長序列更高效和更有效的處理。- 可訓練的塊稀疏注意力：將完整上下文分割成若干塊，每個查詢 token 都會學習關注最相關的 KV 塊，從而實現長序列的高效處理。
- 無引數門控機制：引入了一種新穎的無引數 top-k 門控機制，為每個查詢 token 選擇最相關的塊，確保模型只關注最具資訊量的內容。
- 完全注意力與稀疏注意力的無縫切換：MoBA 被設計為完全注意力機制的靈活替代方案，能夠在完全注意力和稀疏注意力模式之間實現平滑過渡。
總體而言，MoBA 注意力機制使模型能夠自適應且動態地關注上下文中最有資訊量的塊。這在處理涉及長文件或序列的任務時尤其有益，因為在這些任務中，關注整個上下文可能是不必要的，並且計算代價高昂。MoBA 選擇性地關注相關塊的能力，使得資訊處理更加細緻且高效。月之暗面透過整合來自 FlashAttention 和 MoE 的最佳化技術，實現了 MoBA 的高效能版本。圖 2 展示了 MoBA 的高效率。- 根據門控網路和因果掩碼，確定查詢 token 對 KV 塊的分配。
- 根據分配的 KV 塊，安排查詢 token 的順序。
- 計算每個 KV 塊的注意力輸出和分配給它的查詢 token，這一步可以透過不同長度的 FlashAttention 來最佳化。
- 使用線上 Softmax 合併相應的注意力輸出，因為一個查詢 token 可能會關注其當前塊以及多個歷史 KV 塊。
下圖為該演算法的形式化描述，並在圖 1b 中進行了視覺化，展示瞭如何基於 MoE 和 FlashAttention 實現 MoBA。為了驗證 MoBA 的實際效果，月之暗面進行了一系列 Scaling Law 和消融實驗。Full Attention 與 MoBA 對比實驗該團隊訓練了 5 個模型來進行比較，表 1 給出詳細配置。結果見圖 3a，可以看到，Full Attention 與 MoBA 的 scaling 趨勢非常相似。具體來說，這兩種注意力機制之間的驗證損失差異在 1e − 3 的範圍內保持一致。這表明，儘管 MoBA 的稀疏注意力模式稀疏度高達 75%，但它實現了與完全注意力相當的擴充套件效能。此外，該團隊也驗證了 MoBA 的上下文可擴充套件性。在實驗中，該團隊將最大序列長度從 8k 提升到了 32k。結果見圖 3b，可見儘管在所有五項實驗中，MoBA 最後的塊 LM 損失與 Full Attention 相比略高，但損失差距正在逐步縮小。該實驗表明 MoBA 具有長上下文可擴充套件性。該團隊也透過消融實驗對 MoBA 的細粒度塊分割的有效性進行了驗證，結果見圖 4。可以看到，MoBA 的效能受塊粒度的顯著影響。具體來說，最粗粒度設定（從 8 個塊中選擇 2 個塊）與更細粒度的設定之間的效能差異為 1e-2。這些發現表明，細粒度分割似乎是增強 MoE 系列（包括 MoBA）模型效能的普適技術。將 MoBA 與 Full Attention 融合到一起會如何？MoBA 本就是作為 Full Attention 的靈活替代而設計的，可以最小的開銷輕鬆地互相切換，並實現相當的長上下文效能。月之暗面實驗表明，Full Attention 和 MoBA 之間的無縫過渡可以成為高效長上下文預訓練的解決方案。然後他們還討論了分層混合策略，其作用主要是提升監督微調（SFT）的效能。在實驗中的混合訓練時，模型使用的是兩階段方案：第一階段使用 MoBA 完成 90% 的 token 訓練，第二階段模型切換到 Full Attention 使用剩餘 10% 的 token 進行訓練。另外，當然也使用了純 MoBA 和純 Full Attention 方法作為對比。結果見圖 5a，僅使用 MoBA 時會導致 trailing token 的位置損失更高。重要的是，混合方案的損失幾乎與 Full Attention 一樣。這一結果凸顯了混合訓練方案在平衡訓練效率和模型效能方面的有效性。更有趣的是，在 MoBA 和 Full Attention 之間切換時沒有觀察到明顯的損失峰值，這再次證明了 MoBA 的靈活性和穩健性。月之暗面也研究了分層混合的方案。這種方案當然更加複雜，研究動機是他們觀察到 MoBA 有時會導致 SFT 期間的效能不佳，如圖 5b 所示。他們推測這可能歸因於 SFT 中使用的損失掩蔽（loss masking）—— 提示詞 token 通常被排除在 SFT 期間的損失計算之外，這可能會對 MoBA 等稀疏注意力方法造成稀疏梯度難題。因為它可能會阻礙梯度的反向傳播，而這些梯度最初是根據未掩蔽的 token 計算出來的，並貫穿整個上下文。為了解決這個問題，他們提出了另一種混合方法 —— 將最後幾層 Transformer 從 MoBA 換成 Full Attention，而其餘層繼續採用 MoBA。如圖 5b 和圖 5c 所示，此策略可以顯著減少 SFT 損失。研究團隊基於 Llama 3.1 8B 模型開發了 Llama-8B1M-MoBA，透過持續預訓練將上下文長度從 128K 逐步擴充套件到 1M token。模型採用塊大小 4096 和 top-K 引數 12 的設定，實現了高達 95.31% 的注意力稀疏度。為平衡效能，他們保留最後三層為 Full Attention，其餘 29 層使用 MoBA。評估結果顯示，MoBA 模型與 Full Attention 模型（Llama-8B-1M-Full）效能相當。特別是在最具挑戰性的 RULER 基準測試中，即使在 62.5% 的稀疏度下，MoBA 模型仍達到了 0.7818 的得分，幾乎與 Full Attention 模型的 0.7849 持平。對於需要處理 100 萬 token 的「大海撈針」測試集中，模型同樣展現出優秀的能力。實驗證明，MoBA 在保持效能的同時顯著提升了效率：在處理 1M token 時，比 Full Attention 模型快 6.5 倍。在擴充套件到 1000 萬 token 時，較標準 Flash Attention 實現了 16 倍加速；透過塊稀疏注意力機制和最佳化實現，MoBA 將計算複雜度從二次方降低到了亞二次方級別。這使得 MoBA 的優勢在序列長度增加時更為明顯，特別適合處理超長文字場景。https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf 
