Stabilizzazione dell'Apprendimento per Rinforzo MoE mediante Allineamento dei Router di Addestramento e Inferenza

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

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

cs.AI

Abstract

L'apprendimento per rinforzo (RL) è emerso come un approccio cruciale per potenziare le capacità dei grandi modelli linguistici. Tuttavia, nei modelli a miscela di esperti (MoE), il meccanismo di instradamento spesso introduce instabilità, arrivando persino a causare un collasso catastrofico dell'addestramento RL. Analizziamo la coerenza addestramento-inferenza dei modelli MoE e identifichiamo una discrepanza significativa nei comportamenti di instradamento tra le due fasi. Inoltre, anche in condizioni identiche, il framework di instradamento può produrre selezioni di esperti divergenti in passaggi in avanti ripetuti. Per affrontare questa incoerenza fondamentale, proponiamo Rollout Routing Replay (R³), un metodo che registra le distribuzioni di instradamento dal motore di inferenza e le riproduce durante l'addestramento. R³ riduce significativamente la divergenza KL delle politiche tra addestramento e inferenza e mitiga discrepanze estreme senza compromettere la velocità di addestramento. Esperimenti estesi su vari contesti confermano che R³ riesce a stabilizzare l'addestramento RL, prevenendo il collasso e superando metodi come GSPO e TIS. Riteniamo che questo lavoro possa offrire una nuova soluzione per stabilizzare l'RL nei modelli 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.