近年來,開源技術不斷演進,注意力機制的優化成為研究的重點。 2023年2月,Kimi和DeepSeek分別推出了MoBA和NSA這兩項重要進展,均著眼於改進注意力機制,提升大語言模型的表現。 MoBA的核心是稀疏注意力機制,旨在解決標準全注意力機制在長文本處理中的計算複雜度問題,強調在訓練和推理階段都能有效處理上下文信息。
研發過程中的挑戰不少,項目負責人Andrew Lu形像地稱其經歷為“三入思過崖”,指的是在不斷的迭代與改進中解決出現的問題。 MoBA項目經過多次設計優化,從串行的雙層結構調整為並行的單層結構,並利用在線softmax技術提高模型的穩定性,最終實現了良好的訓練效果和性能。儘管在特定應用場景中仍面臨挑戰,如長文本總結任務中的稀疏學習效率問題,團隊通過改進方法使MoBA成功上線。
總體來看,MoBA的開發不僅展示了團隊的技術實力和創新能力,也彰顯了開源文化在推動AI領域突破和發展的重要性。
“從開源論文、開源代碼出發,現在已經進化到開源思維鏈了嘛!”
作者:Andrew Lu,晚點團隊
圖片來源:由無界AI生成
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
Andrew Lu 的研發自述
應章老師(清華大學助理教授章明星)邀請,來答一下之前做MoBA 起起伏伏的心路歷程,我戲稱為“三入思過崖”。 (Andrew Lu 回答的問題是:“如何評價Kimi 開源的稀疏注意力框架MoBA?與DeepSeek 的NSA 相比,二者各有哪些亮點?”)
MoBA 的開始
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。
(編者註:MoBA 的靈感來自當前主流的大模型MoE 結構。MoE 指大模型工作時,每次只激活部分專家的參數,而不是全部,從而節省算力;MoBA 的核心思路是“每次只看最相關上下文,而不是全部上下文,從而節省計算和成本”。)
隨著時間推進到23 年8 月初,主模型Pre-Train 已經訓練了大量token,再來一次成本不低。顯著改變了結構並增加了額外參數的MoBA 至此第一次進入思過崖。
一個非常簡單的MoBA v0.5 的示意圖
編者註:
History Tokens(歷史標記)——在自然語言處理等場景中,代表之前處理過的文本單元集合。
Gate(門)——在神經網絡中,用於控制信息流動的結構
Input(輸入)——模型接收的數據或信息
V(Value)——在註意力機制中,包含實際要處理或關注的數據內容,比如語義向量等
K(Key 標籤)——在註意力機制中,用於標識數據特徵或位置等信息的標識標籤,以便與其他數據匹配關聯
Q(Querry 查詢)——在註意力機制中,用於從鍵– 值對中檢索相關信息的向量
Cross Attention(交叉注意力)——一種注意力機制,關注不同來源的輸入,如將輸入與歷史信息關聯
Self Attention(自註意力)——一種注意力機制,模型對自身輸入進行關注,捕捉輸入內部的依賴關係
一入思過崖
進入思過崖當然是一種戲稱,是停下來尋找改進方案的時間,也是深入理解新結構的時間。第一次進思過崖悟道,進去的快,出來的也快。 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 至此二入思過崖。
一個非常簡單的MoBA v1 的示意圖
編者註:
Self Attention to History(對歷史的自註意力)——一種注意力機制,模型關注歷史標記,捕捉當前輸入與歷史信息間的依賴關係
Share weights(共享權重)——神經網絡中不同部分使用相同的權重參數,以減少參數數量和提高模型泛化能力
FFN(Feed – Forward Neural Network,前饋神經網絡)——一種基本的神經網絡結構,數據沿單一方向從輸入層經過隱藏層流向輸出層
Weighted Sum(加權和)——將多個值按照各自的權重進行求和的操作
二入思過崖
第二次留在思過崖時間比較長,從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 設計
編者註:
MoBA Gating(MoBA 門控)——MoBA 中的特定門控機制
RoPE(Rotary Position Embedding 旋轉位置嵌入)——一種為序列添加位置信息的技術
Partition to blocks(劃分為塊)——將數據分割成不同的塊
Mean Pooling(平均池化)——一種在深度學習中對數據進行降採樣的操作,計算區域內數據的平均值
MatMul(Matrix – Multiply 矩陣乘法)——一種數學運算,用於計算兩個矩陣的乘積
TopK Gating(Top – K 門控)——一種門控機制,選擇前K 個重要的元素等操作
Selected Block Index(選定的塊索引)——表示被選中的塊的編號
Index Select(索引選擇)——根據索引從數據中選取相應元素
Varlen Flash-Attention(可變長度快速注意力)——一種注意力機制,適用於可變長度序列且計算效率較高
Attention Output(注意力輸出)——注意力機制計算後的輸出結果
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,思維鏈)能給大家帶來一些幫助和價值。
FAQ
順帶借地方回答一些這兩天頻繁被人問到的問題,這兩天基本麻煩章老師和蘇神當客服回答問題了,實在過意不去,這裡提取了幾個常見問題一併回答一下。
1.MoBA 對Decoding (模型推理階段的文本生成過程)無效嗎?
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 都有自己獨特的“品味”,強行合併不如從頭重新訓練。
2.MoBA 默認是必選self attention(自註意力機制),那麼self 的鄰居會必選嗎?
不會必選,這個是已知會產生些許疑惑的地方,我們最後選擇相信SGD(Stochastic Gradient Descent,隨機梯度下降)。現在的MoBA gate 實現非常直接,有感興趣的同學可以簡單改造gate 使其必選上一個chunk(數據塊),但我們自己測試這個改動帶來的收益比較margin(收益不大)。
3.MoBA 有Triton (編寫高性能GPU 代碼的框架,由OpenAI 開發)實現麼?
我們實現過一個版本,端到端性能提升10%+,但是Triton 的實現想要持續維護跟上主線的成本比較高昂,我們因此多次迭代後暫緩進一步優化它。
*文章開篇提及的幾個成果的項目地址(GitHub 頁面中均包含技術paper 鏈接,DeepSeek 暫未上線NSA 的GitHub 頁面):
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