PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Liangyao Jin, Zhixin Wang, Liming Liu, Mingjun Zhang, Zhengzhou Zhu, Fangcheng Fu, Jinlong Hou, Yuan Cheng, Xiaohe Hu, Ping Zhang, Siyuan Feng

arXiv:2610.21682v1cs.LGcs.AI

TL;DR

Off-policy mismatch can destabilize MoE reinforcement learning, where rollout–training routing differences alter expert-level gradient support. RDC filters high-drift responses while preserving native training routes; with TIS, it improves average peak AIME validation accuracy over TIS and R3 in asynchronous training.

  • Problem

    Off-policy mismatch can degrade and destabilize MoE reinforcement learning, where routing changes alter expert-level gradient support.

    Source: §1 · p. 1, §3.3 · p. 6, §5 · p. 9

  • Method

    RDC measures rollout–training routing disagreement and filters high-drift responses while allowing the training router to select experts for retained responses.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    5.5 points over TIS and 4.5 points over R3 in average peak validation accuracy on AIME24 and AIME25 with RDC combined with TIS during asynchronous training.

    Source: §1 · p. 3, §4.2 · p. 9

  • Takeaways & Limitations

    RDC offers routing-aware filtering as an alternative to route replay, improving validation performance and training stability under mini-batch updates and fully asynchronous training.

    Source: §5 · p. 9

    The setup considers separate rollout and training engines, while the routing-identifiability analysis assumes equal routing weights and identical outputs for the replacing experts.

    Source: §3.1 · p. 4, §3.4 · p. 6

Abstract

from arXiv · show

Reinforcement learning (RL) can become off-policy due to training-inference mismatch, mini-batch updates, and asynchronous training.The resulting mismatch is especially harmful in Mixture-of-Experts (MoE) models, where rollout--training discrepancies can change routing decisions and lead to router drift. Rollout Routing Replay (R3) mitigates this mismatch by reusing expert selections from the rollout stage during training. Our analysis shows that routing replay introduces gradient bias by losing contributions from experts selected by the training router but not the rollout router. To address this, we propose Router Drift Control (RDC), which uses routing disagreement between rollout and training as a signal of off-policy staleness and selectively rejects high-drift responses. The sequence-level filtering avoids gradient bias introduced by replay while preserving native expert selection and the associated learning signals for retained responses. Experiments show that RDC achieves larger gains as policy staleness increases and remains effective in asynchronous RL. In asynchronous training, RDC combined with Truncated Importance Sampling (TIS) improves the average peak validation accuracy across AIME 2024 and AIME 2025 by 5.5 points over TIS and 4.5 points over R3.

1 INTRODUCTION

Rollout–training mismatch can make reinforcement learning effectively off-policy, and in MoE models it can change which experts receive learning signals. The introduction presents Router Drift Control (RDC) as a way to filter high-drift responses while preserving native training routes.

1 INTRODUCTION

  • Training–inference differences, sequential mini-batch updates, and asynchronous policy lag can make the generation policy diverge from the optimization policy.Later mini-batches become increasingly stale after policy updates, while asynchronous training may use trajectories from earlier policy versions.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • In MoE models, rollout–training routing differences can send the same token to different experts, changing the response’s computation path and expert-level learning signals.The router determines which experts participate in computation and receive learning signals.Source: §1 · p. 2, Figure 2 · p. 3
  • R3 reuses rollout-selected experts during training, but this can alter gradient attribution by omitting contributions from experts the current training router would select.Token-probability corrections such as TIS do not directly identify changes in expert assignments.Source: §1 · p. 2
  • RDC measures routing disagreement by response and filters high-drift responses while retaining native training routes and their associated learning signals.The approach uses response-level variation in routing disagreement to exclude a small fraction of responses from each update.Source: §1 · p. 2
  • 5.5 points over TIS and 4.5 points over R3: RDC combined with TIS improves average peak validation accuracy on AIME 2024 and AIME 2025 in asynchronous training.The introduction also reports that gains become more pronounced as policy staleness increases and that training dynamics remain stable.Source: §1 · p. 3

2 RELATED WORK

Prior work addresses rollout–training mismatch through system, probability-space, and routing-level corrections. The section contrasts those approaches and positions RDC as response-level filtering that preserves native training routes.

2 RELATED WORK

  • The related work groups mismatch corrections by whether they operate at the system, probability, or routing level.Rollout–training mismatch includes both engine differences and policy staleness from sequential updates or asynchronous generation.Source: §2 · p. 3

2.1 SYSTEM-LEVEL CORRECTIONS

  • System-level methods improve numerical consistency between rollout and training engines, but cannot resolve staleness from sequential mini-batch updates or asynchronous training alone.The section therefore identifies a continuing need for algorithm-level corrections.Source: §2.1 · p. 3

2.2 PROBABILITY-SPACE CORRECTIONS

  • Probability-space corrections adjust likelihood ratios or token contributions, but do not directly represent the computation paths taken by MoE routing.The section lists token-level ratio methods and sequence-level or selective corrections, then notes that routing can assign the same response to different experts.Source: §2.2 · p. 4

2.3 ROUTING CORRECTIONS FOR MOE MODELS

  • Routing methods replay or predict routes, or use router-aware weights; RDC instead selects responses by router drift while preserving retained responses’ native training routes.RDC combined with TIS addresses routing disagreement and token-level probability mismatch through complementary mechanisms.Source: §2.3 · p. 4

3 METHOD

The method measures rollout–training routing disagreement at the response level and filters high-drift responses while retaining native training routes. It also explains why routing can change expert-level learning signals even when token probabilities do not reveal the route.

3.1 SETUP AND NOTATION

  • The setup distinguishes the rollout behavior policy from the current training policy, which evaluates sampled responses under teacher forcing.Each token is routed to K experts across L MoE layers.Source: §3.1 · p. 4

3.2 ROUTER DRIFT CONTROL

  • RDC compares rollout and training expert-usage distributions across layers, averages their differences into a response-level drift score, and filters responses with the highest scores.Responses are ranked within groups of similar lengths.Source: §3.2 · p. 5

3.3 ROUTING DETERMINES GRADIENT SUPPORT

  • Routing determines which expert parameter blocks can receive direct gradients, so differing rollout and training routes can redistribute learning signals across experts.Experts selected during rollout may receive no direct training gradient, while experts selected only during training may receive one.Source: §3.3 · p. 6

3.4 TOKEN PROBABILITIES DO NOT IDENTIFY ROUTING

  • Under stated conditions, distinct routing assignments can yield identical token probabilities, so likelihood ratios and token-level KL divergence cannot identify which route was used.The assignments can nevertheless direct gradients to different expert parameter blocks.Source: §3.4 · pp. 6–7

3.5 SELECTIVE CONTROL VERSUS ROUTE REPLAY

  • Unlike route replay, RDC retains the native training route and filters responses by routing disagreement; route substitution can alter expert-gradient support and bias the native-routing objective.RDC can be combined with TIS, which adjusts token-level contributions while RDC filters responses based on router drift.Source: §3.5 · p. 7

4 EXPERIMENTS

The experiments compare RDC with GRPO, TIS, and R3 under synchronous and asynchronous training. RDC’s gains increase with staleness, and asynchronous results show improved AIME accuracy over TIS and R3 alongside stable training dynamics.

4.1 EXPERIMENTAL SETUP

  • The evaluation compares GRPO, R3, TIS, and RDC+TIS on mathematical and logical reasoning benchmarks, reporting best validation accuracy during training.The setup uses Qwen3-30B-A3B trained on DAPO-Math-17K and evaluates AIME 2024, AIME 2025, MATH-500, and ZebraLogic.Source: §4 · p. 7, §4.1 · p. 7
  • Synchronous experiments vary update-batch size while holding the rollout batch at 2,048 responses; asynchronous training uses independent rollout and optimization processes with queued samples.In asynchronous training, the Trainer recomputes routing under the current policy and filters high-drift responses before updating.Source: §4.1 · p. 8

4.2 EXPERIMENTAL RESULTS

  • Under (2048, 512), RDC achieves the highest accuracy on all four benchmarks; versus TIS, it improves ZebraLogic, AIME24, and AIME25 by 7.9, 6.4, and 4.3 percentage points under (2048, 128).The passage reports larger gains as the update batch becomes smaller.Source: §4.2 · p. 8, Table 1 · p. 8
  • 5.5 points over TIS and 4.5 points over R3 on average peak AIME24 and AIME25 accuracy are achieved by RDC under asynchronous training.RDC also has the highest accuracy on AIME24, AIME25, and MATH-500, while R3 performs slightly better on ZebraLogic.Source: Table 2 · p. 8, §4.2 · p. 9
  • During asynchronous training, GRPO’s reward and response length collapse late, while TIS, R3, and RDC avoid that collapse; GRPO’s training–inference KL and gradient norm rise sharply.TIS and RDC keep KL and gradient norm low and stable, while R3 shows a more pronounced late-stage increase in both.Source: §4.2 · p. 9, Figure 5 · p. 9

5 CONCLUSION

The conclusion identifies router drift as an off-policy mismatch that changes expert-level gradient support. RDC filters high-drift responses while retaining native routing, and combined with TIS improves validation performance and training stability in the evaluated settings.

5 CONCLUSION

  • Router changes alter expert-level gradient support, and replaying rollout routes changes the training gradient relative to the current policy’s native routing.The paper presents routing-aware filtering as an alternative to rollout routing replay.Source: §5 · p. 9
  • RDC filters high-drift responses while preserving the current training policy’s routing decisions.Source: §5 · p. 9
  • Combined with TIS, RDC improves validation performance and training stability under mini-batch updates and fully asynchronous training.Source: §5 · p. 9

A EXPERIMENTAL DETAILS

The appendix gives additional hyperparameter and implementation details for the experiments, including evaluation scoring and asynchronous RDC processing.

A EXPERIMENTAL DETAILS

  • Table 3 summarizes additional experimental hyperparameters and implementation details for RDC and the two training regimes.Source: A EXPERIMENTAL DETAILS · p. 13
  • AIME 2024 and AIME 2025 accuracy averages binary correctness across 32 sampled responses per problem and across problems.Training rewards are +1 for a correct answer and −1 otherwise; evaluation correctness is scored as 1 or 0.Source: A EXPERIMENTAL DETAILS · p. 13
  • For each asynchronous Trainer batch, RDC computes response-level expert-usage difference from BF16 routing assignments recorded during old-log-prob computation, then filters responses before the update.The asynchronous setup uses separate resource pools, require batches of 1, and a staleness threshold of 0.5.Source: A EXPERIMENTAL DETAILS · p. 13

B GRADIENT BIAS UNDER ROLLOUT ROUTING REPLAY

The appendix analyzes gradient bias under rollout routing replay by deriving local MoE-layer gradients and comparing replayed with native routing.

B GRADIENT BIAS UNDER ROLLOUT ROUTING REPLAY

B.1 SETUP

The setup considers a sparsely gated MoE layer and compares native routing with replay that fixes rollout expert identities while recomputing gate weights. Its local-gradient analysis holds the input fixed and excludes auxiliary router losses and token dropping.

B.1 SETUP

  • The analysis considers one token at a sparsely gated MoE layer, omitting token and layer indices for clarity.Source: §B.1 · p. 13
  • Routing replay fixes the rollout-selected expert identities but recomputes gate weights from the current router logits.Source: §B.1 · p. 13
  • Local gradients are analyzed with the layer input fixed, excluding auxiliary router losses and token dropping.The selected expert set is treated as fixed during backpropagation, so the derivatives apply away from routing-selection boundaries.Source: §B.1 · pp. 13–14

B.2 GRADIENTS UNDER FIXED ROUTING

The section derives local gradients under a fixed selected expert set, distinguishing selected experts’ gate and expert-parameter contributions from unselected experts’ router-logit contributions. These are one-token, one-layer contributions; an unselected expert may receive gradients from other tokens.

B.2 GRADIENTS UNDER FIXED ROUTING

  • For experts in the selected set, the analysis differentiates the normalized gate weights with respect to router logits.Source: §B.2 · p. 14
  • An unselected expert has no direct output dependence on its router logit while the selected set is held fixed.Source: §B.2 · p. 14
  • The section gives local parameter-gradient expressions for the linear router and for expert parameters using the expert Jacobian.Source: §B.2 · p. 14
  • These gradients describe one token at one layer, and an unselected expert may still receive gradients from other tokens.Source: §B.2 · p. 14

B.3 CHANGES IN LOCAL GRADIENT SUPPORT

Routing replay can remove local gradient contributions for experts selected by native routing but not by rollout routing. Experts selected only by replay can receive gradients instead, and shared experts may also receive different gradients.

B.3 CHANGES IN LOCAL GRADIENT SUPPORT

  • When an expert is selected by the current router but not the rollout router, replay gives zero local gradients with respect to its router logit and parameters.Under native routing, the corresponding expert-parameter and router-logit gradients are nonzero under the conditions stated in Proposition 1.Source: §B.3 · p. 14
  • Replay changes local gradient support: native-only selected experts lose direct contributions, while replay-only selected experts can receive gradients.Shared experts may also receive different gradients because gate weights and downstream gradients can change.Source: §B.3 · p. 14

B.4 GRADIENT BIAS RELATIVE TO NATIVE ROUTING

The replay and native-routing objectives use the same policy-optimization loss and fixed loss inputs, differing only in their forward-pass routing assignments. Their gradients can therefore differ, yielding nonzero bias relative to native routing in some cases.

B.4 GRADIENT BIAS RELATIVE TO NATIVE ROUTING

  • The objectives hold data, advantages, old log probabilities, and other loss inputs fixed, differing only in the forward-pass routing assignments.Source: §B.4 · p. 15
  • Routing replay can produce nonzero gradient bias relative to the native-routing objective.Source: §B.4 · p. 15
  • The bias can arise when native routing selects an expert excluded by replay, producing different expert-parameter gradients under the same training loss.Source: §B.4 · p. 15

Related papers

Loading 2610.21682v1…