Estabilización del Aprendizaje por Refuerzo en MoE mediante la Alineación de los Enrutadores de Entrenamiento e Inferencia

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

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

cs.AI

Resumen

El aprendizaje por refuerzo (RL) se ha consolidado como un enfoque crucial para mejorar las capacidades de los modelos de lenguaje grandes. Sin embargo, en los modelos de Mezcla de Expertos (MoE), el mecanismo de enrutamiento a menudo introduce inestabilidad, llegando incluso a provocar un colapso catastrófico del entrenamiento por RL. Analizamos la consistencia entrenamiento-inferencia de los modelos MoE e identificamos una discrepancia notable en los comportamientos de enrutamiento entre ambas fases. Además, incluso en condiciones idénticas, el marco de enrutamiento puede producir selecciones de expertos divergentes en diferentes pasadas hacia adelante. Para abordar esta inconsistencia fundamental, proponemos Rollout Routing Replay (R3), un método que registra las distribuciones de enrutamiento del motor de inferencia y las reproduce durante el entrenamiento. R3 reduce significativamente la divergencia KL de las políticas entre entrenamiento e inferencia y mitiga las discrepancias extremas sin comprometer la velocidad de entrenamiento. Experimentos exhaustivos en varios entornos confirman que R3 logra estabilizar el entrenamiento por RL, previniendo el colapso y superando a métodos como GSPO y TIS. Creemos que este trabajo puede ofrecer una nueva solución para estabilizar el RL en modelos 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.