2 月 18 日，Kimi 和 DeepSeek 同一天釋出新進展，分別是 MoBA 和 NSA，二者都是對 “注意力機制”（Attention Mechanism）的改進。

今天，MoBA 的一位主要研發同學 Andrew Lu 在知乎發帖，自述研發過程的三次踩坑，他稱為 “三入思過崖”。他在知乎的簽名是“新晉 LLM 訓練師”。

這條回答下的一個評論是：“從開源論文、開原始碼出發，現在已經進化到開源思維鏈了嘛。”

注意力機制之所以重要，是因為它是當前大語言模型（LLM）的核心機制。回到 2017 年 6 月那篇開啟 LLM 革命的 Transformer 八子論文，標題就是：Attention Is All You Need（注意力就是你所需要的一切），該論文被引用次數至今已達 15.3 萬。

注意力機制能讓 AI 模型像人類一樣，知道在處理資訊時該 “重點關注” 什麼、“忽略” 什麼，抓住資訊中最關鍵的部分。

在大模型的訓練階段和使用（推理）階段，注意力機制都會發揮作用。它的大致工作原理是，當輸入一段資料，如 “我喜歡吃蘋果”，大模型會計算句子中每一個詞（Token）與其他詞的關係，從而理解語義等資訊。

而當大模型需要處理的上下文越來越長，標準 Transformer 最初採用的 Full Attention（全注意力機制）對計算資源的佔用變得不可忍受，因為原初流程是，需要全部計算所有輸入詞的重要性分數，再加權計算後得到哪些是最重要的詞，它的計算複雜度會隨文字變長而平方級（非線性）增長。如 MoBA 論文的 “摘要” 部分所寫：

同時，研究者們又追求大模型處理的上下文能足夠長——多輪對話、複雜推理、記憶能力……這些設想中 AGI 應該具備的特性都需要更長長長長的上下文能力。

如何找到一個既沒那麼佔用計算資源和記憶體，又不損失模型效能的注意力機制最佳化方法，於是成為大模型研究的重要課題。

這是數家公司將注意力交匯到 “注意力” 上的技術背景。

在 DeepSeek NSA 和 Kimi MoBA 之外，今年 1 月中旬，另一家中國大模型創業公司 MiniMax 也在其首個開源模型 MiniMax-01 中大規模實現了一種新的注意力機制。MiniMax 創始人閆俊傑當時告訴我們，這是 MiniMax-01 最主要的創新點之一。

面壁智慧聯合創始人、清華大學計算機系副教授劉知遠的團隊也在 2024 年發表過 InfLLM，其中也涉及一種稀疏注意力改進，該論文被 NSA 的論文引用。

這幾個成果中，NSA、MoBA、InfLLm 裡的注意力機制都屬於 “稀疏注意力機制”（Sparse Attention)；而 MiniMax-01 的嘗試則主要是另一個方向：“線性注意力機制”（Linear Attention)。

SeerAttention 作者之一，微軟亞研院高階研究員曹士傑告訴我們：總體來說，線性注意力機制對標準注意力機制的改動更多、更激進，想直接解決隨文字變長，計算度平方爆炸（所以是非線性的）的問題，可能的一個代價是，會損失對長上下文的複雜依賴關係的捕捉；稀疏注意力機制則是利用注意力固有的稀疏性，嘗試尋找一種較為穩健的最佳化方式。

同時在此推薦曹士傑老師在知乎上關於注意力機制的高贊回答：https://www.zhihu.com/people/cao-shi-jie-67/answers

（他回答了 “梁文鋒參與發表的 DeepSeek 新論文 NSA 注意力機制，有哪些資訊值得關注？會帶來哪些影響？” 這個問題。）

MoA（Mixture of Sparse Attention） 的共同一作，清華大學 NICS-EFC 實驗室博士傅天予說，在稀疏注意力機制的大方向下：“NSA 和 MoBA  都引入了動態注意力方法，即能動態選擇需要計算細粒度注意力的 KV Cache 塊，相比一些使用靜態方法的稀疏注意力機制，能提升模型效能。這兩個方法也都是在模型訓練中就引入了稀疏注意力，而非僅在推理時引入，這也進一步提升了模型效能。”

（注：KV Cache 塊是儲存之前計算的 Key 標籤和 Value 值的快取；其中 Key 標籤指注意力機制相關計算中，用於標識資料特徵或資料位置等資訊的標識標籤，以便在計算注意力權重時，能與其他資料匹配和關聯 ，Value 值則與 Key 標籤對應，通常包含了實際要處理的資料內容，如單詞或短語的語義向量等資訊。）

同時，這次月之暗面除了釋出詳細的 MoBA 技術論文，也在 GitHub 該專案網站上釋出了 MoBA 工程程式碼，該套程式碼已在月之暗面自己的產品 Kimi 中線上使用了一年多。

* 以下為 Andrew Lu 在知乎上的自述，已獲得作者授權。原文中有多處 AI 術語，（）中灰色文字解釋部分均為編者注。原貼連結：https://www.zhihu.com/people/deer-andrew

應章老師（清華大學助理教授章明星）邀請，來答一下之前做 MoBA 起起伏伏的心路歷程，我戲稱為 “三入思過崖”。（Andrew Lu 回答的問題是：“如何評價 Kimi 開源的稀疏注意力框架 MoBA？與 DeepSeek 的 NSA 相比，二者各有哪些亮點？”）

MoBA 這個專案開始得非常早，在 2023 年 5 月底月之暗面剛建立不久的時候，入職報到當天就被 Tim （月之暗面聯合創始人周昕宇）拉到了小房間裡，和裘老師（浙江大學 / 之江實驗室裘捷中，MoBA idea 的提出者）以及 Dylan（月之暗面研究員）一起開始搞 Long Context Training（長上下文訓練）。這裡首先要感謝一下 Tim 的耐心和教導，對一個 LLM 新手給予厚望並樂意培養，研發各種上線模型和模型相關技術的諸位大佬裡面，很多人和我一樣基本是從零開始接觸 LLM。

當時業界普遍水平也不是很高，大家都在 4K 預訓練（模型能處理的輸入輸出長度約在 4000 個 Token，幾千個漢字），專案一開始叫 16K on 16B，意思是在 16B （模型引數 160 億）上做 16K 長度的 Pre-train（預訓練） 即可，當然後來很快的這個需求在 8 月變成了需要支援 128K 下 Pre-train 。這也是 MoBA 設計時的第一個要求，能 From Scratch（從頭開始）快速的訓練出一個能支援 128K 長度下的模型，此時還不需要 Continue Training（持續訓練，在已訓模型基礎上繼續訓練）。

這裡也引申出一個有趣的問題，23 年 5/6 月的時候，業界普遍認為訓長得長，端到端訓練長文字（直接用長文字訓練模型）的效果好於訓練一個較短的模型再想辦法搞長它。這種認知到 23 年下半年 long Llama （Meta 開發的、支援長文字處理的大模型）出現的時候才發生了轉變。我們自己也進行了嚴格的驗證，實際上短文字訓練 + 長度啟用具有更好的 token efficiency（每個 token 貢獻的有效資訊量提升，意味著模型能用更少的 token 完成更高質量的任務）。自然 MoBA 設計中的第一個功能就成了時代的眼淚。

在這個時期，MoBA 的結構設計也是更為 “激進” 的，相較於現在的 “極簡化” 結果，初提出的 MoBA 是一個帶 cross attention （處理兩段不同文字資料之間關係的注意力機制）的兩層注意力機制序列方案，gate （控制輸入資料如何在各個專家網路之間分配權重）本身是一個無參的結構（沒有引數，不需要資料訓練），但是為了更好的學習歷史的 token，我們在每個 Transformer 層加了一個機器間的 cross attention 以及對應的引數（可以更好地記住歷史資訊）。此時的 MoBA 設計已經結合了後面為大家熟知的 Context Parallel 思想（完整的上下文序列被存放到不同節點上，在需要計算的時候才集中在一起），我們將整個上下文序列平鋪在資料並行節點之間，將每個資料並行節點內的 context 看成是一個 MoE （Mixture of Experts，專家混合系統）中的 expert（專家），將需要 attention 的 token 傳送到對應 expert 上進行 cross attention 再把結果通訊回來。我們將 fastmoe （一種早期 MoE 訓練的框架）的工作整合進了 Megatron-LM（來自 Nvidia 的現在通用大模型訓練框架）來支援 expert 間通訊能力。

這個思路我們稱之為 MoBA v0.5。

隨著時間推進到 23 年 8 月初，主模型 Pre-Train 已經訓練了大量 token，再來一次成本不低。顯著改變了結構並增加了額外引數的 MoBA 至此第一次進入思過崖。

進入思過崖當然是一種戲稱，是停下來尋找改進方案的時間，也是深入理解新結構的時間。第一次進思過崖悟道，進去的快，出來的也快。Tim 作為月之暗麵點子王掏出了新的改進思路，將 MoBA 從序列兩層注意力方案改並行的單層注意力方案。MoBA 不再增加額外的模型引數，而是利用現有注意力機制引數，同步學習一個序列裡面的所有資訊，這樣就可以儘可能不變動當前結構進行 Continue Training。

這個思路我們稱之為 MoBA v1。

MoBA v1 實際上是 Sparse Attention（稀疏注意力）Context Parallel 的產物，在當時 Context Parallel 並不是大行其道的時候，MoBA v1 體現了極高的端到端加速能力。我們在 3B，7B 上都驗證了它有效之後，在更大模型 scale 水平上撞牆了，訓練過程中出現了非常大的 loss spike（模型訓練時出現的異常現象）。我們初版合併 block attention output（注意力模組處理完資料後輸出結果）的方式過於粗淺，只是簡單累加，導致完全無法和 Full Attention 對分進行 debug，沒有 ground truth（標準答案，這裡指 Full Attention 的結果）的 debug 是極其困難的，我們用盡了各種當時的穩定性手段都不能解決。由於在較大的模型上訓練出了問題，MoBA 至此二入思過崖。

第二次留在思過崖時間比較長，從 23 年 9 月開始，至出思過崖的時候已經到了 24 年初。但是在思過崖裡並不意味著被放棄了，我得以體會到在月之暗面工作的第二大特色，飽和救援。

除了一直就在強力輸出的 Tim 和裘老師，蘇神（蘇劍林，月之暗面研究員）、遠哥（Jingyuan Liu，月之暗面研究員）以及各路大佬都參與進行激烈的討論，開始拆解並修正 MoBA，首先被修正的就是那個簡單的 Weighted Sum（加權和）疊加。我們這裡嘗試過各種和 Gate Matrix 相乘相加的疊加搞法之後，Tim 從故紙堆裡掏出了 Online Softmax（不是看到所有資料才能計算，而是來一個數據處理一個），說這個應該能行。其中最大的一個好處是使用 Online Softmax 之後我們可以透過將稀疏度降低至 0（全選所有的分塊），來和一個數學等價的 Full Attention 進行嚴格對照 debug，這解決了大部分實現中遇到的難題。但是上下文拆分在資料並行節點之間的這個設計依然會導致不均衡問題，一個數據 sample 在資料並行間平鋪之後，第一個資料並行 rank 上的頭部幾個 token 會被後續茫茫多的 Q 傳送進行 attend（注意力計算過程），帶來極差的平衡性，進而拖慢加速效率。這個現象也有一個更廣為人知的名字——Attention Sink（注意力匯聚點）。

此時章老師到訪，聽完我們的想法後提出了新的思路，將 Context Parallel 能力和 MoBA 切分開。Context Parallel 是 Context Parallel，MoBA 是 MoBA ，MoBA 迴歸一個 Sparse Attention 本身而不是分散式 Sparse Attention 訓練框架。只要視訊記憶體放得下，完全可以在單機上處理全部上下文，用 MoBA 來進行計算加速，透過 Context Parallel 的方式來組織和傳遞機器間的上下文。因此我們重新實現了 MoBA v2，基本已經是也是當前大家見到的 MoBA 的樣子。

MoBA v2 穩定可訓，短文字和 Full Attention 可以完全對齊，Scaling Law 看起來非常可靠，而且比較絲滑的支援推廣到線上模型上。我們因此加了更多的資源進來，在經過了一系列的 debug 並消耗了 n 把 infra 組同學頭髮後，我們可以將經過 MoBA 啟用後的 Pretrain 模型做到撈針測試全綠（大模型處理長文字能力測試中達標），在這一步我們已經覺得非常不錯了，開始進行上線。

但是最不意外的，只有意外。SFT （監督微調，在預訓練模型的基礎上，針對特定任務對模型進行進一步訓練，以提高模型在該任務上的效能）階段部分資料帶著非常稀疏的 loss mask（使得只有 1% 甚至更少的 token 有訓練用的梯度）（loss mask 指挑選哪些部分參與衡量模型預測結果與標準答案計算的技術），這導致 MoBA 在大部分 SFT 任務上表現良好，但是越是長文總結型別任務，loss mask 越是稀疏，反應在出來的學習效率越是低下。MoBA 在準上線流程中被按下暫停鍵，第三次進入思過崖。

第三次進入思過崖的時候其實最緊張，此時整個專案已經有了巨大的沉沒成本，公司付出了大量的計算資源和人力資源，如果端到端最後長文應用場景出現問題，那麼前期研究接近於打水漂。幸運的是，由於 MoBA 本身優秀的數學性質，在新一輪飽和救援的實驗 ablation（消融實驗，透過去除模型的某些部分或改變某些設定來研究其對模型效能的影響）中，我們發現去掉 loss mask 表現就非常良好，帶上 loss mask 表現就不盡如人意，進而意識到是帶有 gradient （梯度，在機器學習中用於更新模型引數的方向和步長的一個值）的 token 在 SFT 階段過於稀疏，從而帶來的學習效率低下的問題，因此透過將最後幾層修改為 Full Attention，來提高反向傳播時帶 gradient token 的密度，改善特定任務的學習效率。後續其它實驗證明，這種切換並不會顯著影響切換回來的 Sparse Attention 效果，在 1M （100 萬）長度上和同結構 Full attention 各項指標持平。MoBA 再次從思過崖迴歸，併成功上線服務使用者。

最後的最後，感謝各路大神的拔刀相助，感謝公司的大力支援以及巨量顯示卡。現在我們開放的出來的就是我們在線上用的程式碼，是一個已經經過了長期驗證，因為實際需求砍掉了各種額外設計，保持極簡結構但同時具備足夠效果的 Sparse Attention 結構。希望 MoBA 以及它誕生的 CoT （Chain of Thought，思維鏈）能給大家帶來一些幫助和價值。

順帶借地方回答一些這兩天頻繁被人問到的問題，這兩天基本麻煩章老師和蘇神當客服回答問題了，實在過意不去，這裡提取了幾個常見問題一併回答一下。

MoBA 對 Decoding 是有效的，對 MHA（Multi-Head Attention，多頭注意力）很有效，對 GQA （Grouped Query Attention，分組查詢注意力）效果降低，對 MQA  （Multi-Query Attention，多查詢注意力）效果最差。原理其實很簡單，MHA 的情況下，每個 Q 有一個自己對應的 KV cache，那麼 MoBA 的 gate 在理想情況下是可以透過攤餘計算在 prefill （首次處理輸入時的計算階段）算好並存儲每個 block（資料塊）的代表 token，這個 token 在後續都不會變動，因此所有的 IO （輸入輸出操作）基本可以做到只來自 index select（透過索引選擇資料的操作）之後的 KV cache，這種情況下 MoBA 的稀疏程度就決定了 IO 減少的程度。

但是對於 GQA 和 MQA，由於一組 Q Head 實際上在共享同樣的 一份 KV cache，那麼在每個 Q Head 能自由選擇感興趣 Block 的情況下，很有可能填滿了由稀疏性帶來的 IO 最佳化。比如我們思考這麼一個場景：16 個 Q Head 的 MQA，MoBA 剛好切分整個序列到 16 份，這意味著當最壞情況每個 Q head 感興趣的分別是 序號從 1 到 16 的每個上下文塊，節省 IO 的優勢就會被磨平。能自由選擇 KV Block 的 Q Head 越多，效果越差。

由於 “自由選擇 KV Block 的 Q Head” 這麼個現象存在，天然的改進想法就是合併，假設大家都選一樣的 Block，那不就淨賺 IO 優化了麼。對，但是在我們實際測試下，尤其是對於已經支付了大量成本的預訓練模型，每個 Q head 都有自己獨特的 “品味”，強行合併不如從頭重新訓練。

不會必選，這個是已知會產生些許疑惑的地方，我們最後選擇相信 SGD（Stochastic Gradient Descent，隨機梯度下降）。現在的 MoBA gate 實現非常直接，有感興趣的同學可以簡單改造 gate 使其必選上一個 chunk（資料塊），但我們自己測試這個改動帶來的收益比較 margin（收益不大）。

我們實現過一個版本，端到端效能提升 10%+，但是 Triton 的實現想要持續維護跟上主線的成本比較高昂，我們因此多次迭代後暫緩進一步最佳化它。

MoBA GitHub 頁面：https://github.com/MoonshotAI/MoBA

NSA 技術論文：https://arxiv.org/abs/2502.11089

MiniMax-01 GitHub 頁面：https://github.com/MiniMax-AI/MiniMax-01

InfLLM GitHub 頁面：https://github.com/thunlp/InfLLM?tab=readme-ov-file

SeerAttention GitHub 頁面：https://github.com/microsoft/SeerAttention