Estabilizando o Aprendizado por Reforço MoE através do Alinhamento de Roteadores de Treinamento e Inferência

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

Resumo

O aprendizado por reforço (RL) emergiu como uma abordagem crucial para aprimorar as capacidades de grandes modelos de linguagem. No entanto, em modelos de Mistura de Especialistas (MoE), o mecanismo de roteamento frequentemente introduz instabilidade, podendo até levar a um colapso catastrófico do treinamento por RL. Analisamos a consistência entre treinamento e inferência em modelos MoE e identificamos uma discrepância notável nos comportamentos de roteamento entre as duas fases. Além disso, mesmo sob condições idênticas, a estrutura de roteamento pode produzir seleções divergentes de especialistas em passes *forward* repetidos. Para resolver essa inconsistência fundamental, propomos o Rollout Routing Replay (R3), um método que registra as distribuições de roteamento do motor de inferência e as reproduz durante o treinamento. O R3 reduz significativamente a divergência KL da política entre treinamento e inferência e mitiga discrepâncias extremas sem comprometer a velocidade de treinamento. Experimentos extensos em várias configurações confirmam que o R3 consegue estabilizar o treinamento por RL, prevenindo colapso e superando métodos como GSPO e TIS. Acreditamos que este trabalho pode oferecer uma nova solução para estabilizar o RL em 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.