近日,一篇由小米與北京大學(xué)聯(lián)合署名的論文在arXiv平臺發(fā)布,引發(fā)科技圈關(guān)注。論文中一個引人注目的細節(jié)是,此前被傳以千萬年薪被小米集團創(chuàng)始人兼CEO雷軍招攬的DeepSeek“天才少女”羅福莉,出現(xiàn)在通訊作者名單里,不過論文未標注其是否屬于小米大模型團隊。
羅福莉出生于1995年,本科階段在北京師范大學(xué)計算機專業(yè)深造,碩士畢業(yè)于北京大學(xué)計算語言學(xué)研究所的計算語言學(xué)專業(yè)。畢業(yè)后,她加入阿里巴巴達摩院,主導(dǎo)開發(fā)了多語言預(yù)訓(xùn)練模型VECO,還推動了AliceMind的開源工作。2022年,羅福莉入職DeepSeek,參與MoE大模型DeepSeek - V2的研發(fā)。去年年底,有消息稱小米以千萬年薪挖角羅福莉,此事一度登上熱搜,但雙方至今都未公開聲明其是否正式入職小米。
此次發(fā)表的論文聚焦于提升MoE模型強化學(xué)習(xí)訓(xùn)練的新方法——Rollout Routing Replay(R3)。實驗數(shù)據(jù)顯示,R3的整體性能優(yōu)于GRPO、TIS等強化學(xué)習(xí)領(lǐng)域用于提升模型性能的優(yōu)化算法。引入R3的所有組合方法在訓(xùn)練全程未出現(xiàn)崩盤情況,訓(xùn)練過程中訓(xùn)練 - 推理KL散度始終維持在較低水平,且在不影響訓(xùn)練速度的前提下,使極端token比例減少了一個量級。
當(dāng)下,強化學(xué)習(xí)(RL)在提升大語言模型能力方面扮演著關(guān)鍵角色。通過大規(guī)模強化學(xué)習(xí),大模型能夠進行更深入、廣泛的推理,從而獲得解決復(fù)雜問題所需的高級能力。然而,在MoE模型中,路由機制常常會引入不穩(wěn)定性,甚至導(dǎo)致強化學(xué)習(xí)訓(xùn)練崩潰。現(xiàn)有的引入重要性采樣機制等方法,并不能有效提升訓(xùn)練穩(wěn)定性。與以往采取丟棄差異較大數(shù)據(jù)等變通方法不同,此次論文的研究人員希望通過解決路由分布問題,即R3方法,從根本上解決這一難題。
論文提出的R3方法,其工作原理是在序列生成期間,從推理引擎捕獲路由分布,并直接將其重放到訓(xùn)練引擎中。這一過程縮小了訓(xùn)練和推理之間的差距,顯著特征是不同引擎生成的邏輯向量的KL散度明顯降低,兩個階段之間概率差異顯著的token數(shù)量減少了大約一個數(shù)量級。該方法同時適用于在線策略(on - policy)和小批量(mini - batch)式離線策略強化學(xué)習(xí)(off - policy)場景。
論文闡述了研究團隊的三大主要貢獻:一是系統(tǒng)識別并分析了MoE模型中訓(xùn)練和推理之間的路由分布差異,強調(diào)了這些差異在訓(xùn)練不穩(wěn)定性中的作用;二是提出Rollout Routing Replay方法,重用訓(xùn)練引擎內(nèi)部的推理時間路由分布,以協(xié)調(diào)訓(xùn)練和推理之間的路由行為;三是將R3應(yīng)用于多種RL設(shè)置進行MoE強化學(xué)習(xí),結(jié)果顯示R3在穩(wěn)定性和整體性能方面優(yōu)于GSPO和TIS。
R3的主要設(shè)計思路是在訓(xùn)練前向傳播過程中重用推理路由掩碼I,同時仍將softmax應(yīng)用于訓(xùn)練邏輯以保持梯度流。這一設(shè)計有兩個目的:一是對齊訓(xùn)練和推理,確保訓(xùn)練重放期間使用的專家與推理期間選擇的專家相匹配,消除專家選擇中的不匹配;二是保留梯度數(shù)據(jù)流,僅重放掩碼可使梯度流回logits而不干擾計算圖,有助于有效優(yōu)化路由器。
在效率優(yōu)化方面,R3通過路由掩碼緩存(Router Mask Caching)適配多輪對話場景,降低計算開銷。緩存的路由掩碼具有相似屬性,對于相同的前綴token,MoE路由器應(yīng)產(chǎn)生相同結(jié)果,因此來自推理引擎的路由掩碼可與前綴KVCache一起緩存。對于每個層和token前綴,相應(yīng)的路由掩碼存儲在KVCache中。當(dāng)相同前綴出現(xiàn)并命中緩存時,這些掩碼可被重用,無需重新計算,使R3能與前綴緩存機制無縫集成。研究人員稱,緩存路由掩碼在Agent場景中有較大應(yīng)用空間,如軟件工程和網(wǎng)頁瀏覽等Agent任務(wù),涉及自回歸生成和工具調(diào)用之間的多輪交互,為提高效率,這些過程直接重用前幾輪的KVCache,路由掩碼緩存使R3能在強化學(xué)習(xí)代理任務(wù)中保持高效,無需重新預(yù)填充以生成路由掩碼。
為證明R3在縮小訓(xùn)練 - 推理差異上的有效性,研究人員使用Qwen3 - 30B - A3B模型進行驗證,將推理過程中獲得的路由分布緩存在SGLang上,并在Megatron框架內(nèi)重放。結(jié)果顯示,應(yīng)用R3后,訓(xùn)練和推理之間的KL散度從1.5×10?3減小到7.5×10??,接近稠密模型的6.4×10??水平,表明訓(xùn)練 - 推理差異減少。研究人員還繪制了使用R3的訓(xùn)練 - 推理差異比率的累積分布圖,對于MoE模型,應(yīng)用R3可將具有較大訓(xùn)練推理差異的token的頻率降低一個數(shù)量級。
為評估R3對強化學(xué)習(xí)的性能改進,研究人員從BigMath、ORZ等開源數(shù)據(jù)集篩選約10萬道可驗證數(shù)學(xué)題,采用AIME24、AIME25、AMC23和MATH500作為基準數(shù)據(jù)集進行評估,并在單次訓(xùn)練過程中每5個全局步驟測量一次模型性能。選擇的模型是Qwen3 - 30B - A3B - Base及其微調(diào)模型Qwen3 - 30B - A3B - SFT。評估方式是每5個全局步驟記錄模型性能,最終報告最佳性能及對應(yīng)訓(xùn)練步驟,若模型后期性能驟降,同時追蹤訓(xùn)練崩盤步驟。
實驗結(jié)果顯示,在整體性能上,R3在多步更新場景中表現(xiàn)突出,GRPO + R3平均得分68.05分,比GSPO高出1.29分;GSPO + R3進一步提升至69.00分,比單獨GSPO高2.24分。在單步更新場景中,SFT模型上,GRPO + R3平均得分71.83分,比GRPO(62.23)高9.6分,比GRPO + TIS(66.24)高5.59分;Base模型上,GRPO + R3平均得分70.73,比GRPO(61.69)高9.04分。研究人員還發(fā)現(xiàn),將R3與TIS結(jié)合使用并不能帶來明顯的性能提升,甚至可能降低性能,例如在SFT模型的單小步設(shè)置下,TIS + R3的得分比單獨使用R3低1.69分,這是因為R3已顯著降低了訓(xùn)練和推理之間的策略差異,TIS的額外校正效果微乎其微。
在訓(xùn)練穩(wěn)定性方面,GRPO、GRPO + TIS等無R3的方法在單步更新場景中均出現(xiàn)崩盤,GRPO在60步崩盤、GRPO + TIS在105步崩盤。而引入R3后,所有組合方法均無崩盤,且訓(xùn)練過程中訓(xùn)練 - 推理KL散度等始終較低。
在優(yōu)化與生成行為方面,R3在訓(xùn)練過程中能增強優(yōu)化穩(wěn)定性、探索行為和生成動態(tài)。研究人員繪制的單步 + 基礎(chǔ)模型組訓(xùn)練過程中的序列長度、梯度范數(shù)、生成熵和評估分數(shù)顯示,R3具有更小的梯度范數(shù)、更平滑的序列增長模式和更穩(wěn)定的熵。使用R3時,生成的序列長度在訓(xùn)練開始時迅速上升,表明R3能快速捕捉到正確的優(yōu)化方向,相比之下其他兩個訓(xùn)練過程在第80步之后才緩慢上升,且波動更為明顯;R3始終保持較低的梯度范數(shù),表明優(yōu)化過程更加穩(wěn)定;使用R3時,熵在大約第25步后開始穩(wěn)步上升,表明模型更早地開始探索更優(yōu)策略,不使用R3時,熵上升得更晚,且波動較大。
MoE架構(gòu)已成為擴展現(xiàn)代語言模型的基石,它采用門控網(wǎng)絡(luò),對每個token稀疏地僅激活一部分專家參數(shù),將模型的總參數(shù)數(shù)量與其推理成本分離開來,大幅提升了模型容量。但由于門控網(wǎng)絡(luò)的敏感性,MoE模型容易受到訓(xùn)練不穩(wěn)定性的影響,路由穩(wěn)健性成為有效模型收斂的核心挑戰(zhàn)。此次論文中研究人員在訓(xùn)練過程中重用推理時的路由分布,在保留梯度流的同時對齊專家選擇,為行業(yè)提供了新的研究思路。








