Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Fuli Luo, Zhifang Sui
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:85320-85343, 2026.

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, RL training tends to be unstable and can even collapse entirely. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. To address this issue, we propose Rollout Routing Replay (R3), a novel and effective 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 strong baselines. R3 is orthogonal to most policy optimization algorithm improvements, allowing it to be used in conjunction with them. We believe this work can offer a new solution for stabilizing RL in MoE model.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-ma26aw, title = {Stabilizing {M}o{E} Reinforcement Learning by Aligning Training and Inference Routers}, author = {Ma, Wenhan and Zhang, Hailin and Zhao, Liang and Song, Yifan and Wang, Yudong and Luo, Fuli and Sui, Zhifang}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {85320--85343}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/ma26aw/ma26aw.pdf}, url = {https://proceedings.mlr.press/v306/ma26aw.html}, 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, RL training tends to be unstable and can even collapse entirely. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. To address this issue, we propose Rollout Routing Replay (R3), a novel and effective 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 strong baselines. R3 is orthogonal to most policy optimization algorithm improvements, allowing it to be used in conjunction with them. We believe this work can offer a new solution for stabilizing RL in MoE model.} }
Endnote
%0 Conference Paper %T Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers %A Wenhan Ma %A Hailin Zhang %A Liang Zhao %A Yifan Song %A Yudong Wang %A Fuli Luo %A Zhifang Sui %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-ma26aw %I PMLR %P 85320--85343 %U https://proceedings.mlr.press/v306/ma26aw.html %V 306 %X Reinforcement learning (RL) has emerged as a crucial approach for enhancing the capabilities of large language models. However, in Mixture-of-Experts (MoE) models, RL training tends to be unstable and can even collapse entirely. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. To address this issue, we propose Rollout Routing Replay (R3), a novel and effective 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 strong baselines. R3 is orthogonal to most policy optimization algorithm improvements, allowing it to be used in conjunction with them. We believe this work can offer a new solution for stabilizing RL in MoE model.
APA
Ma, W., Zhang, H., Zhao, L., Song, Y., Wang, Y., Luo, F. & Sui, Z.. (2026). Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:85320-85343 Available from https://proceedings.mlr.press/v306/ma26aw.html.

Related Material