Stabiliser l'apprentissage par renforcement des MoE par l'alignement des routeurs d'entraînement et d'inférence

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

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

cs.AI

papers.abstract

L'apprentissage par renforcement (RL) est apparu comme une approche cruciale pour améliorer les capacités des grands modèles de langage. Cependant, dans les modèles à mélange d'experts (MoE), le mécanisme de routage introduit souvent une instabilité, pouvant même mener à un effondrement catastrophique de l'entraînement par RL. Nous analysons la cohérence entraînement-inférence des modèles MoE et identifions un écart notable dans les comportements de routage entre les deux phases. De plus, même dans des conditions identiques, le cadre de routage peut produire des sélections d'experts divergentes sur des passes avant répétées. Pour résoudre cette incohérence fondamentale, nous proposons Rollout Routing Replay (R³), une méthode qui enregistre les distributions de routage depuis le moteur d'inférence et les rejoue pendant l'entraînement. R³ réduit significativement la divergence de politique KL entre l'entraînement et l'inférence et atténue les écarts extrêmes sans compromettre la vitesse d'entraînement. Des expériences approfondies dans divers contextes confirment que R³ réussit à stabiliser l'entraînement RL, en prévenant l'effondrement et en surpassant des méthodes telles que GSPO et TIS. Nous pensons que ces travaux peuvent offrir une nouvelle solution pour stabiliser le RL dans les modèles MoE.

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.