提出PR2方法，解决MoE大模型强化学习中的路由漂移问题。
PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning
- 用轻量预测器预估路由器短期演化，指导推理时的专家选择。
- 实验显示PR2显著降低路由不匹配，提升强化学习稳定性与推理性能。
- 适合研究高效训练MoE模型或强化学习优化的开发者参考。
基于混合专家（MoE）的大语言模型在大规模下表现优异，但其强化学习（RL）训练常因路由漂移而出现不稳定。根本原因在于：路由器在模型更新中激活模式剧烈变化，导致推理与训练阶段的专家激活差异显著，引发重要性采样权重波动。现有路由重播方法虽冻结推理轨迹中的路由路径，却忽视了路由器在离策略更新下的动态演化，造成路由过时。为此，我们提出预测性路由重播（PR2），为每个路由器添加轻量级演化预测器，学习短时序内路由器的演化趋势。推理阶段，使用预测的路由分布进行top-k路由，使梯度能传递至未来可能激活的专家；训练阶段，重播预测后的路由路径，维持重要性估计的一致性。理论分析与实验证明，PR2有效减少路由引起的不匹配，提升训练稳定性，并在多个推理基准上实现更强性能。
原文摘要 · Abstract (English)
Mixture of Experts (MoE) Large Language Models (LLMs) achieve strong performance at scale. However, reinforcement learning (RL) on MoE-based LLMs often suffers from training instability. A root cause is router drift, i.e., expert activations can change drastically across model updates and differ between disaggregated rollout and training phases, causing large rollout--training mismatch and unstable importance sampling weights in PPO-style RL algorithms. Routing replay mitigates this issue by freezing the replay route within each reasoning trajectory, but it ignores how the router evolves under off-policy updates and thus causes router staleness. To address this limitation, we propose Predictive Routing Replay (PR2), which augments each router with a lightweight evolution predictor that learns to anticipate short-horizon router evolution. During the rollout phase, we use the predictive routing distribution to apply top-$k$ routing, enabling gradients to reach experts that are likely to become active after updates. During the training phase, we replay the resulting predicted route to retain consistency for stable importance estimation. Theoretical analysis and experiments support that PR2 reduces routing-induced mismatch, improves RL stability, and yields stronger performance across various reasoning benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。