Stabilisatie van MoE-versterkend leren door afstemming van trainings- en inferentierouters

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

Auteurs: Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Zhifang Sui, Fuli Luo

cs.AI

Samenvatting

Versterkend leren (RL) is naar voren gekomen als een cruciale aanpak om de mogelijkheden van grote taalmodel(len) te verbeteren. Bij Mixture-of-Experts (MoE)-modellen introduceert het routeringsmechanisme echter vaak instabiliteit, wat zelfs kan leiden tot een catastrofale ineenstorting van de RL-training. Wij analyseren de consistentie tussen training en inferentie van MoE-modellen en identificeren een opmerkelijke discrepantie in routeringsgedrag tussen deze twee fasen. Bovendien kan het routeringsraamwerk, zelfs onder identieke omstandigheden, uiteenlopende expertselecties opleveren bij herhaalde forward passes. Om deze fundamentele inconsistentie aan te pakken, stellen wij Rollout Routing Replay (R³) voor, een methode die routeringsverdelingen van de inferentie-engine vastlegt en deze tijdens de training afspeelt. R³ vermindert de KL-divergentie van het train-inferentie-beleid aanzienlijk en verzacht extreme discrepanties zonder de trainingssnelheid te compromitteren. Uitgebreide experimenten in diverse settings bevestigen dat R³ erin slaagt de RL-training te stabiliseren, ineenstorting voorkomt en methoden zoals GSPO en TIS overtreft. Wij geloven dat dit werk een nieuwe oplossing kan bieden voor het stabiliseren van RL in MoE-modellen.

English

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.