Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

• URL: http://arxiv.org/abs/2510.11370v2
• Date: Tue, 21 Oct 2025 17:19:46 GMT
• Title: Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers
• Abstract summary: We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors.<n>We propose Rollout Replay Routing (R3), a method that records routing distributions from the inference engine and replays them during training.<n>R3 significantly reduces training-inference policy KL divergence and mitigates extreme discrepancies without compromising training speed.
• Score: 33.18639109838351
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Expert-Space Exploration in MoE Reinforcement Learning [41.1632714379556]
  Expert-Space Exploration Reinforcement Learning (ESRL) is an architecture-aware framework that explicitly explores the expert-routing space of MoE models.<n>ESRL achieves the best performance across MoE backbones with top-K, top-1, and shared-expert routing, as well as across mathematics, science, and code tasks without additional sampling or computational cost.
  arXiv  Detail & Related papers  (2026-09-11T16:58:41Z)
• Enough is as good as a feast: A Comprehensive Analysis of How Reinforcement Learning Mitigates Task Conflicts in LLMs [52.129885179495226]
  This study systematically explores the merging behavior of RL-trained large language models (LLMs) compared to those trained with traditional SFT.<n>We find that RL significantly reduces task conflicts and results in less performance degradation after merging.
  arXiv  Detail & Related papers  (2026-07-24T07:08:57Z)
• When Model Merging Breaks Routing: Training-Free Calibration for MoE [67.27839627141972]
  Hessian-Aware Router (HARC) is a training-free framework that leverages second-order curvature information to realign the merged router.<n> Experiments on mathematical reasoning and code generation tasks show that HARC effectively mitigates routing breakdown across diverse MoE merging baselines.
  arXiv  Detail & Related papers  (2026-06-02T09:33:33Z)
• PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning [40.05865857566415]
  We propose Predictive Routing Replay (PR2), which augments each router with a lightweight evolution predictor that learns to anticipate short-horizon router evolution.<n>PR2 reduces routing-induced mismatch, improves RL stability, and yields stronger performance across various reasoning benchmarks.
  arXiv  Detail & Related papers  (2026-05-29T22:28:08Z)
• On the Interplay of Pre-Training, Mid-Training, and RL on Reasoning Language Models [73.10315509190623]
  Recent reinforcement learning techniques have yielded impressive reasoning improvements in language models.<n>It remains unclear whether post-training truly extends a model's reasoning ability beyond what it acquires during pre-training.<n>We develop a fully controlled experimental framework that isolates the causal contributions of pre-training, mid-training, and RL-based post-training.
  arXiv  Detail & Related papers  (2025-12-08T18:12:10Z)
• Stabilizing Reinforcement Learning with LLMs: Formulation and Practices [61.361819972410046]
  We show why and under what conditions the true sequence-level reward can be optimized via a surrogate token-level objective in policy gradient methods such as REINFORCE.<n>This insight provides a principled explanation for the crucial role of several widely adopted techniques in stabilizing RL training.
  arXiv  Detail & Related papers  (2025-12-01T07:45:39Z)
• Towards Stable and Effective Reinforcement Learning for Mixture-of-Experts [113.0656076371565]
  We propose a novel router-aware approach to optimize importance sampling weights in off-policy reinforcement learning (RL)<n>Specifically, we design a rescaling strategy guided by router logits, which effectively reduces gradient variance and mitigates training divergence.<n> Experimental results demonstrate that our method significantly improves both the convergence stability and the final performance of MoE models.
  arXiv  Detail & Related papers  (2025-10-27T05:47:48Z)
• Adversarial Diffusion for Robust Reinforcement Learning [46.44328012099217]
  We introduce Adversarial Diffusion for Robust Reinforcement Learning (AD-RRL)<n>AD-RRL guides the diffusion process to generate worst-case trajectories during training, effectively optimizing the Conditional Value at Risk (CVaR) of the cumulative return.<n> Empirical results across standard benchmarks show that AD-RRL achieves superior robustness and performance compared to existing robust RL methods.
  arXiv  Detail & Related papers  (2025-09-28T12:34:35Z)
• Balanced Actor Initialization: Stable RLHF Training of Distillation-Based Reasoning Models [27.0496567592082]
  The development of alignment and reasoning capabilities in large language models has seen remarkable progress.<n>The third paradigm of applying RLHF to distillation-trained models presents significant challenges.<n>We propose Balanced Actor Initialization (BAI), a two-stage weighted model merging approach.
  arXiv  Detail & Related papers  (2025-08-30T01:53:25Z)
• Ring-lite: Scalable Reasoning via C3PO-Stabilized Reinforcement Learning for LLMs [51.21041884010009]
  Ring-lite is a Mixture-of-Experts (MoE)-based large language model optimized via reinforcement learning (RL)<n>Our approach matches the performance of state-of-the-art (SOTA) small-scale reasoning models on challenging benchmarks.
  arXiv  Detail & Related papers  (2025-06-17T17:12:34Z)
• RAGEN: Understanding Self-Evolution in LLM Agents via Multi-Turn Reinforcement Learning [125.96848846966087]
  Training large language models (LLMs) as interactive agents presents unique challenges.<n>While reinforcement learning has enabled progress in static tasks, multi-turn agent RL training remains underexplored.<n>We propose StarPO, a general framework for trajectory-level agent RL, and introduce RAGEN, a modular system for training and evaluating LLM agents.
  arXiv  Detail & Related papers  (2025-04-24T17:57:08Z)
• Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining [74.83412846804977]
  Reinforcement learning (RL)-based fine-tuning has become a crucial step in post-training language models.<n>We present a systematic end-to-end study of RL fine-tuning for mathematical reasoning by training models entirely from scratch.
  arXiv  Detail & Related papers  (2025-04-10T17:15:53Z)
• Avoiding mode collapse in diffusion models fine-tuned with reinforcement learning [0.0]
  Fine-tuning foundation models via reinforcement learning (RL) has proven promising for aligning to downstream objectives. We exploit the hierarchical nature of diffusion models (DMs) and train them dynamically at each epoch with a tailored RL method. We show that models trained with HRF achieve better preservation of diversity in downstream tasks, thus enhancing the fine-tuning robustness and at uncompromising mean rewards.
  arXiv  Detail & Related papers  (2024-10-10T19:06:23Z)
• Symmetric Reinforcement Learning Loss for Robust Learning on Diverse Tasks and Model Scales [13.818149654692863]
  Reinforcement learning (RL) training is inherently unstable due to factors such as moving targets and high gradient variance.<n>In this work, we improve the stability of RL training by adapting the reverse cross entropy (RCE) from supervised learning for noisy data to define a symmetric RL loss.
  arXiv  Detail & Related papers  (2024-05-27T19:28:33Z)
• Robust Reinforcement Learning as a Stackelberg Game via Adaptively-Regularized Adversarial Training [43.97565851415018]
  Robust Reinforcement Learning (RL) focuses on improving performances under model errors or adversarial attacks. Most of the existing literature models RARL as a zero-sum simultaneous game with Nash equilibrium as the solution concept. We introduce a novel hierarchical formulation of robust RL - a general-sum Stackelberg game model called RRL-Stack.
  arXiv  Detail & Related papers  (2022-02-19T03:44:05Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.