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

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

• 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