# Router Drift Control: Mitigating Off-Policy Degradation in Mixture-of-Experts Reinforcement Learning

> Router Drift Control filters high-drift responses in MoE reinforcement learning, improving AIME accuracy by 5.5 points over TIS and 4.5 over R3 under asynchronous training.

- **Source:** [arXiv](https://arxiv.org/abs/2610.21682)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/KtcGiT
- **Whiteboard:** https://picx.dev/p/KtcGiT/image

## Summary

# Summary of "Router Drift Control: Mitigating Off-Policy Degradation in Mixture-of-Experts Reinforcement Learning"

## Summary (Overview)

- **Problem**: In Mixture-of-Experts (MoE) reinforcement learning, off-policy mismatch (from training-inference discrepancies, mini-batch updates, and asynchronous training) causes **router drift**—where the same response activates different experts during rollout versus training, altering computation paths and gradient attribution.

- **Key Insight**: Token probabilities alone cannot identify routing decisions. Two different routing assignments can produce identical output distributions while directing gradients to completely different expert parameter blocks.

- **Proposed Method**: **Router Drift Control (RDC)**—a response-level filtering approach that quantifies rollout–training routing disagreement and selectively excludes high-drift responses from policy updates, while preserving native training routes for retained responses.

- **Key Results**: RDC combined with Truncated Importance Sampling (TIS) improves average peak validation accuracy on AIME 2024 and AIME 2025 by **5.5 points over TIS** and **4.5 points over R3** in asynchronous training, with gains increasing as policy staleness grows.

- **Contribution**: Introduces routing-aware response filtering as an alternative to route replay (R3), avoiding the gradient bias introduced by replay while maintaining stable training dynamics.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Reinforcement learning with verifiable rewards is widely used to improve reasoning in large language models. In practical systems, rollout generation and policy optimization are handled by **separate inference and training engines**, creating three sources of off-policy mismatch:

1. **Training-Inference Mismatch**: Different numerical precision, computational kernels, and parallel strategies produce different token probabilities even with identical parameters.

2. **Mini-Batch Updates**: Large rollout batches are partitioned into mini-batches for sequential updates—later mini-batches become increasingly stale relative to the current policy.

3. **Asynchronous Training**: Decoupling rollout and optimization for throughput introduces policy staleness.

### The Role of Routing in MoE Models

MoE architectures activate only a subset of experts per token via a router. Routing decisions directly determine:
- The computation path for each token
- Which experts receive learning signals

Under off-policy conditions, small policy differences can change routing decisions, causing the same token to be assigned to different experts during rollout and training—this is **router drift**.

### Key Theoretical Insight: Routing Determines Gradient Support

At layer $l$ and token $t$, an MoE layer computes:

$$
h_{l}^{\mathrm{out}}(t) = \sum_{e \in R_{l}(t)} \mathrm{gate}_{l,e}(t) E_{l,e}\big(h_{l}^{\mathrm{in}}(t)\big) \tag{7}
$$

where $E_{l,e}$ denotes expert $e$ and $\mathrm{gate}_{l,e}(t)$ is its routing weight. Critically, if the router does not select an expert, the token provides **no direct gradient** to that expert's parameters:

$$
e \notin R_{l}(t) \quad \Longrightarrow \quad \frac{\partial h_{l}^{\mathrm{out}}(t)}{\partial \theta_{l,e}} = 0 \tag{8}
$$

### Why Token Probabilities Cannot Identify Routing

The paper provides a formal argument: consider a top-K MoE layer where $S$ contains $K-1$ shared experts, and $e_1$, $e_2$ are two distinct experts outside $S$. Consider two routing assignments:

$$
R_{1} = S \cup \{e_{1}\}, \qquad R_{2} = S \cup \{e_{2}\} \tag{9}
$$

If the shared experts have identical routing weights under both assignments and the replacing experts produce identical outputs:

$$
E_{l,e_{1}}\big(h_{l}^{\mathrm{in}}(t)\big) = E_{l,e_{2}}\big(h_{l}^{\mathrm{in}}(t)\big) \tag{10}
$$

Then both assignments produce identical logits and probability distributions, yet:

$$
\frac{\partial h_{l}^{\mathrm{out}}(t; R_{1})}{\partial \theta_{l,e_{2}}} = 0, \quad \frac{\partial h_{l}^{\mathrm{out}}(t; R_{2})}{\partial \theta_{l,e_{1}}} = 0 \tag{11}
$$

**Conclusion**: Likelihood ratios and token-level KL divergence cannot identify which routing assignment was used, but the gradient support differs fundamentally.

---

## Methodology

### Formal Setup

- Let $\mu(\cdot \mid x)$ be the **behavior policy** (rollout engine) and $\pi_{\theta}(\cdot \mid x)$ the **training policy**.
- For a response $y = (y_1, \ldots, y_T)$ with token positions $\mathcal{T} = \{1, \ldots, T\}$:
  - MoE model with $L$ layers, $E$ experts per layer, $K$ experts routed per token
  - Expert sets selected during rollout and training:

$$
R_{l}^{\mu}(t), R_{l}^{\pi}(t) \subseteq \{1, \dots, E\}, \quad |R_{l}^{\mu}(t)| = |R_{l}^{\pi}(t)| = K \tag{1}
$$

### Response-Level Router Drift Score

For engine $a \in \{\mu, \pi\}$, the expert-usage distribution at layer $l$ over response $y$ is:

$$
u_{l}^{a}(e \mid y) = \frac{1}{TK} \sum_{t=1}^{T} \mathbb{I}[e \in R_{l}^{a}(t)], \qquad e \in \{1, \ldots, E\} \tag{2}
$$

The routing difference at each layer uses **total variation distance**:

$$
d_{l}(y) = \frac{1}{2} \left\| u_{l}^{\mu}(\cdot | y) - u_{l}^{\pi}(\cdot | y) \right\|_{1} \tag{3}
$$

The response-level router drift score averages layer-wise distances:

$$
D_{\mathrm{route}}(y) = \frac{1}{L} \sum_{l=1}^{L} d_{l}(y) \tag{4}
$$

The score lies in $[0, 1]$, with larger values indicating greater mismatch in aggregate expert usage.

### Selective Response Filtering

Let $q \in [0, 1)$ be the filtering ratio. RDC ranks responses by $D_{\mathrm{route}}(y)$ **within groups of similar response lengths** (length buckets) and filters the highest-drift fraction. The response-level mask is:

$$
m_{\mathcal{B}}(y) = \mathbb{I}\left[ y \notin \mathcal{F}_{q}(\mathcal{B}) \right] \tag{5}
$$

The RDC objective:

$$
\mathcal{L}_{\mathrm{RDC}}(\theta) = \frac{1}{N_{\mathcal{B}}} \sum_{y \in \mathcal{B}} m_{\mathcal{B}}(y) \ell_{\theta}(y; R^{\pi}(y)) \tag{6}
$$

where $\ell_{\theta}(y; R)$ is the sum of token-level policy losses for response $y$ under routing assignment $R$, and $N_{\mathcal{B}}$ is the total number of valid tokens in retained responses.

### Comparison with Route Replay

For route substitution, the gradient contribution of response $y$ under routing assignment $R$ is:

$$
g_{\theta}(y; R) = \nabla_{\theta} \ell_{\theta}(y; R) \tag{12}
$$

Replacing the training route with the rollout route changes the gradient by:

$$
\Delta_{\mathrm{route}}(y) = g_{\theta}(y; R^{\mu}(y)) - g_{\theta}(y; R^{\pi}(y)) \tag{13}
$$

This substitution can alter expert-level gradient support and introduce bias relative to the native routing objective. RDC avoids this by **retaining native training routes** and using routing disagreement only to determine which responses contribute.

### Algorithm

```
Algorithm 1: Router Drift Control
Require: Batch B, rollout routes R^μ, training policy π_θ, ratio q
1: for all y ∈ B do
2:   Obtain training routes R^π(y)
3:   Compute D_route(y) using Eqs. 2–4
4: end for
5: Partition B into length buckets
6: for all buckets B_b do
7:   Select the top ⌈q|B_b|⌉ responses
8: end for
9: Construct m_B using Eq. 5
10: Compute L_RDC(θ) using Eq. 6
11: Update θ using ∇_θ L_RDC(θ)
```

---

## Empirical Validation / Results

### Experimental Setup

- **Model**: Qwen3-30B-A3B (MoE model)
- **Training Data**: DAPO-Math-17K
- **Evaluation**: AIME 2024, AIME 2025, MATH-500, ZebraLogic
- **Baseline**: GRPO with DAPO Clip-Higher objective ($\epsilon_{\mathrm{low}} = 0.2$, $\epsilon_{\mathrm{high}} = 0.27$)
- **Methods Compared**: GRPO, R3, TIS, RDC+TIS
- **Filtering Ratio**: $q = 0.08$ (default)
- **Infrastructure**: Megatron (BF16 training), vLLM (FP8 rollout)

### Synchronous Training Results

**Table 1: Best validation accuracy under different update batch sizes**

| $(\mathcal{B}_{\text{rollout}}, \mathcal{B}_{\text{update}})$ | Method | ZebraLogic | AIME24 | AIME25 | MATH-500 |
|---|---|---|---|---|---|
| (2048, 2048) | GRPO | 61.2 | 77.7 | 65.2 | 96.6 |
| | GRPO+R3 | 63.0 | 74.9 | 65.4 | 97.6 |
| | GRPO+TIS | 66.1 | 80.0 | 71.3 | 98.2 |
| | **GRPO+TIS+RDC** | **73.9** | 78.5 | **72.4** | 97.8 |
| (2048, 512) | GRPO | 46.3 | 62.2 | 50.8 | 96.0 |
| | GRPO+R3 | 55.2 | 69.9 | 57.9 | 96.6 |
| | GRPO+TIS | 61.3 | 70.6 | 60.3 | 97.0 |
| | **GRPO+TIS+RDC** | **68.1** | **71.5** | **62.5** | **97.6** |
| (2048, 128) | GRPO | 61.1 | 53.4 | 40.0 | 92.2 |
| | GRPO+R3 | 53.1 | 61.0 | 48.7 | 95.0 |
| | GRPO+TIS | 78.2 | 70.6 | 62.9 | 97.6 |
| | **GRPO+TIS+RDC** | **86.1** | **77.0** | **67.2** | **97.8** |

**Key finding**: RDC's advantage becomes more pronounced as the update batch shrinks (increasing staleness). Under (2048, 128), RDC improves ZebraLogic, AIME24, and AIME25 by **7.9, 6.4, and 4.3 percentage points** over TIS alone.

### Asynchronous Training Results

**Table 2: Best validation accuracy under asynchronous RL**

| Method | ZebraLogic | AIME24 | AIME25 | MATH-500 |
|---|---|---|---|---|
| GRPO | 46.9 | 63.4 | 50.7 | 96.2 |
| GRPO+R3 | 53.9 | 68.7 | 53.1 | 96.2 |
| GRPO+TIS | 53.7 | 65.0 | 54.8 | 96.4 |
| **GRPO+TIS+RDC** | 53.0 | **71.2** | **59.6** | **97.2** |

- RDC achieves the highest accuracy on AIME24, AIME25, and MATH-500
- Average peak improvement: **+5.5 points over TIS** and **+4.5 points over R3** on AIME benchmarks
- R3 exhibits late-stage degradation; RDC maintains strong late-stage performance

### Training Dynamics

The uncorrected GRPO baseline collapses after ~100 training steps (reward becomes negative, response length drops from ~6,000 to ~2,000 tokens). This coincides with:
- Rapid growth in training–inference KL divergence
- Sharply rising gradient norms

TIS and RDC keep both KL and gradient norm **low and stable** throughout training, while R3 shows more pronounced increases in later stages.

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Routing as a distinct off-policy signal**: The paper formally demonstrates that token probabilities do not uniquely determine routing decisions, making routing disagreement an independent source of off-policy information that probability-based corrections cannot capture.

2. **Gradient bias of route replay**: The analysis shows that replaying rollout routes (R3) alters gradient attribution relative to native routing, potentially omitting contributions from experts the current training router would select.

3. **Routing-aware filtering as an alternative**: RDC provides a principled middle ground—it uses routing information for *selection* (which responses to train on) rather than *modification* (changing computation paths), avoiding replay-induced bias.

### Practical Implications

1. **Scalability to asynchronous RL**: As RL systems increasingly adopt asynchronous training for throughput, RDC provides a robust correction mechanism that maintains stability and performance under significant policy lag.

2. **Complementary corrections**: RDC combines naturally with TIS—they address orthogonal aspects of off-policy mismatch (routing disagreement vs. token-level probability mismatch).

3. **Implementation simplicity**: RDC requires no changes to the training computation graph—only a filtering step before optimization, making it easy to integrate into existing RL pipelines.

4. **Length-bucket-aware filtering**: Ranking within length groups prevents systematic bias against longer responses, which naturally exhibit more routing variation.

---

## Conclusion

### Main Takeaways

- **Router drift** is a distinct and consequential source of off-policy degradation in MoE reinforcement learning, affecting expert-level gradient support in ways that token probabilities cannot reveal.

- **RDC** (Router Drift Control) selectively filters high-drift responses based on response-level routing disagreement while preserving native training routes for retained responses, avoiding the gradient bias of route replay (R3).

- RDC's gains **increase with policy staleness**, making it particularly valuable for mini-batch updates and fully asynchronous training regimes.

- Combined with TIS, RDC achieves substantial improvements over both TIS alone and R3 on mathematical reasoning benchmarks while maintaining stable training dynamics.

### Future Directions

- Investigating adaptive filtering ratios that respond to the degree of policy staleness over time
- Exploring finer-grained filtering (e.g., token-level or layer-level drift signals) beyond response-level selection
- Applying RDC to other MoE architectures and domains beyond mathematical reasoning
- Studying the interaction between routing-aware filtering and other off-policy correction mechanisms

---

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