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