PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning

• URL: http://arxiv.org/abs/2606.00395v2
• Date: Tue, 02 Jun 2026 03:28:31 GMT
• Title: PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning
• Abstract summary: 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.
• Score: 40.05865857566415
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Mixture of Experts (MoE) Large Language Models (LLMs) achieve strong performance at scale. However, reinforcement learning (RL) on MoE-based LLMs often suffers from training instability. A root cause is router drift, i.e., expert activations can change drastically across model updates and differ between disaggregated rollout and training phases, causing large rollout--training mismatch and unstable importance sampling weights in PPO-style RL algorithms. Routing replay mitigates this issue by freezing the replay route within each reasoning trajectory, but it ignores how the router evolves under off-policy updates and thus causes router staleness. To address this limitation, we propose Predictive Routing Replay (PR2), which augments each router with a lightweight evolution predictor that learns to anticipate short-horizon router evolution. During the rollout phase, we use the predictive routing distribution to apply top-$k$ routing, enabling gradients to reach experts that are likely to become active after updates. During the training phase, we replay the resulting predicted route to retain consistency for stable importance estimation. Theoretical analysis and experiments support that PR2 reduces routing-induced mismatch, improves RL stability, and yields stronger performance across various reasoning benchmarks.

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)
• Dynamic Rollout Editing for Reducing Overthinking in RL-Trained Reasoning Models [102.76983747945836]
  Long-form chain-of-thought reasoning can improve performance on complex tasks.<n>But models often continue generating unnecessary reasoning after a correct answer has emerged.<n>We study this phenomenon from the perspective of GRPO-style reinforcement learning.
  arXiv  Detail & Related papers  (2026-06-16T13:10:30Z)
• SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models [51.99912169291891]
  Reinforcement learning (RL) has emerged as a promising paradigm for training reasoning-oriented models.<n>We propose a training paradigm that interleaves conventional RL with inverse reinforcement learning (IRL)<n>Our findings suggest that alternating between RL and IRL offers an effective pathway toward extending the exploration capacity of reasoning-oriented large language models.
  arXiv  Detail & Related papers  (2026-04-18T13:49:47Z)
• ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning [85.39146836671897]
  Low-rank adapters (LoRAs) are a parameter-efficient finetuning technique that injects trainable low-rank matrices into pretrained models to adapt them to new tasks.<n>Existing Mixture-of-LoRAs routers assign a learned routing weight to each LoRA to enable end-to-end training of the router.<n>We propose a new router designed that we call Reinforcement Routing for Mixture-of-LoRAs (ReMix)
  arXiv  Detail & Related papers  (2026-03-10T18:51:27Z)
• Dr.LLM: Dynamic Layer Routing in LLMs [55.11953638340419]
  Dr.LLM is a retrofittable framework that equips pretrained models with lightweight per-layer routers deciding to skip, execute, or repeat a block.<n>On ARC (logic) and DART (math), Dr.LLM improves accuracy by up to +3.4%p while saving 5 layers per example on average.
  arXiv  Detail & Related papers  (2025-10-14T17:51:26Z)
• Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers [33.18639109838351]
  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.
  arXiv  Detail & Related papers  (2025-10-13T13:11:27Z)
• ProxRouter: Proximity-Weighted LLM Query Routing for Improved Robustness to Outliers [14.831117443453165]
  Large language model (LLM) query routers are critical to modern AI platforms.<n>We propose Prox, which applies an exponentially tilted aggregation mechanism to balance bias and variance in nonparametric routers.
  arXiv  Detail & Related papers  (2025-10-10T20:28:14Z)
• Shuffle-R1: Efficient RL framework for Multimodal Large Language Models via Data-centric Dynamic Shuffle [65.14124923451077]
  Reinforcement learning (RL) has emerged as an effective post-training paradigm for enhancing the reasoning capabilities of multimodal large language model (MLLM)<n>However, current RL pipelines often suffer from training inefficiencies caused by two underexplored issues: Advantage Collapsing and Rollout Silencing.<n>We propose Shuffle-R1, a simple yet principled framework that improves RL fine-tuning efficiency by dynamically restructuring trajectory sampling and batch composition.
  arXiv  Detail & Related papers  (2025-08-07T17:53:47Z)
• DRL4Route: A Deep Reinforcement Learning Framework for Pick-up and Delivery Route Prediction [21.335721424944257]
  We present the first attempt to generalize Reinforcement Learning (RL) to the route prediction task, leading to a novel RL-based framework called DRL4Route. DRL4Route can serve as a plug-and-play component to boost the existing deep learning models. It follows the actor-critic architecture which is equipped with a Generalized Advantage Estimator.
  arXiv  Detail & Related papers  (2023-07-30T14:50:31Z)

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.