同一天，DeepSeek和月之暗面發(fā)表了關于改進注意力機制的論文-引自機器之心

昨天下午，DeepSeek 發(fā)布了一篇新論文，提出了一種改進版的注意力機制 NSA；加上還有創(chuàng)始人兼 CEO 梁文鋒親自參與，一時之間吸引眼球無數(shù)，參閱報道

但其實就在同一天，月之暗面也發(fā)布了一篇主題類似的論文，并且同樣巧合的是，月之暗面創(chuàng)始人兼 CEO 楊植麟也是該論文的署名作者之一。并且不同于 DeepSeek 只發(fā)布了一篇論文，月之暗面還發(fā)布了相關代碼。且這些代碼已有一年的實際部署驗證，有效性和穩(wěn)健性都有保證。

這篇論文提出了一種名為 MoBA 的注意力機制，即 Mixture of Block Attention，可以直譯為「塊注意力混合」。據(jù)介紹，MoBA 是「一種將混合專家（MoE）原理應用于注意力機制的創(chuàng)新方法。」該方法遵循「更少結(jié)構(gòu)」原則，并不會引入預定義的偏見，而是讓模型自主決定關注哪些位置。

img

• 
  

  論文標題：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 論文也收獲了諸多好評。

img

下面我們就來解讀一下這篇論文。

最近一段時間，大語言模型（LLM）發(fā)展到了前所未有的規(guī)模，我們希望它們能夠處理模擬人類認知的復雜任務。具體來說，實現(xiàn) AGI 的一個關鍵能力是處理、理解和生成長序列的能力，這對于從歷史數(shù)據(jù)分析到復雜推理和決策過程的各種應用至關重要。這種對擴展上下文處理能力的需求不僅體現(xiàn)在長輸入提示理解的普及上，還體現(xiàn)在最近 Kimi k1.5、DeepSeek-R1 和 OpenAI o1/o3 中展示的，對長鏈思維（CoT）輸出能力的探索中。

然而，由于傳統(tǒng)注意力機制帶來的計算復雜度呈二次函數(shù)增長，擴展 LLM 的序列長度并非易事。這一挑戰(zhàn)催生了一波研究浪潮，其中一個重要方向是利用注意力分數(shù)的固有稀疏性。這種稀疏性既源于數(shù)學上的 softmax 操作，也源于生物學上的觀察，即人們已在與記憶存儲相關的大腦區(qū)域中觀察到了稀疏連接。

現(xiàn)有方法通常利用預定義的結(jié)構(gòu)約束（例如基于 sink 的注意力機制或滑動窗口注意力機制）來利用這種稀疏性。雖然這些方法可能有效，但它們往往高度依賴于特定任務，可能會限制模型的整體泛化能力。另一種方法是采用動態(tài)稀疏注意力機制，例如 Quest、Minference 和 RetrievalAttention，這些方法在推理時選擇部分 token 進行計算。盡管這些方法可以減少長序列的計算量，但它們并未顯著緩解長上下文模型的高昂訓練成本。

最近，線性注意力模型作為一種有前景的替代方案出現(xiàn)，例如 Mamba、RWKV 和 RetNet。這些方法用線性近似替代了傳統(tǒng)的基于 softmax 的注意力機制，從而降低了長序列處理的計算開銷。然而，由于線性注意力與傳統(tǒng)注意力機制存在顯著差異，適配現(xiàn)有的 Transformer 模型通常需要高昂的轉(zhuǎn)換成本，或者需要從頭開始訓練全新的模型。更重要的是，它們在復雜推理任務中的有效性證據(jù)仍然有限。

因此，一個關鍵的研究問題隨之而來：我們?nèi)绾卧O計一種穩(wěn)健且適應性強的注意力架構(gòu)，既能保留原始 Transformer 框架，又能遵循「更少結(jié)構(gòu)」的原則，使模型能夠在不依賴預定義偏差的情況下自主決定關注哪些部分？理想情況下，這種架構(gòu)能夠在完全注意力和稀疏注意力模式之間無縫切換，從而最大限度地與現(xiàn)有預訓練模型兼容，并在不犧牲性能的前提下實現(xiàn)高效推理和加速訓練。

為此研究人員提出了塊注意力混合機制（Mixture of Block Attention, MoBA），這是一種全新架構(gòu)，它基于混合專家系統(tǒng)（MoE）的創(chuàng)新原理，并將其應用于 Transformer 模型的注意力機制中。MoE 此前主要應用于 Transformer 的前饋網(wǎng)絡（FFN）層，而 MoBA 首次將其引入長上下文注意力機制中，允許動態(tài)選擇與每個查詢 token 相關的歷史關鍵塊和值塊。這種方法不僅提高了 LLM 的效率，還使其能夠處理更長、更復雜的提示，而無需按比例增加資源消耗。

相比之下，MoBA 通過將上下文劃分為塊并采用門控機制選擇性地將查詢 token 路由到最相關的塊，解決了傳統(tǒng)注意力機制的計算效率低下的問題。這種塊稀疏注意力顯著降低了計算成本，為更高效地處理長序列鋪平了道路。模型能夠動態(tài)選擇最具信息量的關鍵塊，從而提高了性能和效率，這對于涉及大量上下文信息的任務尤為有益。

在該工作中，月之暗面詳細介紹了 MoBA 的架構(gòu)，首先是其塊劃分和路由策略，其次是與傳統(tǒng)注意力機制相比的計算效率。他們也進行了實驗，證明了 MoBA 在處理長序列任務中的卓越性能。

MoBA：將 MoE 原理應用于注意力機制

MoBA 通過動態(tài)選擇歷史片段（塊）進行注意力計算，擴展了 Transformer 模型的能力。

MoBA 的靈感來源于混合專家（Mixture of Experts，簡稱 MoE）技術和稀疏注意力技術。前者主要應用于 Transformer 架構(gòu)中的前饋網(wǎng)絡（Feedforward Network，簡稱 FFN）層，而后者已被廣泛用于擴展 Transformer 以處理長文本上下文。

本文創(chuàng)新點在于月之暗面將 MoE 原理應用于注意力機制本身，從而實現(xiàn)對長序列更高效和更有效的處理。

MoBA 主要包括如下部分：

• 可訓練的塊稀疏注意力：將完整上下文分割成若干塊，每個查詢 token 都會學習關注最相關的 KV 塊，從而實現(xiàn)長序列的高效處理。

• 無參數(shù)門控機制：引入了一種新穎的無參數(shù) top-k 門控機制，為每個查詢 token 選擇最相關的塊，確保模型只關注最具信息量的內(nèi)容。

• 完全注意力與稀疏注意力的無縫切換：MoBA 被設計為完全注意力機制的靈活替代方案，能夠在完全注意力和稀疏注意力模式之間實現(xiàn)平滑過渡。

總體而言，MoBA 注意力機制使模型能夠自適應且動態(tài)地關注上下文中最有信息量的塊。這在處理涉及長文檔或序列的任務時尤其有益，因為在這些任務中，關注整個上下文可能是不必要的，并且計算代價高昂。MoBA 選擇性地關注相關塊的能力，使得信息處理更加細致且高效。

月之暗面通過整合來自 FlashAttention 和 MoE 的優(yōu)化技術，實現(xiàn)了 MoBA 的高性能版本。圖 2 展示了 MoBA 的高效率。

img

MoBA 實現(xiàn)過程包含五個步驟：

• 根據(jù)門控網(wǎng)絡和因果掩碼，確定查詢 token 對 KV 塊的分配。

• 根據(jù)分配的 KV 塊，安排查詢 token 的順序。

• 計算每個 KV 塊的注意力輸出和分配給它的查詢 token，這一步可以通過不同長度的 FlashAttention 來優(yōu)化。

• 將注意力輸出重新排列回其原始順序。

• 使用在線 Softmax 合并相應的注意力輸出，因為一個查詢 token 可能會關注其當前塊以及多個歷史 KV 塊。

下圖為該算法的形式化描述，并在圖 1b 中進行了可視化，展示了如何基于 MoE 和 FlashAttention 實現(xiàn) MoBA。

img

MoBA 的實驗表現(xiàn)

為了驗證 MoBA 的實際效果，月之暗面進行了一系列 Scaling Law 和消融實驗。

Full Attention 與 MoBA 對比實驗

該團隊訓練了 5 個模型來進行比較，表 1 給出詳細配置。

img

結(jié)果見圖 3a，可以看到，F(xiàn)ull Attention 與 MoBA 的 scaling 趨勢非常相似。具體來說，這兩種注意力機制之間的驗證損失差異在 1e ? 3 的范圍內(nèi)保持一致。這表明，盡管 MoBA 的稀疏注意力模式稀疏度高達 75%，但它實現(xiàn)了與完全注意力相當?shù)臄U展性能。

img

此外，該團隊也驗證了 MoBA 的上下文可擴展性。在實驗中，該團隊將最大序列長度從 8k 提升到了 32k。結(jié)果見圖 3b，可見盡管在所有五項實驗中，MoBA 最后的塊 LM 損失與 Full Attention 相比略高，但損失差距正在逐步縮小。該實驗表明 MoBA 具有長上下文可擴展性。

該團隊也通過消融實驗對 MoBA 的細粒度塊分割的有效性進行了驗證，結(jié)果見圖 4。

img

可以看到，MoBA 的性能受塊粒度的顯著影響。具體來說，最粗粒度設置（從 8 個塊中選擇 2 個塊）與更細粒度的設置之間的性能差異為 1e-2。這些發(fā)現(xiàn)表明，細粒度分割似乎是增強 MoE 系列（包括 MoBA）模型性能的普適技術。

將 MoBA 與 Full Attention 融合到一起會如何？

MoBA 本就是作為 Full Attention 的靈活替代而設計的，可以最小的開銷輕松地互相切換，并實現(xiàn)相當?shù)拈L上下文性能。月之暗面實驗表明，F(xiàn)ull Attention 和 MoBA 之間的無縫過渡可以成為高效長上下文預訓練的解決方案。然后他們還討論了分層混合策略，其作用主要是提升監(jiān)督微調(diào)（SFT）的性能。

在實驗中的混合訓練時，模型使用的是兩階段方案：第一階段使用 MoBA 完成 90% 的 token 訓練，第二階段模型切換到 Full Attention 使用剩余 10% 的 token 進行訓練。另外，當然也使用了純 MoBA 和純 Full Attention 方法作為對比。

結(jié)果見圖 5a，僅使用 MoBA 時會導致 trailing token 的位置損失更高。重要的是，混合方案的損失幾乎與 Full Attention 一樣。這一結(jié)果凸顯了混合訓練方案在平衡訓練效率和模型性能方面的有效性。更有趣的是，在 MoBA 和 Full Attention 之間切換時沒有觀察到明顯的損失峰值，這再次證明了 MoBA 的靈活性和穩(wěn)健性。

img

月之暗面也研究了分層混合的方案。這種方案當然更加復雜，研究動機是他們觀察到 MoBA 有時會導致 SFT 期間的性能不佳，如圖 5b 所示。

他們推測這可能歸因于 SFT 中使用的損失掩蔽（loss masking）—— 提示詞 token 通常被排除在 SFT 期間的損失計算之外，這可能會對 MoBA 等稀疏注意力方法造成稀疏梯度難題。因為它可能會阻礙梯度的反向傳播，而這些梯度最初是根據(jù)未掩蔽的 token 計算出來的，并貫穿整個上下文。

為了解決這個問題，他們提出了另一種混合方法 —— 將最后幾層 Transformer 從 MoBA 換成 Full Attention，而其余層繼續(xù)采用 MoBA。如圖 5b 和圖 5c 所示，此策略可以顯著減少 SFT 損失。

大語言模型評估

研究團隊基于 Llama 3.1 8B 模型開發(fā)了 Llama-8B1M-MoBA，通過持續(xù)預訓練將上下文長度從 128K 逐步擴展到 1M token。模型采用塊大小 4096 和 top-K 參數(shù) 12 的設置，實現(xiàn)了高達 95.31% 的注意力稀疏度。為平衡性能，他們保留最后三層為 Full Attention，其余 29 層使用 MoBA。

img

評估結(jié)果顯示，MoBA 模型與 Full Attention 模型（Llama-8B-1M-Full）性能相當。特別是在最具挑戰(zhàn)性的 RULER 基準測試中，即使在 62.5% 的稀疏度下，MoBA 模型仍達到了 0.7818 的得分，幾乎與 Full Attention 模型的 0.7849 持平。對于需要處理 100 萬 token 的「大海撈針」測試集中，模型同樣展現(xiàn)出優(yōu)秀的能力。

img

效率和可擴展性

實驗證明，MoBA 在保持性能的同時顯著提升了效率：在處理 1M token 時，比 Full Attention 模型快 6.5 倍。在擴展到 1000 萬 token 時，較標準 Flash Attention 實現(xiàn)了 16 倍加速；通過塊稀疏注意力機制和優(yōu)化實現(xiàn)，MoBA 將計算復雜度從二次方降低到了亞二次方級別。

這使得 MoBA 的優(yōu)勢在序列長度增加時更為明顯，特別適合處理超長文本場景。

更多詳情，請參見原論文。

https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf

哪吒歸來

“是魔是仙 我自己說了算”