撞車DeepSeek NSA，Kimi楊植麟署名的新注意力架構MoBA發布，代碼也公開
昨天下午，DeepSeek 發布了一篇新論文，提出了一種改進版的注意力機制 NSA；加上還有創始人兼 CEO 梁文鋒親自參與，一時之間吸引眼球無數。
但其實就在同一天，月之暗面也發布了一篇主題類似的論文，并且同樣巧合的是，月之暗面創始人兼 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：將 MoE 原理應用于注意力機制
MoBA 通過動態選擇歷史片段（塊）進行注意力計算，擴展了 Transformer 模型的能力。
MoBA 的靈感來源于混合專家（Mixture of Experts，簡稱 MoE）技術和稀疏注意力技術。前者主要應用于 Transformer 架構中的前饋網絡（Feedforward Network，簡稱 FFN）層，而后者已被廣泛用于擴展 Transformer 以處理長文本上下文。
本文創新點在于月之暗面將 MoE 原理應用于注意力機制本身，從而實現對長序列更高效和更有效的處理。
MoBA 主要包括如下部分：
- 可訓練的塊稀疏注意力：將完整上下文分割成若干塊，每個查詢 token 都會學習關注最相關的 KV 塊，從而實現長序列的高效處理。
- 無參數門控機制：引入了一種新穎的無參數 top-k 門控機制，為每個查詢 token 選擇最相關的塊，確保模型只關注最具信息量的內容。
- 完全注意力與稀疏注意力的無縫切換：MoBA 被設計為完全注意力機制的靈活替代方案，能夠在完全注意力和稀疏注意力模式之間實現平滑過渡。
總體而言，MoBA 注意力機制使模型能夠自適應且動態地關注上下文中最有信息量的塊。這在處理涉及長文檔或序列的任務時尤其有益，因為在這些任務中，關注整個上下文可能是不必要的，并且計算代價高昂。MoBA 選擇性地關注相關塊的能力，使得信息處理更加細致且高效。
月之暗面通過整合來自 FlashAttention 和 MoE 的優化技術，實現了 MoBA 的高性能版本。圖 2 展示了 MoBA 的高效率。
MoBA 實現過程包含五個步驟：
- 根據門控網絡和因果掩碼，確定查詢 token 對 KV 塊的分配。
- 根據分配的 KV 塊，安排查詢 token 的順序。
- 計算每個 KV 塊的注意力輸出和分配給它的查詢 token，這一步可以通過不同長度的 FlashAttention 來優化。
- 將注意力輸出重新排列回其原始順序。
- 使用在線 Softmax 合并相應的注意力輸出，因為一個查詢 token 可能會關注其當前塊以及多個歷史 KV 塊。
下圖為該算法的形式化描述，并在圖 1b 中進行了可視化，展示了如何基于 MoE 和 FlashAttention 實現 MoBA。
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 的優勢在序列長度增加時更為明顯，特別適合處理超長文本場景。
更多詳情，請參見原論文。
