PR²: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning

Authors: Daize Dong, Junlin Chen, Haolong Jia, Jiang Liu, Jiawei Wu, Huanwei Di, Jialian Wu, Zhengzhong Liu, Zicheng Liu, Emad Barsoum, Dimitris N. Metaxas, Hongyi Wang

Affiliations: Rutgers University, AMD, MBZUAI

---

Summary (Overview)

• Problem: Reinforcement learning (RL) on Mixture of Experts (MoE) LLMs suffers from training instability caused by router drift—expert activations change drastically across model updates, causing mismatch between rollout and training phases in PPO-style algorithms.
• Limitation of Existing Solutions: Routing replay freezes replay routes within reasoning trajectories to stabilize importance sampling, but ignores how the router evolves under off-policy updates, causing router staleness.
• Proposed Method: PR² augments each router with a lightweight evolution predictor that anticipates short-horizon router evolution, enabling predicted expert indices during rollout that are replayed during training for consistency.
• Key Results: PR² substantially reduces routing mismatch, improves RL stability, and achieves superior performance across multiple reasoning benchmarks. On AIME24, GRPO with PR² achieves 40.31% accuracy on Qwen3-30B-A3B-Base, improving over routing replay and GSPO by 12.29% and 9.38% points respectively (off-8 setting).
• Theoretical Contribution: Derives a divergence-based bound showing that replay staleness controls route-induced gradient deviation, motivating the predictive KL objective.

---

Introduction and Theoretical Foundation

Background

Large Language Models (LLMs) have demonstrated strong reasoning capabilities when scaled through increased parameters and training compute. Mixture of Experts (MoE) models have emerged as a promising architecture enabling training compute to scale sublinearly with model size, activating typically less than 10% of parameters per token. Beyond pretraining, reinforcement learning (RL) has become central for optimizing LLMs to improve reasoning, agentic capabilities, and alignment with human preferences.

The Challenge: Router Drift in MoE RL

A key difficulty for stable RL on MoE-based LLMs stems from learned routers that dynamically assign tokens to experts. When policy updates reuse trajectories from a stale snapshot, MoE models expose this off-policy gap as router drift:

The same token may be routed to different experts by the old snapshot and the current training policy, altering computation paths behind PPO-style importance ratios and amplifying their variance.

Formalization

Off-Policy Objective. Given an old policy snapshot $\pi_{\theta_{\mathrm{old}}}$ generating trajectories and current training policy $\pi_\theta$, the off-policy objective is:

$$J(\theta) = \mathbb{E}_{x, y \sim \pi_{\theta_{\mathrm{old}}}}\left[ w_{\theta}(x, y) R(x, y) \right]$$

with token-level importance ratios:

$$w_{\theta}(x, y) = \prod_{t=1}^{T} r_t(\theta), \quad r_t(\theta) = \frac{\pi_{\theta}\left(y_t \mid x, y_{<t}\right)}{\pi_{\theta_{\mathrm{old}}}\left(y_t \mid x, y_{<t}\right)} \tag{1}$$

Router Staleness. For a replay scheme with route distribution $\rho_t^{\mathrm{rep}}$, token-level router staleness is defined via TV distance:

$$\mathcal{S}_t(\rho_t^{\mathrm{rep}}) = D_{\mathrm{TV}}\bigl(\rho_{\theta,t}^{\pi}(\cdot \mid x, y_{<t}), \rho_t^{\mathrm{rep}}(\cdot \mid x, y_{<t})\bigr)$$

Key Theoretical Result. For any bounded fixed-route gradient kernel $g_t$, replay staleness controls route-induced gradient deviation:

$$\left\| g_t(\rho_{\theta,t}^{\pi}, \theta) - g_t(\rho_t^{\mathrm{rep}}, \theta) \right\| \leq 2M_t \cdot \mathcal{S}_t(\rho_t^{\mathrm{rep}})$$

This reveals a fundamental tension:

• Stable importance estimation favors frozen routing
• Effective learning requires routing distributions that track policy evolution

Related Work

The paper builds on three research threads:

1. MoE architectures with learned routing strategies, load balancing, and gradient estimation
2. PPO-style RL algorithms for LLM alignment (GRPO, GSPO, etc.)
3. Stabilization techniques for MoE RL including routing replay and its variants

---

Methodology

PR²: Predictive Routing Replay

The design principle of PR² is to keep deterministic replay consistency while predicting the expert index likely to become active after short-horizon policy updates.

Route Prediction During Rollout

For each MoE layer $l$ at token $t$, let $h_{\mathrm{old},t}^{(l)}$ be the router input and $p_{\mathrm{old},t}^{(l)} = h_{\mathrm{old},t}^{(l)} W_{\mathrm{old}}^{(l)}$ be the router logits under the old snapshot.

PR² introduces a lightweight evolution predictor $W_p^{(l)}$ (initialized at 0) that outputs an additive logit bias:

$$b_t^{(l)} = h_{\mathrm{old},t}^{(l)} W_p^{(l)}$$

The corrected logits define the predictive routing distribution:

$$\hat{\rho}_t^{(l)} = \mathrm{Softmax}\left(p_{\mathrm{old},t}^{(l)} + b_t^{(l)}\right) \tag{4}$$

Predicted expert indices are selected via top-k:

$$\hat{\mathcal{I}}_t^{(l)} = \mathrm{TopK}\left(\hat{\rho}_t^{(l)}, k\right), \quad \hat{\mathcal{R}}_t = \left\{\hat{\mathcal{I}}_t^{(l)}\right\}_{l=1}^{L} \tag{5}$$

Replay During Training

During training, the current policy computes router logits but reuses the cached expert indices $\hat{\mathcal{I}}_t^{(l)}$, restricting MoE computation:

$$o_t^{(l)} = \sum_{j \in \hat{\mathcal{I}}_t^{(l)}} \rho_{t,j}^{(l)} E_j^{(l)}\left(h_t^{(l)}\right) \tag{7}$$

This freezes only the indices of selected experts while expert outputs and policy gradients always use current parameters $\theta$.

Evolution Predictor Training

The predictor is trained with a KL divergence objective:

$$\mathcal{L}_{\mathrm{Pr}^2} = \sum_{l=1}^{L} \mathbb{E}_t \left[ D_{\mathrm{KL}}\left(\langle \rho_t^{(l)} \rangle \| \hat{\rho}_t^{(l)}\right) \right] \tag{8}$$

where $\langle \cdot \rangle$ denotes stop-gradient. The current router acts as a teacher (stop-gradient), and gradients update only the predictor parameters $\{W_p^{(l)}\}_{l=1}^L$.

Staleness Surrogate Bound: The predictive loss bounds route-induced gradient deviation:

$$\mathbb{E}_t\left[ \left\| g_t(P_t^{\rho}, \theta) - g_t(P_t^{\hat{\rho}}, \theta) \right\| \right] \leq M\sqrt{2\mathcal{L}_{\mathrm{Pr}^2}}$$

Training Details

• Off-policy training begins after the first inner update when the current router has moved from the route-recording snapshot
• Dedicated learning-rate multiplier $\alpha$ for evolution predictors to track router evolution
• No modification to the base PPO objective

Algorithm Pseudocode

Algorithm 1: Predictive Routing Replay (PR²)
1: Route Prediction During Rollout on old snapshot θ_old.
2: for each token t and MoE layer l do
3:   Compute p_old,t^(l) ← h_old,t^(l) · W_old^(l)
4:   Compute b_t^(l) ← h_old,t^(l) · W_p^(l)
5:   Set ρ̂_t^(l) ← Softmax(p_old,t^(l) + b_t^(l))
6:   Select Î_t^(l) ← TopK(ρ̂_t^(l), k) and dispatch to experts
7:   Cache Î_t^(l) and (h_old,t^(l), p_old,t^(l))
8: end for
9: Replay During Training on training model θ.
10: for each inner RL update do
11:   for each token t and layer l do
12:     Compute ρ_t^(l) ← Softmax(p_t^(l))
13:     Replay Î_t^(l) as the MoE expert index
14:   end for
15:   Calculate L_Pr² and update {W_p^(l)}_{l=1}^{L}
16:   Calculate the base RL loss and update θ
17: end for

---

Empirical Validation / Results

Experimental Setup

Models and Data:

• Qwen3-30B-A3B-Base trained on DAPO-17K
• Moonlight-16B-A3B trained on GSM8K
• OLMoE-1B-7B trained on RLVR-GSM

Off-Policy Strength: $\kappa = \frac{B_{\mathrm{global}}}{B_{\mathrm{update}}}$ where larger $\kappa$ = stronger off-policy reuse

Baselines: GRPO, GSPO, GRPO + Routing Replay (R²)

Main Results: Qwen3-30B-A3B-Base

Table 1: Downstream reasoning accuracy on Qwen3-30B-A3B-Base. PR² achieves the best average accuracy in all settings, with gains of 9.25%, 13.00%, and 8.16% points over routing replay under off-2, off-4, and off-8 respectively.

Cross-Model Evaluation

On Moonlight-16B-A3B, PR² improves over routing replay by:

• +1.19% (off-2)
• +1.19% (off-4)
• +3.44% (off-8)

The largest gains appear under the strongest rollout reuse, consistent with the router-staleness view. On OLMoE-1B-7B, PR² achieves the best average accuracy across all off-policy strengths.

Training Stability

PR² exhibits:

• Lower clipping rates and reduced policy-gradient loss spikes
• Smoother reward growth without abrupt entropy collapse
• Steadier response length increase, indicating more stable optimization

Routing Prediction Analysis

Key Findings:

• PR² maintains higher top-k agreement and lower route KL across all settings
• Zero-deviation ratio increases from 76.9% (GRPO) and 88.8% (routing replay) to 91.5–92.2% with PR²
• 1-slot mismatch decreases from 19.6% (GRPO) and 11.0% (routing replay) to 7.7–8.5%
• Cases with ≥2 mismatched slots reduced to only 0.1%

---

Theoretical and Practical Implications

Theoretical Implications

1. Formalization of router staleness: The paper provides the first formal treatment of router staleness as a bound on route-induced gradient deviation, showing that replay staleness directly controls the quality of fixed-route PPO gradients.

2. The stability-learning trade-off: The work articulates a fundamental tension: stable importance estimation favors frozen routing, while effective learning requires routing distributions that track policy evolution. This provides a principled framework for understanding MoE RL instability.

3. Staleness-controlled predictive loss: The KL divergence objective is theoretically motivated by the gradient-deviation bound, offering a principled training signal for the evolution predictor.

Practical Implications

1. Drop-in integration: PR² can be integrated into existing MoE RL frameworks (VeRL, Slime, TRL) without modifying the policy optimization objective, making it practical for real-world deployment.

2. Scalability: The evolution predictor adds negligible computational overhead while enabling more stable training, particularly valuable for large-scale MoE models where training instability is costly.

3. Robustness across settings: PR² consistently improves performance across different model scales (1B–30B), architectures, and off-policy strengths, suggesting broad applicability.

---

Conclusion

This paper identifies router drift as a key source of instability in off-policy RL for MoE-based LLMs and introduces Predictive Routing Replay (PR²) to address it. By predicting short-horizon expert indices while preserving deterministic replay, PR² resolves the central tension between stable importance estimation and effective learning.

Key contributions:

1. Formalization of router staleness with a divergence-based bound on route-induced gradient deviation
2. A lightweight evolution predictor trained with a theoretically motivated KL objective
3. Comprehensive empirical validation showing significant improvements in reasoning accuracy, training stability, and route tracking across multiple MoE backbones

Future directions include extending PR² to other RL paradigms, exploring adaptive prediction horizons, and investigating the interaction between router prediction and other stability mechanisms in large-scale MoE training.