智东西(公众号:zhidxcom)
编译 | 程茜
编辑 | 李水青
智东西 10 月 15 日消息,10 月 14 日,小米和北京大学联合署名的论文发表于 arXiv,曾被曝获小米集团创始人兼 CEO 雷军以千万年薪招募的 DeepSeek" 天才少女 " 罗福莉,出现在了这篇论文的通讯作者之列,但值得注意的是,论文作者中并没有标注罗福莉属于小米大模型团队。
通讯作者中的罗福莉是 95 后,她本科就读于北京师范大学计算机专业,硕士毕业于北京大学计算语言学研究所计算语言学专业。随后罗福莉曾在阿里巴巴达摩院主导开发了多语言预训练模型 VECO,并推动了 AliceMind 的开源工作,2022 年入职 DeepSeek,参与了 MoE 大模型 DeepSeek-V2 的研发。去年年底,小米被曝以千万年薪挖角 DeepSeek-V2 核心开发者之一罗福莉,使其冲上热搜,但双方至今都未公开声明是否正式入职小米。
▲ DeepSeek" 天才少女 " 罗福莉(图源:罗福莉个人公众号)
当下,强化学习(RL)已成为提升大语言模型能力的关键方法。然而,在 MoE 模型中,路由机制往往会引入不稳定性,甚至导致强化学习训练崩溃,但现有的引入重要性采样机制等并不能提升训练稳定性。不同于此前采取诸如丢弃差异较大的数据之类的变通方法,这篇论文的研究人员希望通过解决路由分布也就是 R3 来根本性解决这个问题。
论文地址:https://arxiv.org/pdf/2510.11370
一、破解强化学习崩溃的关键方法,小米团队提出 R3
强化学习已成为大语言模型后期训练的基石,利用大规模强化学习,大模型更深入、更广泛推理,获得解决复杂问题所需的高级能力,但其面临的关键挑战是如何平衡效率和稳定性。
现代强化学习框架通常使用不同的引擎进行推理和训练用于部署,但这种架构上的分离可能导致 token 概率出现分歧,甚至可能导致灾难性的强化学习崩溃。然而,现有的改进方法并不能完全解决 MoE 模型上进行强化学习训练时出现的强化学习离线策略问题。
研究人员提出的 R3,其工作原理是在序列生成期间从推理引擎捕获路由分布,并将其直接重放到训练引擎中。这一过程可以缩小训练和推理之间的差距,其显著特征是不同引擎生成的逻辑向量的 KL 散度(量化两个概率分布之间的差异程度,值越小说明两个分布越接近)显著降低,两个阶段之间概率差异显著的 token 数量减少了大约一个数量级。
此外,该方法同时适用于在线策略(on-policy)和小批量(mini-batch)式离线策略强化学习(off-policy)场景。
论文提到了研究团队的三大主要贡献:
1、系统识别和分析了 MoE 模型中训练和推理之间的路由分布差异,强调了它们在训练不稳定性中的作用;
2、提出 Rollout Routing Replay,它重用训练引擎内部的推理时间路由分布,以协调训练和推理之间的路由行为;
3、将 R3 应用于多种 RL 设置进行 MoE 强化学习,并表明 R3 在稳定性和整体性能方面优于 GSPO 和 TIS。
二、可显著缩小训练 - 推理差异,对 Agent 任务大有裨益
R3 的主要思路是在训练前向传播过程中重用推理路由掩码 I,同时仍将 softmax 应用于训练逻辑以保持梯度流。
这种设计主要有两个目的:一是对齐训练和推理,确保训练重放期间使用的专家与推理期间选择的专家相匹配,从而消除专家选择中的不匹配;二是保留梯度数据流,通过仅重放掩码,梯度仍然可以流回 logits 而不会干扰计算图,这有助于有效地优化路由器。
▲重放门控权重、重放输出 y 的计算方式
其论文提到,缓存的路由掩码具有相似的属性,对于相同的前缀 token,MoE 路由器应该产生相同的结果,因此来自推理引擎的路由掩码可以与前缀 KVCache 一起缓存。
对于每个层和 token 前缀,相应的路由掩码都存储在 KVCache 中。当相同的前缀出现并命中缓存时,这些掩码可以被重用,从而无需重新计算,这使得 R3 能够与前缀缓存机制无缝集成。
研究人员称,缓存路由掩码在 Agent 场景中有较大应用空间。例如软件工程和网页浏览等 Agent 任务,都涉及自回归生成和工具调用之间的多轮交互,为了提高效率,这些过程直接重用了前几轮的 KVCache,因此无需重新生成已计算的数据。路由掩码缓存使 R3 能够在强化学习代理任务中保持高效,而无需重新预填充以生成路由掩码。
为了证明 R3 在缩小训练 - 推理差异上的有效性,研究人员使用 Qwen3-30B-A3B 模型进行了验证,其将推理过程中获得的路由分布缓存在 SGLang 上,并在 Megatron 框架内重放它们。
▲使用 Megatron 进行两次前向传播获得的概率
研究人员还绘制了使用 R3 的训练 - 推理差异比率的累积分布图,对于 MoE 模型,应用 R3 可将具有较大训练推理差异的 token 的频率降低一个数量级。
▲ a、MoE 模型中训练 - 推理差异的说明,b、MoE+R3 模型中训练 - 推理差异的说明,c、稠密模型中训练 - 推理差异的说明,d、极端 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 分。
▲主要评估结果
训练稳定性方面:如 GRPO、GRPO+TIS 等无 R3 的方法在单步更新场景中均出现崩盘,GRPO 在 60 步崩盘、GRPO+TIS 在 105 步崩盘。
引入 R3 后,所有组合方法均无崩盘,且训练过程中训练 - 推理 KL 散度等始终较低。
▲多步更新训练 - 推理崩溃分析
▲ wen3-30B-A3B-Base 训练动态
结语:聚焦 MoE 模型训练难题,小米提出新思路
MoE 架构如今已成为扩展现代语言模型的基石,其采用门控网络,对每个 token 稀疏地仅激活一部分专家参数,从而将模型的总参数数量与其推理成本分离开来,从而大幅提升了模型容量。然而,由于门控网络的敏感性,MoE 模型容易受到训练不稳定性的影响,这使得路由稳健性成为有效模型收敛的核心挑战。
在这篇论文中,研究人员在训练过程中重用推理时的路由分布,以在保留梯度流的同时对齐专家选择。这种思路或为行业提供了新的研究思路。
登录后才可以发布评论哦
打开小程序可以发布评论哦