Стабилизация обучения с подкреплением для MoE за счет согласования маршрутизаторов обучения и вывода

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

October 13, 2025

Авторы: Wenhan Ma, Hailin Zhang, Liang Zhao, Yifan Song, Yudong Wang, Zhifang Sui, Fuli Luo

cs.AI

Аннотация

Обучение с подкреплением (RL) стало ключевым подходом для расширения возможностей больших языковых моделей. Однако в моделях со смесью экспертов (MoE) механизм маршрутизации часто вызывает нестабильность, иногда приводя к катастрофическому коллапсу RL-обучения. Мы анализируем согласованность обучения и вывода в MoE-моделях и выявляем значительное расхождение в поведении маршрутизации между этими двумя фазами. Более того, даже в идентичных условиях фреймворк маршрутизации может давать различные выборы экспертов при повторных прямых проходах. Для решения этой фундаментальной несогласованности мы предлагаем метод Rollout Routing Replay (R3), который записывает распределения маршрутизации из механизма вывода и воспроизводит их во время обучения. R3 существенно сокращает KL-дивергенцию политик между обучением и выводом и смягчает экстремальные расхождения без ущерба для скорости обучения. Многочисленные эксперименты в различных условиях подтверждают, что R3 успешно стабилизирует RL-обучение, предотвращает коллапс и превосходит такие методы, как GSPO и TIS. Мы полагаем, что данная работа может предложить новое решение для стабилизации RL в 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.