為什麼 MiniMax M2 是一個 Full Attention 模型? - 新聞 詳情

來源:為什麼MiniMax M2是一個Full Attention模型?

作為MiniMax M2預訓練的負責人,我收到了很多來自社區的詢問:“為什麼你們在MiniMax M2上開倒車,採用了 Full Attention 機制?” 在一次又一次的聊天中解釋了背後的故事後,我覺得是時候在一篇blog裏寫下我們的心路歷程。

我可以花一整個下午來討論為什麼應該構建應該做 Linear/Sparse Attention。同樣,我也可以反過來花一整個下午來討論為什麼不應該去做。但所有這些紙上談兵又有什麼意義呢?回到實際情況裏,我們要不要做呢?

先説結論:我們一直在研究它。但在一個現實的工業系統中,Efficient Attention想要打敗Full Attention還有些距離。LLM發展到今天,整個鏈路越來越複雜,服務的場景越來越多,結構設計上需要考慮的點也爆炸式增長;Code/Math場景效果咋樣、Agent場景下效果怎麼樣、多模態是什麼表現、Long CoT行不行、RL能不能Scale上去、低精度運算下有沒有坑、Interleaved Thinking怎麼弄、線上怎麼做cache、怎麼做Speculative Decoding等等。

簡而言之,現實與理論存在着較大差異,要為 Linear/Sparse Attention 正名,就要在滿足了條件1到條件n,並解決了問題1到問題n之後。

為什麼要做Efficient Attention?

如果你有無限算力,你會選擇研究Linear Attention或者Sparse Attention嗎?也許有人會説infinite context場景下,Softmax Attention的Attention Score會oversmoothing....但誰知道呢?在目前算力bound的前提下,還沒有哪個模型真的把Softmax Attention打到能力上限。因此,從實際應用的角度來看,目前大家做Linear Attention或者Sparse Attention都是奔着省算力去的。

有沒有可能奔着省token去呢——達到同樣效果,需要的token量更少。如果你相信Scaling Law,想達到這個目的,應該不會選擇走Efficient Attention這條路,而是其他途徑。

説穿了就是,算力有限,我們需要有一個更省算力的結構,同算力消耗下(training compute and inference compute)有更好的效果。

需要解決的問題

我們希望做一個能被商用的模型,我們必須從用户所關心的問題出發:效果、速度(TPS)和價格。效果是底線,一個效果差的模型,即便免費也毫無價值。那麼,怎麼做一個效果足夠好的Linear/Sparse/Hybrid Attention模型呢?這裏最大的問題不是結構設計,而是評測的侷限性。速度和價格受推理系統影響,當然優秀的模型自然有優秀的工程師來幫忙優化。

1.觀測侷限性

“只要你把bench建出來,我自然能找到辦法打上去。”

縱觀大模型發展的幾年,榜單分數提升的速度是驚人的,不管多難的榜,就算榜單剛出來的時候SOTA只有個位數分數,只要入了各家的法眼,總能在幾版迭代後刷爆。

怎麼建一個全面、真能反應模型能力差異的評測鏈路?這是一個很難的課題,也是大模型迭代的重中之重。這個問題在模型結構,特別是Attention迭代中,將變得更加嚴峻。

  • Benchmark不夠全面

“沒有免費的午餐”,把attention的複雜度降下去,付出的代價是什麼?

在做MiniMax-Text-01的時候,大家還普遍在看MMLU/BBH/Math/LongBench這類的榜單(現在已經被刷爆了)。以一年多前的視角來看,Ligntning Attention + Full Attention完全能打全Full Attention,畢竟這些榜單上都不差(我們端到端訓了個Hybrid架構的小模型來驗證)。

難道真有免費午餐?其實不然。這個代價在更大的模型上暴露出來了:複雜多跳推理任務有明顯缺陷。

當問題暴露出來就好辦了,對代理指標迭代優化!經過幾輪迭代,代理指標看起來能打MHA了。但是scale上去後,代理指標和下游實際場景還能match上嗎?還有其他的問題嗎?誰知道呢,還沒實驗到這裏。

模型越進步,評測越難做。但這是必經之路,加油評測人!

  • 觀測成本高

針對複雜多跳推理任務,我們能找到訓練更早期可以被觀測的代理指標,但並不是所有任務在預訓練階段都可以被觀測(起碼現在還不行)。同時隨着任務越來越難,想要到對應指標觀測的置信區間,需要付出的算力代價也越來越大,這也導致了實驗迭代比較緩慢(算力不夠才來研究這玩意,研究這玩意吃算力也不少)。

除了評測榜單,還有模型優化問題,不scale上去,永遠不知道會發生什麼,很多問題在小規模試驗中無法暴露。看過M1論文的朋友應該會發現M1 RL訓練過程中有嚴重的精度問題,不做到這一步確實很難發現這個雷。基於這個現象再回過頭去對Lightning Attention做數值收斂性分析,要怎麼解決真的很通透了。發現問題,真的比解決問題要難得的多的多。

  • 其他變量

訓練模型的變量太多太多,不同結構在不同數據分佈下的表現大相徑庭,不同結構適配的優化器也差異巨大。在數據高速迭代的週期裏,用一個月前的數據做實驗可能會得完全相反的結論。

我們很難做完備的觀測,但在試圖找更靠譜的實驗策略。

2.Efficient Attention的基建

相比Full Attention,Linear Attention和Sparse Attention的基建要差的多,想要真的拿到收益,要補不少課。

以Linear Attention的基建問題為例:如果對現在已有的線性結構做計算強度分析,會發現不少結構都是訪存bound(沒錯,訓練的時候)。如果不能做極致的IO優化,從GPU算力利用來講是吃虧的。

把視角再轉到推理,這裏需要解決的問題比訓練要多不少了:如何提供一個真正更快更便宜的推理服務?Linear Attention的優勢體現在線性計算強度,常數顯存佔用。那麼和Full Attention的計算消耗和顯存消耗必然存在一個交點,通常這個交點理論值在幾K的大小,對於今天的大模型,這個長度並不長。但是注意,這裏是理論值,我們需要解決下面幾個問題來逼近這個數值:

  • States的低精度存儲:當前Linear Attention對精度要求比Full Attention高得多;

  • 如何解決Prefix Cache:正常業務命中Cache的概率是很高的;

  • 如何優化Linear Attention上的投機解碼

幸好,這些問題目前看起來都是可以解決的。

下一步是什麼

Scaling這件事依舊是主旋律,Context Length是其中的關鍵之一,不管是Pretrain還是PostTrain,Context Length增長的趨勢越來越明顯。當GPU算力的增速慢於Data長度增長帶來的算力壓力增加的時候,Linear/Sparse Attention的收益會被逐漸釋放。我們需要提前準備些東西:

  • 更多模態、信息更加豐富的長文數據

  • 更合理的評測體系和實驗範式,幫助更快的迭代

  • 更完善的訓推基建,榨乾GPU的潛力

補充

開源推理代碼swa的實現忘記刪掉了,看到有人問為什麼最後沒有用。這裏也簡單回覆下:當然是效果不行

這個地方實驗的比較早,當時GPT-OSS還沒有開源,看到GPT-OSS結構長這樣還挺吃驚的。這裏可以簡單講講我們的一些失敗經驗。我們是CPT範式變成Hybrid SWA的思路,這裏考慮了做層間混合和層內混合兩種,做層內混合的出發點是這樣層間的計算強度是均衡的,不管是訓練做PP策略,還是推理的時候做PP或者AFD分離都更友好。

當然都沒work,具體表現為Context越長性能下降越顯著,這在Agent場景是不太能接受的。

在我們的分析裏,這裏有很多Global Attention Pattern(如retrieval head和induction head)在前期預訓練階段已經形成,通過CPT很難調整這些Attention Pattern。如果構建數據探針去檢索對應的head並將其保留為Full Attention能極大的緩解對應問題,但是不幸的是,根據人類先驗很難把這些Pattern全部探出來。

另外,這個問題和Attention sink沒有關係。

如果大家對這種思路感興趣的話,GPT-OSS、CWM、Gemma的性能大家可以分析下,特別是長文。