---
title: Rollout Routing Replay (R3)
url: https://www.emergentmind.com/topics/rollout-routing-replay-r-3
type: topic
---

# Rollout Routing Replay (R3)

Rollout Routing Replay (R$^3$) is a methodology for stabilizing reinforcement learning (RL) in Mixture-of-Experts (MoE) transformer models by eliminating discrepancies between routing decisions made during inference (rollout generation) and those during training. The approach addresses the fundamental instability induced by non-deterministic or diverging expert selection across these phases, which manifests acutely in RL settings, leading to inflated policy KL divergence, high-variance importance weights, and frequent policy collapse. R$^3$ rectifies the mismatch by logging the expert selection mask during inference and deterministically replaying it in training, thus yielding a training policy closely aligned to inference-time decisions and dramatically improving RL stability and final task performance [2510.11370].

## 1. Training–Inference Discrepancy in MoE Routing

Reinforcement learning for large language models commonly separates the rollout (inference) and gradient (training) computation engines. In dense LMs, these pipelines produce nearly identical token-level probability distributions. However, MoE models use a sparse expert gating mechanism: at each token position, the output $y = \sum_{i=1}^M g_i \cdot E_i(x)$ depends on a sparse gate vector $g \in \mathbb{R}^M$ determined by a routing mask applied to the router logits. Minor numerical or implementation differences across inference and training can flip which experts are selected, causing dramatically different token outputs.

Empirical observations quantify this effect:
- On Qwen3-30B-A3B, $\mathrm{D_{KL}}(\pi_\mathrm{train} \| \pi_\mathrm{infer}) \approx 1.5 \times 10^{-3}$ for MoE (nearly 2$\times$ that of dense baselines).
- 10% of router calls select different experts between training and inference; at the token level, 94% of tokens differ in at least one layer.

This expert selection mismatch inflates the variance of the PPO importance weights $w_t = \pi_\mathrm{train}(y_t|x, y_{<t}) / \pi_\mathrm{train,old}(y_t|x, y_{<t})$; high variance in $w_t$ leads to unstable policy optimization, clipped ratios, and frequent RL collapse, where policy updates drift to trivial or degenerate behaviors [2510.11370].

## 2. Formal Specification of Rollout Routing Replay

At each token step $t$ in MoE layer $\ell$, denote:
- $x_\mathrm{infer}$: activation in the inference engine
- $s_\mathrm{infer} = x_\mathrm{infer} W_r$: router logits
- $I_\mathrm{infer} = \mathrm{TopKMask}(s_\mathrm{infer}, K) \in \{0,1\}^M$: top-$K$ mask of selected experts

The inference-time routing distribution:
$$\pi^\mathrm{inf}(e | x_{\leq t}) \propto I_\mathrm{infer}(e) \cdot \exp s_\mathrm{infer}(e)$$

During training, the system traditionally recomputes router logits $s_\mathrm{train} = x_\mathrm{train} W_r$ and applies its own mask $I_\mathrm{train}$. Rollout Routing Replay instead overrides the training mask with $I_\mathrm{infer}$ and computes a replayed gate:
$$g_\mathrm{replay}(e) = \frac{I_\mathrm{infer}(e) \cdot \exp s_\mathrm{train}(e)}{\sum_{j=1}^M I_\mathrm{infer}(j) \cdot \exp s_\mathrm{train}(j)}$$
resulting in the training routing distribution:
$$\pi^\mathrm{train}_{R^3}(e | x_{\leq t}) = g_\mathrm{replay}(e)$$

This framework views R$^3$ as minimizing:
$$L_{KL} = \mathbb{E}_{x}[\,KL(\pi^\mathrm{inf}(\cdot | x) \| \pi^\mathrm{train}_{R^3}(\cdot | x))\,]$$

Empirical measurements show R$^3$ halves the MoE train–infer KL divergence (from $1.5 \times 10^{-3}$ to $7.5 \times 10^{-4}$), thereby matching dense model behavior ($\approx 6.4 \times 10^{-4}$).

## 3. Algorithmic Workflow and Implementation

R$^3$ executes two phases per global RL step:

**1. Rollout (Inference Engine):**
- For each sample in a batch, sequence tokens are generated autoregressively.
- At each token and MoE layer, compute $s_\mathrm{infer}$ and $I_\mathrm{infer}$ (top-$K$ experts), cache mask per token per layer.
- Store $(x, y, \text{rewards}, \{I_\mathrm{infer}\})$ for training.

**2. Training (Training Engine with R$^3$):**
- For each minibatch, at each token and layer:
  - Compute $s_\mathrm{train}$, retrieve cached $I_\mathrm{infer}$
  - Form $g_\mathrm{replay}$ as above
  - Evaluate policy and PPO loss using $\pi^\mathrm{train}_{R^3}$
  - Backpropagate gradients

Crucially, only the mask is overridden; gradients still flow into $s_\mathrm{train}$. R$^3$ is applied in both "old" and "new" policy branches during on-policy PPO. Pseudocode and further details can be found in [2510.11370].

## 4. Theoretical Properties and Stability Effects

R$^3$ enforces near-identity between the training and inference routing distributions under small parameter updates ($s_\mathrm{train} \approx s_\mathrm{infer}$ as $\theta\to\theta_{\text{old}}$), causing $D_{KL}(\pi_\mathrm{inf} \| \pi^\mathrm{train}_{R^3}) \to 0$. Empirically, the fraction of "extreme tokens," defined as $F(\tau) = \text{fraction}\{t | \max(\pi_\mathrm{train}/\pi_\mathrm{infer}, \pi_\mathrm{infer}/\pi_\mathrm{train}) > \tau\}$, drops by an order of magnitude for $\tau > 2$ with R$^3$. 

As a result, importance weights $w_t$ become tightly centered around 1, correlating with stable PPO policy updates and eliminating the policy collapse observed in non-replayed MoE RL [2510.11370].

## 5. Empirical Evaluation and Benchmarks

R$^3$ has been evaluated on mathematical reasoning tasks over roughly 100,000 problems, with metrics reported on AIME24/25 (Avg@32), AMC23 (Avg@16), and MATH500 Lv5 (Avg@4). Using Qwen3-30B-A3B (Base and SFT), the following baselines were compared:
- GRPO: token-level clipped PPO
- GSPO: sequence-level importance sampling
- TIS: truncated importance sampling

Results indicate:
- Multi-step SFT: GSPO baseline achieves Avg $\approx 66.8$, GSPO+R$^3$: Avg $\approx 69.0$ (+2.2), GRPO+R$^3$: Avg $\approx 68.0$ (+1.3)
- Single-step SFT: GRPO collapses at step 60 (Avg 62.2), GRPO+TIS: Avg 66.2, GRPO+R$^3$: Avg 71.8 (+5.6 over TIS), with no collapse
- Similar gains observed with Base model

Stability curves demonstrate that D$_{KL}$ and F(2) escalate during collapsing runs but remain below $10^{-4}$ with R$^3$. Generation dynamics show smoother gradient norms, steadier entropy, and faster reward gains with R$^3$ [2510.11370].

## 6. Ablations and Sensitivity Analysis

Key ablation findings include:
- R$^3$ in combination with TIS provides no additive benefit or may slightly degrade performance, suggesting R$^3$ nearly closes the off-policy gap alone.
- R$^3$ stabilizes both single- and multi-mini-step PPO, with especially pronounced advantages in the fragile single-step regime.
- Caching $I_\mathrm{infer}$ masks (vs. recomputing) yields no fidelity loss but improves computational efficiency for long-context or agent scenarios.
- Default $K=2$ is robust, with preliminary tests showing similar gains for $K\in\{1,2,4\}$.

## 7. Practical Implications and Limitations

Implementing R$^3$ entails certain costs and constraints:
- Memory/storage overhead increases as $O(|y| \cdot \text{layers} \cdot M)$ due to per-token, per-layer caching of router masks; prefix caching is recommended for long sequences or dialogs.
- Framework support must enable extraction and injection of routing masks between rollout and training engines; this is straightforward if router code is shared but may require engineering otherwise.
- R$^3$ addresses routing-induced discrepancies alone; other nondeterministic elements (e.g., kernel-level or architectural mismatches) may still produce residual train–infer gaps.
- While demonstrated primarily for MoE Transformers on math reasoning tasks, the methodology generalizes to any sparse routing module (including conditional computation) and other RL applications (e.g., code generation, logical reasoning).

Future work may adapt R$^3$ to diverse router architectures (e.g., noisy top-$K$, Gumbel-Softmax), low-precision kernels, or distributed multi-node inference, but the protocol—record routing decisions at inference, replay in training—remains conceptually constant.

In summary, Rollout Routing Replay (R$^3$) directly targets MoE router-induced train–infer mismatch, delivering a two-fold reduction in train–infer KL divergence, a substantial decrease in outlier tokens, eliminating PPO collapse, and yielding consistent 1–6 point improvements in downstream RL task performance [2510.11370].

Source: https://www.emergentmind.com/topics/rollout-routing-replay-r-3