Towards Stable and Effective Reinforcement Learning for Mixture-of-Experts
Di Zhang, Xun Wu, Shaohan Huang, Yaru Hao, Li Dong, Zewen Chi, Zhifang Sui, Furu Wei
TL;DR
This work tackles instability in off-policy reinforcement learning for Mixture-of-Experts language models caused by router fluctuations. It introduces Router-Shift Policy Optimization (RSPO), which adds a router shift ratio to softly down-weight updates when routing drift is large and retains sequence-level importance sampling to align with sequence rewards. By combining token-level clipping with a multiplicative router-shift weight and a geometric-mean aggregation, RSPO achieves greater training stability and improved final performance on multi-benchmark mathematical reasoning tasks compared with GRPO, GSPO, and GMPO. Experiments on both a small MoE model and a large MoE model demonstrate that RSPO maintains higher token entropy and progressively stable routing, making it a robust and scalable approach for MoE RL in large-scale reasoning systems.
Abstract
Recent advances in reinforcement learning (RL) have substantially improved the training of large-scale language models, leading to significant gains in generation quality and reasoning ability. However, most existing research focuses on dense models, while RL training for Mixture-of-Experts (MoE) architectures remains underexplored. To address the instability commonly observed in MoE training, we propose a novel router-aware approach to optimize importance sampling (IS) weights in off-policy RL. Specifically, we design a rescaling strategy guided by router logits, which effectively reduces gradient variance and mitigates training divergence. Experimental results demonstrate that our method significantly improves both the convergence stability and the final performance of MoE models, highlighting the potential of RL algorithmic innovations tailored to MoE architectures and providing a promising direction for efficient training of large-scale expert models.