Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers
Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Fuli Luo, Zhifang Sui
Abstract
Reinforcement learning (RL) has emerged as a crucial approach for enhancing the capabilities of large language models. However, in Mixture-of-Experts (MoE) models, the routing mechanism often introduces instability, even leading to catastrophic RL training collapse. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. Moreover, even under identical conditions, the routing framework can yield divergent expert selections across repeated forward passes. To address this foundational inconsistency, we propose Rollout Routing Replay (R3), a method that records routing distributions from the inference engine and replays them during training. R3 significantly reduces training-inference policy KL divergence and mitigates extreme discrepancies without compromising training speed. Extensive experiments on various settings confirm that R3 succeeds in stabilizing RL training, preventing collapse and outperforming methods such as GSPO and TIS. We believe this work can offer a new solution for stabilizing RL in MoE models.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 73323fb5-edcf-4f7b-a6d1-db49f953f7deCited by top-tier papers8
- Rethinking the Trust Region in LLM Reinforcement LearningPenghui Qi, Xiangxin Zhou, Zichen Liu, Tianyu Pang et al.ICML 2026 · 22 citations
- PyVision-RL: Forging Open Agentic Vision Models via RLShitian Zhao, Shaoheng Lin, Ming Li, Haoquan Zhang et al.ICML 2026 · 10 citations
- Grouter: Decoupling Routing from Representation for Accelerated MoE TrainingYuqi Xu, Rizhen Hu, zihan liu, Mou Sun et al.ICML 2026 · 7 citations
- Graph-GRPO: Training Graph Flow Models with Reinforcement LearningBaoheng Zhu, Deyu Bo, Delvin Zhang, Xiao WangICML 2026 · 3 citations
- Stable Asynchrony: Variance-Controlled Off-Policy RL for LLMsLuke Huang, Zhuoyang Zhang, Qinghao Hu, Shang Yang et al.ICML 2026 · 3 citations
Builds on9
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- DAPO: An Open-Source LLM Reinforcement Learning System at ScaleQiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan et al.NeurIPS 2025 · 2,828 citations
- SWE-bench: Can Language Models Resolve Real-world Github Issues?Carlos E. Jimenez, John Yang, Alexander Wettig, Shunyu Yao et al.ICLR 2024 · 2,082 citations
- SGLang: Efficient Execution of Structured Language Model ProgramsLianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun et al.NeurIPS 2024 · 1,586 citations
- Efficient Memory Management for Large Language Model Serving with PagedAttentionWoosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng et al.SOSP 2023 · 1,016 citations
Related papers
- Towards Stable and Effective Reinforcement Learning for Mixture-of-ExpertsDi Zhang, Xun Wu, Shaohan Huang, Lingjie Jiang et al.ACL 2026 · 3 citations
- StableMoE: Stable Routing Strategy for Mixture of ExpertsDamai Dai, Li Dong, Shuming Ma, Bo Zheng et al.ACL 2022
- MoE-GRPO: Optimizing Mixture-of-Experts via Reinforcement Learning in Vision-Language ModelsDohwan Ko, Jinyoung Park, Seoung Choi, Sanghyeok Lee et al.CVPR 2026 · 3 citations
- Routing Experts: Learning to Route Dynamic Experts in Existing Multi-modal Large Language ModelsQiong Wu, Zhaoxi Ke, Yiyi Zhou, Xiaoshuai Sun et al.ICLR 2025
- Breaking the Echo Chamber: A Dynamic Ensemble Pruning Perspective on MoEXinlai Kang, Dunyao Xue, Zhengbo Wang, Chengshuo Du et al.ICML 2026
