小米与北大联合发文!“天才少女”罗福莉参与,提出MoE模型新解法

   时间:2025-10-16 18:00 来源:快讯作者:顾青青

近日,一篇由小米与北京大学联合署名的论文在arXiv平台发布,引发科技圈关注。论文中一个引人注目的细节是,此前被传以千万年薪被小米集团创始人兼CEO雷军招揽的DeepSeek“天才少女”罗福莉,出现在通讯作者名单里,不过论文未标注其是否属于小米大模型团队。

罗福莉出生于1995年,本科阶段在北京师范大学计算机专业深造,硕士毕业于北京大学计算语言学研究所的计算语言学专业。毕业后,她加入阿里巴巴达摩院,主导开发了多语言预训练模型VECO,还推动了AliceMind的开源工作。2022年,罗福莉入职DeepSeek,参与MoE大模型DeepSeek - V2的研发。去年年底,有消息称小米以千万年薪挖角罗福莉,此事一度登上热搜,但双方至今都未公开声明其是否正式入职小米。

此次发表的论文聚焦于提升MoE模型强化学习训练的新方法——Rollout Routing Replay(R3)。实验数据显示,R3的整体性能优于GRPO、TIS等强化学习领域用于提升模型性能的优化算法。引入R3的所有组合方法在训练全程未出现崩盘情况,训练过程中训练 - 推理KL散度始终维持在较低水平,且在不影响训练速度的前提下,使极端token比例减少了一个量级。

当下,强化学习(RL)在提升大语言模型能力方面扮演着关键角色。通过大规模强化学习,大模型能够进行更深入、广泛的推理,从而获得解决复杂问题所需的高级能力。然而,在MoE模型中,路由机制常常会引入不稳定性,甚至导致强化学习训练崩溃。现有的引入重要性采样机制等方法,并不能有效提升训练稳定性。与以往采取丢弃差异较大数据等变通方法不同,此次论文的研究人员希望通过解决路由分布问题,即R3方法,从根本上解决这一难题。

论文提出的R3方法,其工作原理是在序列生成期间,从推理引擎捕获路由分布,并直接将其重放到训练引擎中。这一过程缩小了训练和推理之间的差距,显著特征是不同引擎生成的逻辑向量的KL散度明显降低,两个阶段之间概率差异显著的token数量减少了大约一个数量级。该方法同时适用于在线策略(on - policy)和小批量(mini - batch)式离线策略强化学习(off - policy)场景。

论文阐述了研究团队的三大主要贡献:一是系统识别并分析了MoE模型中训练和推理之间的路由分布差异,强调了这些差异在训练不稳定性中的作用;二是提出Rollout Routing Replay方法,重用训练引擎内部的推理时间路由分布,以协调训练和推理之间的路由行为;三是将R3应用于多种RL设置进行MoE强化学习,结果显示R3在稳定性和整体性能方面优于GSPO和TIS。

R3的主要设计思路是在训练前向传播过程中重用推理路由掩码I,同时仍将softmax应用于训练逻辑以保持梯度流。这一设计有两个目的:一是对齐训练和推理,确保训练重放期间使用的专家与推理期间选择的专家相匹配,消除专家选择中的不匹配;二是保留梯度数据流,仅重放掩码可使梯度流回logits而不干扰计算图,有助于有效优化路由器。

在效率优化方面,R3通过路由掩码缓存(Router Mask Caching)适配多轮对话场景,降低计算开销。缓存的路由掩码具有相似属性,对于相同的前缀token,MoE路由器应产生相同结果,因此来自推理引擎的路由掩码可与前缀KVCache一起缓存。对于每个层和token前缀,相应的路由掩码存储在KVCache中。当相同前缀出现并命中缓存时,这些掩码可被重用,无需重新计算,使R3能与前缀缓存机制无缝集成。研究人员称,缓存路由掩码在Agent场景中有较大应用空间,如软件工程和网页浏览等Agent任务,涉及自回归生成和工具调用之间的多轮交互,为提高效率,这些过程直接重用前几轮的KVCache,路由掩码缓存使R3能在强化学习代理任务中保持高效,无需重新预填充以生成路由掩码。

为证明R3在缩小训练 - 推理差异上的有效性,研究人员使用Qwen3 - 30B - A3B模型进行验证,将推理过程中获得的路由分布缓存在SGLang上,并在Megatron框架内重放。结果显示,应用R3后,训练和推理之间的KL散度从1.5×10⁻³减小到7.5×10⁻⁴,接近稠密模型的6.4×10⁻⁴水平,表明训练 - 推理差异减少。研究人员还绘制了使用R3的训练 - 推理差异比率的累积分布图,对于MoE模型,应用R3可将具有较大训练推理差异的token的频率降低一个数量级。

为评估R3对强化学习的性能改进,研究人员从BigMath、ORZ等开源数据集筛选约10万道可验证数学题,采用AIME24、AIME25、AMC23和MATH500作为基准数据集进行评估,并在单次训练过程中每5个全局步骤测量一次模型性能。选择的模型是Qwen3 - 30B - A3B - Base及其微调模型Qwen3 - 30B - A3B - SFT。评估方式是每5个全局步骤记录模型性能,最终报告最佳性能及对应训练步骤,若模型后期性能骤降,同时追踪训练崩盘步骤。

实验结果显示,在整体性能上,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分。研究人员还发现,将R3与TIS结合使用并不能带来明显的性能提升,甚至可能降低性能,例如在SFT模型的单小步设置下,TIS + R3的得分比单独使用R3低1.69分,这是因为R3已显著降低了训练和推理之间的策略差异,TIS的额外校正效果微乎其微。

在训练稳定性方面,GRPO、GRPO + TIS等无R3的方法在单步更新场景中均出现崩盘,GRPO在60步崩盘、GRPO + TIS在105步崩盘。而引入R3后,所有组合方法均无崩盘,且训练过程中训练 - 推理KL散度等始终较低。

在优化与生成行为方面,R3在训练过程中能增强优化稳定性、探索行为和生成动态。研究人员绘制的单步 + 基础模型组训练过程中的序列长度、梯度范数、生成熵和评估分数显示,R3具有更小的梯度范数、更平滑的序列增长模式和更稳定的熵。使用R3时,生成的序列长度在训练开始时迅速上升,表明R3能快速捕捉到正确的优化方向,相比之下其他两个训练过程在第80步之后才缓慢上升,且波动更为明显;R3始终保持较低的梯度范数,表明优化过程更加稳定;使用R3时,熵在大约第25步后开始稳步上升,表明模型更早地开始探索更优策略,不使用R3时,熵上升得更晚,且波动较大。

MoE架构已成为扩展现代语言模型的基石,它采用门控网络,对每个token稀疏地仅激活一部分专家参数,将模型的总参数数量与其推理成本分离开来,大幅提升了模型容量。但由于门控网络的敏感性,MoE模型容易受到训练不稳定性的影响,路由稳健性成为有效模型收敛的核心挑战。此次论文中研究人员在训练过程中重用推理时的路由分布,在保留梯度流的同时对齐专家选择,为行业提供了新的研究思路。

 
 
更多>同类内容
全站最新
热门内容