小米的最新大模型科研成果,對外曝光了。
就在最近,小米AI團隊攜手北京大學聯合發佈了一篇聚焦MoE與強化學習的論文。
而其中,因為更早之前在DeepSeek R1爆火前轉會小米的羅福莉,也赫然在列,還是通訊作者。
羅福莉碩士畢業於北京大學,這次也算是因AI串聯起了小米和北大。
有意思的是,就在今年9月DeepSeek登上《Nature》的時候,羅福莉也出現在了作者名單,不過是以“北京獨立研究者”的身份。
當時還有過風言風語,說當初“雷軍千萬年薪挖來AI天才少女”,當事人可能離職了。
但這篇小米最新AI論文披露後,一切似乎有了答案…
這篇論文大道至簡,提出了一種在MoE架構中提高大模型強化學習的思路。
相對已經共識的是,當前強化學習已成為在預訓練遇到瓶頸後,推動LLM突破能力邊界的關鍵工具。
不過在MoE架構中,情況就沒那麼簡單了,由於需要根據問題分配不同的專家,路由機制會讓訓練過程變得不穩定,嚴重時甚至會直接把模型“整崩”。
為瞭解決這個問題,研究團隊提出了一種全新的思路,讓MoE也能平穩且高效地推進大規模強化學習。
自從預訓練時代告一段落,後訓練成了巨頭們拿起Scaling Law瞄準的的下一個戰場。
靠著大規模強化學習,大模型開始學會更長鏈路的推理,也能搞定那些需要呼叫工具的複雜Agent任務。
不過,強化學習在擴展規模的過程中,總會不可避免地撞上一道鐵幕:效率和穩定性的權衡。
想要高效率,就得訓練得更“猛”——更高的學習率、更大的平行度、更頻繁的樣本更新。可這樣一來,穩定性也更容易出現問題。
但一味追求穩定也不行,效率會被拖住,模型訓練慢得像蝸牛。
想要解決這個問題,得先回到強化學習的底層一探究竟。
LLM的強化學習,通常分兩步:
第一步是推理,模型自己生成內容、和環境互動、拿到反饋分數;
第二步是訓練,根據這些分數去微調自己,並想辦法在下次拿更高分。
不過,這兩步通常不是在同一套系統裡跑的。
比如,現在主流方案是SGLang負責生成內容,追求速度快;而Megatron負責訓練更新,追求算得準。
雖然兩邊用的是同一套模型參數,但底層實現有細微差別,比如像隨機性、精度、平行方式、快取策略,這些看似微不足道的細節波動,都會讓結果出現偏差。
於是就出現了一個尷尬現象:
一模一樣的Prompt,兩套模式下最終生成的結果都能不一樣。
這種「機率漂移」積累多了,模型就會越學越偏,最後學著學著,訓練目標和實際表現徹底牛頭不對馬嘴。
這就是業內常說,強化學習災難性崩潰。
研究團隊指出,導致MoE在強化學習中容易崩掉的罪魁禍首,在於路由分佈。
在MoE模型中,路由器不會把所有參數都用上,而是會根據每個輸入token的特徵,挑幾位在該領域更擅長的“專家”出來幹活,從而可以節省不少資源。
但副作用也很明顯,這種動態模式會讓模型在訓練階段和推理階段得出的最佳策略大相逕庭,比傳統的稠密模型要“飄忽”得多。
對此,這篇論文給出了一種新穎的解決方案。
既然問題出在路由隨機,那為何不直接把路由鎖住呢?
他們的做法是:在推理時把路由分佈記錄下來,等到訓練時再把這些分佈原封不動地“重放”進去。
這樣,訓練和推理就走同一條路線,不再各幹各的。
根據這種“重放”的特定,研究將這種方法命名為——Rollout Routing Replay(R3)。
解決了穩定性的問題,再來看看如何把效率也穩穩拿下。
在強化學習中,模型會不斷重複“生成→獲得獎勵→更新→再生成”的飛輪,一個完整過程下來,可能要跑上幾十萬、甚至上百萬次推理。
要是每次生成都要從頭計算上下文,算力與時間成本將呈幾何式增長。
為應對這種情況,主流推理引擎普遍採用KVCache前綴快取策略:把之前算好的上下文保存下來,下次直接“接著算”。
不過,除了上下文不一致,MoE架構還涉及到路由選擇不一致的問題——按照傳統的解決方案,即便是重複的上下文,每一次計算,模型還是要重新選專家、啟動專家。
因此,研究團隊在KVCache的基礎上又加了一招——路由掩碼(routing mask)。
他們的想法是,既然對於對相同的上下文,MoE的路由結果應該一樣,那乾脆,把推理階段的路由掩碼和前綴KVCache一起快取起來。
這樣當相同上下文再次出現時,模型就能直接用上次的掩碼,不必重算。
這樣,R3就能夠與現有的前綴快取系統無縫銜接,在大規模強化學習及複雜的Agent任務中,也依然能保持出色的計算效率。
為評估R3的實際效果,研究團隊基於Qwen3-30B-A3B模型進行了一系列實驗。
結果發現,不管在那種場景下,R3的整體成績都更好。
在多mini-step設定下,GRPO+R3的表現比GSPO高出1.29分。
若將R3與GSPO結合,性能還可以進一步提升0.95分。
崩潰情況也少了很多。
不難看出,隨著訓練時間的延長,即便到了第150步,R3依然能保持相對平緩的曲線。
相比之下,如果是用GRPO訓練,到第60步時就已經嚴重跑偏。
而且,R3不光讓模型更穩,也讓它更聰明。
實驗結果結果表明,R3能更快找到正確方向、最佳化過程更絲滑,還能更早開始探索更優策略。
一句話總結,研究團隊在這篇論文提出了一種叫R3的方法,通過在訓練中復用推理階段的路由分佈,能夠讓MoE模型的強化學習更穩定、更高效。
說完論文,再讓我們看看這支由小米系和北京大學攜手牽起的研究團隊。
論文的第一作者叫Wenhan Ma。
資料不多,只知道Wenhan是小米LLM-Core團隊的研究員,而且還是實習生。
此前,他還曾參與過小米MiMo模型與多模態MiMo-VL的研發。
相比起來,這篇論文的兩名通訊作者,大家可能更耳熟能詳一點。
一位是羅福莉。
羅福莉本科畢業於北京師範大學電腦專業,碩士階段進入北京大學計算語言學深造。期間,她在不少NLP頂級會議上都發表過論文。
碩士畢業後,羅福莉加入阿里巴巴達摩院,擔任機器智能實驗室研究員,負責開發多語言預訓練模型VECO,並推動AliceMind項目的開源工作。
2022年,羅福莉加入DeepSeek母公司幻方量化從事深度學習相關工作,後又擔任DeepSeek的深度學習研究員,參與研發DeepSeek-V2等模型。
截至目前,羅福莉的學術論文總引用次數已超過1.1萬次,僅在今年一年內就新增了約八千次引用。
而另一名通訊作者,正是羅福莉的北大碩士導師——穗志方。
穗教授是北京大學資訊科學技術學院的教授、博士生導師,長期從事計算語言學、文字挖掘與知識工程研究,在NLP與AI領域發表了大量高水平論文。
但稍有有個新問題,在這篇論文成果的單位註釋中,羅福莉的單位沒有被明確,她既不是北大的,也沒有被歸入小米。
咦……依然是獨立研究者? (量子位)