推論ルーターと学習ルーターの整合によるMoE強化学習の安定化

要旨

強化学習（RL）は大規模言語モデルの能力向上における重要な手法として台頭してきた。しかし、Mixture-of-Experts（MoE）モデルでは、ルーティング機構が不安定性を引き起こし、RL訓練の破綻を招くことさえある。本論文では、MoEモデルの訓練と推論の一貫性を分析し、両フェーズ間のルーティング挙動に顕著な不一致が存在することを明らかにする。さらに同一条件下であっても、ルーティングフレームワークは繰り返しのフォワードパスにおいて異なるエキスパート選択を生じ得る。この根本的な不一致を解決するため、我々は推論エンジンからルーティング分布を記録し訓練中に再生するRollout Routing Replay（R3）を提案する。R3は訓練速度を損なうことなく、訓練-推論方策のKLダイバージェンスを大幅に低減し、極端な不一致を緩和する。多様な設定での大規模実験により、R3がRL訓練を安定化させ、破綻を防止し、GSPOやTISなどの手法を上回る性能を発揮することを確認した。本研究がMoEモデルにおけるRL安定化の新たな解決策を提供できると考える。

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.

推論ルーターと学習ルーターの整合によるMoE強化学習の安定化

Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

要旨

Support