MiniMax M2采用“Full Attention”是開倒車？

本文轉載自「MiniMax稀宇科技」

為什么M2是Full Attention？ Guixianren，公眾號：MiniMax 稀宇科技

作為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的性能大家可以分析下，特別是長文。