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

> PR² introduces a lightweight evolution predictor to anticipate router drift in MoE LLMs, reducing route-induced gradient deviation and improving reasoning accuracy by up to 13% over existing replay methods.

- **Source:** [arXiv](https://arxiv.org/abs/2606.00395)
- **Published:** 2026-08-18
- **Permalink:** https://picx.dev/p/b8iqVt
- **Whiteboard:** https://picx.dev/p/b8iqVt/image

## Summary

# 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

| Policy | Method | AIME24 (Avg@32) | AIME25 (Avg@32) | AMC23 (Avg@16) | HMMT25 (Avg@16) | Average |
|--------|--------|-----------------|-----------------|----------------|-----------------|---------|
| **Off-2** | GRPO | 31.04 | 24.68 | 74.06 | 9.16 | 34.74 |
| | GSPO | 30.42 | 23.54 | 77.18 | 11.88 | 35.76 |
| | GRPO + R² | 35.73 | 25.73 | 77.50 | 11.25 | 37.55 |
| | **GRPO + PR²** | **47.71** | **32.81** | **87.50** | **19.17** | **46.80** |
| **Off-4** | GRPO | 25.42 | 17.71 | 70.93 | 8.33 | 30.60 |
| | GSPO | 32.50 | 22.08 | 79.38 | 12.08 | 36.51 |
| | GRPO + R² | 28.13 | 21.25 | 73.59 | 10.00 | 33.24 |
| | **GRPO + PR²** | **47.40** | **31.67** | **85.47** | **20.42** | **46.24** |
| **Off-8** | GRPO | 25.00 | 15.93 | 72.81 | 7.50 | 30.31 |
| | GSPO | 30.93 | 22.18 | 72.03 | 7.50 | 33.16 |
| | GRPO + R² | 28.02 | 21.77 | 76.88 | 8.30 | 33.74 |
| | **GRPO + PR²** | **40.31** | **28.54** | **83.13** | **15.63** | **41.90** |

**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.

---

_Markdown view of https://picx.dev/p/b8iqVt, served by PicX — AI-generated visual whiteboard summaries of research papers._