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

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

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

• 論文標(biāo)題：MoBA: Mixture of Block Attention for Long-Context LLMs
• 論文地址：https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf
• 項(xiàng)目地址：https://github.com/MoonshotAI/MoBA

與 DeepSeek 的 NSA 注意力機(jī)制新論文一樣，月之暗面這篇 MoBA 論文也收獲了諸多好評。

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

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

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

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

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

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

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

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

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

MoBA：將 MoE 原理應(yīng)用于注意力機(jī)制

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

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

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

MoBA 主要包括如下部分：

• 可訓(xùn)練的塊稀疏注意力：將完整上下文分割成若干塊，每個(gè)查詢 token 都會(huì)學(xué)習(xí)關(guān)注最相關(guān)的 KV 塊，從而實(shí)現(xiàn)長序列的高效處理。
• 無參數(shù)門控機(jī)制：引入了一種新穎的無參數(shù) top-k 門控機(jī)制，為每個(gè)查詢 token 選擇最相關(guān)的塊，確保模型只關(guān)注最具信息量的內(nèi)容。
• 完全注意力與稀疏注意力的無縫切換：MoBA 被設(shè)計(jì)為完全注意力機(jī)制的靈活替代方案，能夠在完全注意力和稀疏注意力模式之間實(shí)現(xiàn)平滑過渡。

總體而言，MoBA 注意力機(jī)制使模型能夠自適應(yīng)且動(dòng)態(tài)地關(guān)注上下文中最有信息量的塊。這在處理涉及長文檔或序列的任務(wù)時(shí)尤其有益，因?yàn)樵谶@些任務(wù)中，關(guān)注整個(gè)上下文可能是不必要的，并且計(jì)算代價(jià)高昂。MoBA 選擇性地關(guān)注相關(guān)塊的能力，使得信息處理更加細(xì)致且高效。

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

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

• 根據(jù)門控網(wǎng)絡(luò)和因果掩碼，確定查詢 token 對 KV 塊的分配。
• 根據(jù)分配的 KV 塊，安排查詢 token 的順序。
• 計(jì)算每個(gè) KV 塊的注意力輸出和分配給它的查詢 token，這一步可以通過不同長度的 FlashAttention 來優(yōu)化。
• 將注意力輸出重新排列回其原始順序。
• 使用在線 Softmax 合并相應(yīng)的注意力輸出，因?yàn)橐粋€(gè)查詢 token 可能會(huì)關(guān)注其當(dāng)前塊以及多個(gè)歷史 KV 塊。

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

MoBA 的實(shí)驗(yàn)表現(xiàn)

為了驗(yàn)證 MoBA 的實(shí)際效果，月之暗面進(jìn)行了一系列 Scaling Law 和消融實(shí)驗(yàn)。

Full Attention 與 MoBA 對比實(shí)驗(yàn)

該團(tuán)隊(duì)訓(xùn)練了 5 個(gè)模型來進(jìn)行比較，表 1 給出詳細(xì)配置。

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

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

該團(tuán)隊(duì)也通過消融實(shí)驗(yàn)對 MoBA 的細(xì)粒度塊分割的有效性進(jìn)行了驗(yàn)證，結(jié)果見圖 4。

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

將 MoBA 與 Full Attention 融合到一起會(huì)如何？

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

在實(shí)驗(yàn)中的混合訓(xùn)練時(shí)，模型使用的是兩階段方案：第一階段使用 MoBA 完成 90% 的 token 訓(xùn)練，第二階段模型切換到 Full Attention 使用剩余 10% 的 token 進(jìn)行訓(xùn)練。另外，當(dāng)然也使用了純 MoBA 和純 Full Attention 方法作為對比。

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

月之暗面也研究了分層混合的方案。這種方案當(dāng)然更加復(fù)雜，研究動(dòng)機(jī)是他們觀察到 MoBA 有時(shí)會(huì)導(dǎo)致 SFT 期間的性能不佳，如圖 5b 所示。

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

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

大語言模型評估

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

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

效率和可擴(kuò)展性

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

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

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