훈련 및 추론 라우터 정렬을 통한 MoE 강화 학습 안정화

초록

강화학습(RL)은 대규모 언어 모델의 능력을 향상시키는 중요한 접근법으로 부상했습니다. 그러나 전문가 혼합(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