Layerwise Recurrent Router for Mixture-of-Experts

🎯 Pitch

Standard MoE models waste capacity because each layer’s router blindly assigns tokens without knowing what previous layers selected, leaving a 52B MoE competing with a mere 6.7B dense model. RMoE fixes this by threading a lightweight GRU across layers to share routing history, delivering consistent gains over strong baselines with negligible slowdown.

1. Executive Summary

This paper introduces the Layerwise Recurrent Router for Mixture-of-Experts (RMoE), a architectural modification that integrates a Gated Recurrent Unit (GRU) into the routing process of MoE-based language models to establish dependencies between routing decisions across consecutive layers (e.g., informing layer i's router which experts were selected in layer i-1, enabling cross-layer coordination that standard independent routers lack). Evaluated on language modeling benchmarks (Enwiki8, WikiText-103) and scaled up to 0.91B-parameter models pre-trained on 40B tokens with supervised fine-tuning, RMoE consistently outperforms a spectrum of router baselines — including standard linear routers, MLP routers, XMoE, and HyperMoE — while introducing negligible training overhead (49.07 s/step vs. 48.87 s/step for SMoE at 0.91B scale). The gains are attributed to the GRU's facilitation of cross-layer mutual information sharing and the provision of an auxiliary Recurrent Gradient pathway that improves router optimization, establishing that layerwise recurrence benefits MoE models primarily through better gradient flow rather than simply increased router capacity.

2. Context and Motivation

The Core Problem: Parameter Inefficiency in Mixture-of-Experts Models

The fundamental problem this paper addresses is a persistent and somewhat embarrassing fact about Mixture-of-Experts (MoE) language models: they are dramatically parameter-inefficient. Despite their appealing property — scaling total model parameters without proportionally increasing the per-token computational cost (because each token only activates a sparse subset of experts) — current pre-trained MoE models consistently underperform dense models with far fewer total parameters. The paper provides concrete, sobering examples right in the introduction:

"Rajbhandari et al. (2022) shows that with the same training data, an MoE with 52B parameters and 1.3B activated ones for each token performs similarly to a 6.7B standard model."

This is a staggering gap: a model with nearly 8× more total parameters achieves parity with a dense model less than one-eighth its size. The problem persists across different MoE designs. Another example:

"Komatsuzaki et al. (2023) demonstrates that upcycling a standard T5-base (248M) into its MoE counterpart (2B) by copying existing FFN can bring some improvements, but it still lags behind the T5-large with 783M parameters."

And even with more sophisticated architectures incorporating fine-grained and shared experts:

"Dai et al. (2024) use fine-grained and shared experts to improve the effectiveness, but the 16B MoE performs comparably with the 7B standard model (Bi et al., 2024)."

The practical implication is stark: if you have a fixed inference compute budget (proportional to activated parameters), you are paying for the memory and storage costs of a massively larger model without getting proportional performance benefits. The paper does not frame this as a fatal flaw of MoE — the ability to scale total parameters while keeping FLOPs roughly constant is genuinely valuable — but rather as an inefficiency that demands explanation and remediation. Why are all those extra parameters not pulling their weight?

Why This Matters: The Router as the Suspected Bottleneck

The paper homes in on a specific component as the likely culprit: the router. In standard Sparse Mixture-of-Experts (SMoE), each transformer layer contains one router — typically a single linear layer followed by a softmax and a Top-k selection — that determines which experts process each token. This router is lightweight by design (its parameter count is a tiny fraction of the total model), but its decisions have outsized consequences: if the router makes suboptimal assignments, entire banks of expert parameters may be systematically underutilized or misapplied.

The paper marshals compelling evidence that current routers are far from optimal:

Evidence 1: Routing collapses to shallow, token-ID-based patterns almost immediately. The paper cites Xue et al. (2024), who find that "the routing results converge to the token-id-based routing very quickly during the early phase of pre-training." This means the router essentially learns a static mapping — certain tokens always go to certain experts, regardless of context, position in the sequence, or the semantic demands of the input. Such a shallow routing mechanism cannot possibly exploit the full combinatorial richness of expert assignments across layers and tokens.

Evidence 2: Non-learned and even fixed-random routers are competitive with learned ones. This is perhaps the most damning finding. The paper notes that "hash functions (Roller et al., 2021), stochastic routing policy (Zuo et al., 2021), and fixed-random router (Chen et al., 2023) achieves competitive performance with the learnable router." If a random, untrained router performs roughly as well as one that has been optimized through thousands of training steps alongside the rest of the model, this strongly suggests that the learning signal available to the router is insufficient to discover meaningful token-expert assignments. The router is effectively a deadweight component — present and trainable, but not learning anything useful beyond what a static hash function could provide.

Evidence 3: Simply increasing router capacity does not help. The paper's own experiments bear this out. Comparing SMoE-MLP (a two-layer MLP router with GELU activation, following Shen et al., 2023) against the standard linear router, the authors find:

"replacing the original simple linear layer with a more capable MLP does not improve performance. It even underperforms the fixed random routing (RandomMoE) on Enwikik8, suggesting that naively increasing model capacity can't result in a more powerful router."

This is a crucial negative result: the problem is not that routers lack parameters or representational capacity. A larger router with more degrees of freedom does not automatically learn better assignments. Something more fundamental is broken about how routers receive information and gradients during training.

Where Prior Router Designs Fall Short: The Isolation Problem

The paper's central diagnosis is that the problem is primarily informational, not computational. Specifically, the router at each layer operates in isolation:

"current routers in different MoE layers still operate independently without comprehensive investigations into the decisions of other layers. This isolation may lead to suboptimal expert utilization, as each layer manages its routing based solely on local information."

Think about what this means. At layer 5 of a 24-layer transformer, the router sees only the hidden state representation produced by layer 4. It has no direct knowledge of which experts were activated at layer 3, layer 2, or layer 1 — information that could be highly relevant for making coordinated assignments. In principle, the hidden state residual stream could carry some routing-relevant information forward through the layers, but the paper argues this is unreliable:

"though vanilla MoE models could technically share the routing information via hidden states residual, this information may be overshadowed by the language modelling loss, requiring routing-relevant information to 'compete' for its representation."

The hidden states exist in a shared representational space optimized for predicting the next token. Any routing-relevant signal they contain must coexist with — and likely be dominated by — the semantic and syntactic features needed for language modeling. There is no dedicated channel for conveying "here's which experts were activated previously, and whether those activations were effective" to downstream routers.

Prior router improvements have focused on other aspects:

XMoE (Chi et al., 2022): Projects hidden states to a lower dimension and uses cosine similarity to expert embeddings, which the paper says "can prevent the hidden states from collapsing to a linear combination of expert embeddings." This addresses a representation collapse problem but does nothing to share information across layers.
HyperMoE (Do et al., 2023): Uses a fixed random hypernetwork to generate router weights conditioned on input, increasing "capacity" in a sense but still operating independently per layer.
SMoE-Dropout / RandomMoE (Chen et al., 2023): Uses a fixed random-initialized router with gradually increasing Top-k. This circumvents the router learning problem entirely by abandoning learned routing, which is a clever insight into its difficulty but a retreat rather than a solution.

The underlying structural problem — layerwise isolation — persists across all these designs. Each router makes its decisions in a vacuum, with no mechanism for coordination, no feedback channel for downstream layers to influence upstream routing, and no explicit memory of past routing decisions within a token's trajectory through the network.

A Philosophical Parallel: The Argument for Dedicated Routing State

The paper's solution has a clear conceptual motivation that it states explicitly: if the standard residual stream is a poor medium for carrying routing information (because it is already saturated with language-modeling-relevant features), then introduce a dedicated, purpose-built channel that exists solely to propagate routing state across layers. This is the deep design principle behind RMoE.

This is analogous to several well-established architectural patterns in deep learning: (a) LSTMs and GRUs introduced dedicated cell states separate from the main activation stream, allowing information to flow across time steps without being repeatedly transformed and attenuated by nonlinearities; (b) the original transformer introduced residual connections to provide a dedicated "highway" for gradient flow, insulated from the transformations of attention and FFN sublayers; (c) in memory-augmented neural networks, external memory banks provide a separate read-write channel for information that does not need to be compressed into the network's internal representations.

RMoE applies this same logic to the routing problem: the GRU provides a separate state vector $h_i$ that lives in its own low-dimensional space (typically 128 dimensions, much smaller than the hidden state dimension of 352 or 1280) and is updated at each layer using a dedicated projector and recurrent unit. This state is not burdened with representing syntactic structure or world knowledge; it exists only to track routing-relevant information — which experts were chosen, how confident those choices were, and potentially whether those activations proved useful. The paper positions this not as an incremental tweak but as a new computation stage in the MoE architecture, "orthogonal to and compatible with most existing methods."

How This Paper Positions Itself Relative to Existing Work

The paper's framing in Section 2 makes clear that it sees the router design space as under-explored relative to other aspects of MoE. Many prior works focus on what happens after routing — expert specialization, training stability, load balancing — but relatively few question the routing mechanism itself. Among those that do, most treat routing as a per-layer classification problem (which expert is best for this token right now?) without considering its sequential, cross-layer structure.

The paper acknowledges one concurrent work that is similar in spirit: Gong et al. (2024), which also introduces a GRU in sequential routing stages. However, the paper explicitly distinguishes itself:

"it does not view such a recurrent mechanism as a general and composable method with broad MoE fields or provide relative ablation or analysis."

The implication is that RMoE's contribution is not the specific choice of GRU — which any practitioner could think of — but rather (1) the systematic validation that layerwise recurrence consistently helps across scales, architectures, and training paradigms; (2) the careful disentanglement of why it helps (cross-layer information vs. recurrent gradients vs. additional parameters); and (3) the demonstration that it composes orthogonally with existing router designs (e.g., XMoE + GRU router outperforms either alone). The paper positions itself not as a one-off router variant but as a design principle — layerwise recurrence — that can be retrofitted onto any MoE architecture with minimal cost.

The Practical Stakes: Why This Problem Is Worth Solving

The paper does not dwell on the broader significance of fixing MoE parameter inefficiency, but the implications are substantial:

Economic efficiency of large-scale training and deployment: MoE models are popular precisely because they offer a path to scaling total parameters without scaling per-token FLOPs. If those extra parameters are largely wasted, the entire value proposition of MoE is undermined. Making MoE parameters more effective directly improves the return on investment for training and serving large models.
Enabling on-device and edge deployment: If a fixed inference budget achieves higher quality through better expert utilization, smaller activated-parameter models can handle tasks previously requiring larger ones, reducing hardware requirements.
Modularity and interpretability: Better routers that coordinate across layers could lead to more meaningful expert specialization — experts that consistently handle specific linguistic phenomena, domain knowledge, or reasoning patterns — which has downstream benefits for model interpretability, controlled generation, and continual learning.
Composability with future MoE innovations: Because RMoE introduces a novel computation stage orthogonal to existing MoE designs, improvements in expert architectures (fine-grained experts, shared experts), training strategies (auxiliary-loss-free balancing, sparse upcycling), and routing policies (expert choice, soft MoE) can all potentially benefit from the addition of cross-layer routing state. RMoE is not a competing method but a complement.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

This paper proposes a modification to the router inside Mixture-of-Experts transformer layers: instead of each layer's router picking experts independently based only on the current hidden state, the authors add a small recurrent neural network (a GRU) that carries a dedicated "routing memory" vector forward from layer to layer, so the router at layer 5 knows which experts were chosen at layers 1–4 and can coordinate accordingly. The core problem being solved is that standard MoE routers are isolated per layer — they have no mechanism to share information about which experts have already been activated for a token, leading to suboptimal expert utilization and the well-documented parameter inefficiency where 52B-parameter MoE models perform like 6.7B dense models. The solution is to add a lightweight, cross-layer-shared GRU that maintains a routing-specific hidden state, creating an explicit communication channel that does not compete with language modeling signals in the main residual stream.

3.2 Big-picture architecture (diagram in words)

The RMoE architecture inserts a new computation stage between the transformer's hidden state output and the router's expert selection, at every MoE layer. Here are the components and their roles:

Layerwise Hidden State Projector (one per layer): At layer $i$ , takes the transformer's hidden state $x_i$ (e.g., dimension 352 or 1280) and projects it down to a much smaller dimension $p$ (typically 128) via a learned linear layer. This produces $x'_i$ , a compressed representation that captures routing-relevant features from the current token's representation. Each layer has its own projector because hidden states vary substantially in norm and distribution across layers.
Cross-Layer-Shared GRU (one per model): A Gated Recurrent Unit that is shared across all layers. It takes the projected state $x'_i$ and the GRU hidden state from the previous layer $h_{i-1}$ as input, and produces the current layer's GRU output $h_i$ . This is the mechanism that establishes explicit cross-layer dependency — $h_i$ encodes information about routing decisions across all previous layers, not just the immediate predecessor. The GRU is shared (same parameters at every layer) because the operation being modeled — "update routing memory based on current token representation" — is the same abstract function at each layer.
Per-Layer Router (linear layer, one per layer): Takes the GRU output $h_i$ (not the original hidden state $x_i$ ) as input, and computes the standard gating scores over all $N$ experts. The router parameters $G_i$ are layer-specific (not shared), preserving each layer's freedom to have different expert affinities.
Experts (FFNs, one set per layer): The expert networks themselves are unchanged from standard SMoE. They receive the original hidden state $x_i$ as input — not the GRU output — meaning the GRU affects only the selection of experts, not the computation those experts perform.

Information flow (step by step): At layer $i$ , the hidden state $x_i$ arrives from the previous layer's transformer block (attention + FFN/MoE output). First, the layer-specific projector compresses $x_i$ to $x'_i$ . Second, the shared GRU combines $x'_i$ with $h_{i-1}$ (the GRU state from layer $i-1$ ) to produce $h_i$ . Third, the layer-specific linear router processes $h_i$ to produce expert selection scores. Fourth, those scores determine which experts process $x_i$ (via standard Top-k gating). Fifth, the GRU state $h_i$ is passed forward to layer $i+1$ . Notice that the GRU state forward pass is completely independent of the expert computation — it runs in parallel with it, not as a sequential dependency, which is why the paper reports negligible wall-time overhead.

3.3 Roadmap for the deep dive

First, the standard MoE routing formulation (Equation 1) and how RMoE modifies it, to establish precisely what changes in the mathematical definition of the model.
Second, the GRU mechanism itself — its internal gates, how it updates the routing memory $h_i$ , and why this specific recurrent architecture was chosen over alternatives.
Third, the layerwise projector design — why separate projectors per layer are necessary and what happens when they are shared.
Fourth, the integration of the GRU output into the router (Equation 6) and why the original hidden state still feeds the experts while the GRU state feeds the router.
Fifth, the training procedure, loss functions (language modeling loss + load balancing loss), and how the Recurrent Gradient emerges from backpropagation through the GRU's cross-layer connections.
Sixth, the compositionality claim — how RMoE's new computation stage integrates with existing router designs like XMoE, and the specific hyperparameter choices made in experiments.

3.4 Detailed, sentence-based technical breakdown

This is primarily an architectural design paper whose core idea is that routing in MoE models should be treated as a sequential decision-making process across layers, and that a lightweight recurrent mechanism can capture inter-layer routing dependencies that isolated per-layer routers miss. The contribution is not a new training algorithm or loss function, but a structural modification to the MoE computation graph that introduces a dedicated channel for routing memory.

Standard Mixture-of-Experts Routing (Base Formulation)

The paper begins by restating the standard SMoE formulation to establish a clear baseline that RMoE modifies. In a standard SMoE layer with $N$ experts denoted $E_n$ for $n \in [1, N]$ , the router $g(\cdot; G, k)$ is defined by its parameters $G \in \mathbb{R}^{h \times N}$ (a matrix mapping from hidden dimension $h$ to $N$ expert scores) and an integer $k$ (the number of experts to activate). Given an input token representation $x \in \mathbb{R}^h$ , the layer output $y \in \mathbb{R}^h$ is:

$y = \sum_{n \in N} g_n(x; G, k) \, E_n(x)$

where $g_n(x; G, k)$ is the $n$ -th element of the router's output distribution over experts, $E_n(x)$ is the output of expert $n$ on input $x$ , and the sum is over all $N$ experts (but only $k$ experts receive non-zero gating weights due to the Top-k operation, making it sparse).

What this equation computes: The output $y$ is a weighted combination of expert outputs, where the weights come from the router's assessment of which experts are most appropriate for the current token. The router first computes raw scores $x \cdot G$ (a linear transformation producing $N$ unnormalized scores), then applies a softmax to convert these to a probability distribution, then zeros out all but the top $k$ entries (Top-k sparsification), and finally renormalizes. The surviving $k$ experts each process the input $x$ through their FFN, and their outputs are summed with the corresponding normalized gating weights.

Why this form: The linear-then-softmax-then-Top-k router is the de facto standard in SMoE because it is simple, fast, and differentiable (the Top-k operation produces sparse gradients but the softmax is smooth). The $n$ -th column of $G$ can be interpreted as an "expert embedding" — the router assigns high scores when the hidden state $x$ aligns with a particular expert's embedding direction. The key limitation that RMoE addresses is that this computation happens independently at each layer — the router $g(\cdot; G_i, k)$ at layer $i$ sees only $x_i$ , with no access to which experts were chosen at layers $1$ through $i-1$ .

The Layerwise Projector: Compressing Hidden States for the GRU

The first component RMoE adds is a per-layer linear projector that compresses the high-dimensional hidden state $x_i \in \mathbb{R}^h$ into a lower-dimensional representation $x'_i \in \mathbb{R}^p$ suitable for the GRU's internal state (where $p$ is typically 128, much smaller than $h$ which is 352 for the small-scale experiments and 1280 for the 0.91B-scale experiments):

$x'_i = \text{Proj}_i(x_i)$

where $\text{Proj}_i: \mathbb{R}^h \to \mathbb{R}^p$ is a learned linear transformation specific to layer $i$ . Each layer $i$ has its own projector matrix $W_{\text{proj}, i} \in \mathbb{R}^{h \times p}$ (and bias, though this is not explicitly stated, it is standard for linear layers).

What this equation computes: The projector takes the full hidden state $x_i$ (which contains all information needed for language modeling — syntax, semantics, positional information, etc.) and extracts a compact summary that is specifically relevant for routing decisions. The output $x'_i$ lives in a much smaller space (128 dimensions) dedicated to the GRU's routing memory, explicitly separating routing-relevant features from the main representational stream.

Why this form and why per-layer: The paper provides both empirical and analytical justification for per-layer projectors. First, from the analysis in Appendix A.4.5 (Figure 11), the paper shows that router weight norms and standard deviations vary substantially across layers — early layers have different statistical properties than middle and late layers. This means a single shared projector would need to handle very different input distributions, which is challenging. The paper explicitly tests this in the ablation (Table 6): "Layerwise projector in Eq. 5 performs better than standard RNNs using a single shared projector" (validated at 1.141 BPC with per-layer vs. 1.148 BPC with shared on Enwiki8). The paper draws a parallel to the established practice of not sharing LayerNorm parameters when employing shared MoE transformer blocks (Xue et al., 2022) — the same principle applies here. Second, conceptually, different layers process fundamentally different features (lower layers encode local syntax, higher layers encode abstract semantics), and the routing-relevant information at each layer may have a different structure that benefits from layer-specific compression.

The choice of $p=128$ as the GRU hidden dimension is a hyperparameter swept in the paper (Table 7). Testing at the 0.91B scale with 20B training tokens, the paper finds that $p=128$ (GRU with averaged downstream task score of 49.20 pretrain, 51.32 after SFT) outperforms both $p=256$ (49.08 pretrain, 50.04 SFT) and $p=512$ (49.19 pretrain, 50.02 SFT). The sweet spot at $p=128$ suggests that the routing memory channel benefits from moderate capacity — too small loses information, too large introduces trainable parameters that may overfit or destabilize training.

The Gated Recurrent Unit (GRU): Establishing Cross-Layer Routing Memory

The GRU is the computational heart of RMoE. It is shared across all layers (the same GRU parameters are used at every MoE layer), and its role is to maintain and update a routing memory state $h_i \in \mathbb{R}^p$ that propagates forward through the layers. The paper first reviews the standard GRU formulation (Equations 2 and 3, cited from Dey & Salem, 2017), then shows how it is applied in the RMoE context (Equation 5).

Standard GRU Formulation (Reference)

At each time step $i$ , given an input $x'_i$ and the previous hidden state $h_{i-1}$ , the GRU computes:

$s_i = \sigma(W_s x'_i + U_s h_{i-1})$

$z_i = \sigma(W_z x'_i + U_z h_{i-1})$

where $\sigma$ is the sigmoid activation function, $W_s, W_z \in \mathbb{R}^{p \times p}$ are input-to-gate weight matrices, $U_s, U_z \in \mathbb{R}^{p \times p}$ are hidden-to-gate weight matrices, $s_i \in (0, 1)^p$ is the reset gate (controlling how much of the previous state to forget), and $z_i \in (0, 1)^p$ is the update gate (controlling the interpolation between the old state and a candidate new state).

Then the candidate hidden state $\tilde{h}_i$ and the new hidden state $h_i$ are computed as:

$\tilde{h}_i = \tanh(W_h x'_i + s_i \odot (U_h h_{i-1}))$

$h_i = (1 - z_i) \odot \tilde{h}_i + z_i \odot h_{i-1}$

where $W_h, U_h \in \mathbb{R}^{p \times p}$ are additional weight matrices, $\odot$ is element-wise multiplication, and $\tanh$ is the hyperbolic tangent activation.

What these equations compute: The GRU performs a gated combination of two sources of information — the current input $x'_i$ (what routing-relevant features does this token have at this layer?) and the previous hidden state $h_{i-1}$ (what routing decisions were made in previous layers?). The reset gate $s_i$ determines which dimensions of the previous state $h_{i-1}$ are allowed to influence the candidate new state $\tilde{h}_i$ — if $s_i[j]$ is close to 0, dimension $j$ of the previous state is effectively ignored when computing the candidate. The update gate $z_i$ determines the final interpolation between the candidate state $\tilde{h}_i$ (new information) and the previous state $h_{i-1}$ (old information). If $z_i[j]$ is close to 1, dimension $j$ of the state retains its old value (the GRU "remembers"); if $z_i[j]$ is close to 0, it adopts the new candidate value (the GRU "updates").

Why this form: The GRU was chosen over a vanilla RNN and over an LSTM after empirical comparison. Table 6 shows that on Enwiki8, the GRU router achieves 1.141 validation BPC vs. 1.145 for a simple RNN router and 1.148 for an LSTM router. The GRU's gating mechanism is particularly well-suited here because: (1) The update gate provides an explicit mechanism for the model to decide whether to carry forward routing information from previous layers (e.g., if the same expert specialization is useful) or to overwrite it (e.g., if the token's role has changed at higher layers). (2) The reset gate allows the model to selectively ignore irrelevant history when computing a new candidate state, preventing stale routing information from constraining current decisions. (3) Compared to LSTM, the GRU has fewer parameters (no separate cell state, no output gate), which matters because the GRU is replicated (logically) at every layer — even though the parameters are shared, the computational cost is incurred at every layer, so a leaner architecture is preferable.

RMoE's GRU Application (Equation 5)

In the RMoE architecture, the GRU is applied with a crucial twist that distinguishes it from standard sequential RNN usage:

$h_i = \text{GRU}(x'_i, h_{i-1})$

where $h_0$ is initialized to a zero vector at the start of each token's forward pass through the layers.

What this equation computes: For each token and each layer $i$ , the shared GRU takes the projected hidden state $x'_i$ (the token's routing-relevant features at layer $i$ , compressed to $p$ dimensions) and the GRU state from the previous layer $h_{i-1}$ (the accumulated routing memory from layers 1 through $i-1$ ), and produces an updated state $h_i$ that encodes routing information through layer $i$ . Note that this is layerwise recurrence, not sequence recurrence: the GRU unrolls across transformer layers (depth), not across token positions (time). This is a critical design choice with major implementation implications — because the recurrence is across layers, all tokens in a batch can be processed in parallel at each layer; there is no sequential dependency across tokens within a layer. The paper explicitly states this advantage:

"Such operation doesn't introduce sequence-level recurrence and can be efficiently implemented, as shown in Tab. 1 and Tab. 3."

Why layerwise recurrence and not sequence recurrence (or both): The paper's motivation is specifically about inter-layer coordination of expert assignments. When a token passes through layer 1, the router decides which experts process it. When it reaches layer 2, that information is lost in the standard architecture. The GRU preserves it. This is fundamentally different from sequence-level recurrence (e.g., maintaining a routing state across token positions in a sentence), which the paper does not explore. The choice of layerwise recurrence is justified by the diagnosis in the introduction — the problem is "current routers in different layers still operate independently" — and the solution directly addresses that specific isolation. Sequence-level recurrence would address a different problem (e.g., routing consistency for the same token concept across positions) and would introduce genuinely sequential dependencies that hurt training parallelism, a tradeoff the authors explicitly avoid.

Routing with GRU State: Separating Routing Input from Expert Input

The final step in each RMoE layer is the actual expert selection and computation, formalized in Equation 6:

$y_i = \sum_{n \in N} g_n(h_i; G_i, k) \, E_n(x_i)$

where $y_i$ is the output of the $i$ -th MoE layer, $h_i$ is the GRU output at layer $i$ , $g_n(h_i; G_i, k)$ is the router's output for expert $n$ (computed from $h_i$ using layer-specific router parameters $G_i$ ), $E_n(x_i)$ is the output of expert $n$ on the original hidden state $x_i$ , and the sum is over all $N$ experts with sparsity enforced by Top-k.

What this equation computes: The layer output is computed exactly as in standard SMoE (Equation 1), but with one crucial substitution: the router's input is $h_i$ (the GRU output) rather than $x_i$ (the original hidden state). This means the expert selection is informed by the cross-layer routing memory, but the expert computation still operates on the full hidden state $x_i$ . In other words, the GRU influences which experts are chosen, but does not constrain what those experts compute. This is a deliberate separation of concerns.

Why this separation: There are two reasons this design choice matters, one practical and one conceptual. Practically, if the experts also took $h_i$ as input (instead of $x_i$ ), then the GRU's low-dimensional state ( $p = 128$ ) would become an information bottleneck for the entire FFN computation (which needs to process the full richness of the hidden state for language modeling). The expert FFNs have hidden dimensions of 352 (small-scale) or 448 (large-scale) — compressing the 352/1280-dimensional $x_i$ through a 128-dimensional GRU state would destroy the information needed for effective language modeling. Conceptually, the separation embodies the paper's design philosophy: the GRU is a dedicated routing channel that should not have to also carry the full representational load of the language model. It is the routing analog of a cell state in LSTM — a narrow, purpose-built information conduit that exists alongside the main computational pathway.

The layer-specific router parameters $G_i \in \mathbb{R}^{p \times N}$ operate on the $p$ -dimensional GRU state $h_i$ rather than the $h$ -dimensional hidden state $x_i$ . Since $p \ll h$ (128 vs. 352 or 1280), each router has fewer parameters than the standard linear router (which takes $x_i \in \mathbb{R}^h$ as input). For example, for a layer with 16 experts, the standard router has $h \times 16$ parameters, while the RMoE router has $p \times 16$ parameters. The GRU and projectors add parameters back (the GRU has $3 \times (p \times p + p \times p) = 6p^2$ parameters for the weight matrices plus biases, and each layer adds a projector with $h \times p$ parameters), so the net parameter change depends on the configuration. But the key point is that the routing computation is happening in a different representational space, one that is explicitly optimized for propagating routing decisions across layers.

The Recurrent Gradient: How Backpropagation Through the GRU Creates a New Optimization Pathway

A critical insight of the paper, validated through ablation experiments in Section 5 (Tables 4 and 5), is that the GRU provides more than just forward-pass information sharing — it creates a novel backward gradient pathway that the paper terms the Recurrent Gradient. Understanding this requires walking through what happens during backpropagation in both standard SMoE and RMoE.

Standard SMoE gradient flow: In standard SMoE, the router at layer $i$ receives gradients from exactly one source: the language modeling loss, backpropagated through the expert outputs $E_n(x_i)$ and the gating weights $g_n(x_i; G_i, k)$ . More precisely, the gradient with respect to the router parameters $G_i$ comes from $\partial \mathcal{L}_{\text{LM}} / \partial g_n$ , which flows through the gating score's contribution to the weighted sum in Equation 1. There is an additional gradient from the load balancing loss $\mathcal{L}_{\text{LB}}$ , but this is a simple auxiliary loss that pushes the router toward uniform expert assignment and does not carry meaningful signal about which specific experts are good for which specific tokens. The language modeling gradient through the gating weights is sparse (only $k$ out of $N$ experts receive non-zero gradients because of Top-k), and the signal is entirely local to layer $i$ — there is no gradient flow between router $i$ and router $i-1$ .

RMoE gradient flow (Recurrent Gradient): In RMoE, the router at layer $i$ receives gradients from the language modeling loss through the same path as standard SMoE (through $g_n(h_i; G_i, k)$ ), but it additionally receives gradients through a completely different path: the GRU hidden state $h_i$ . Here is how this works:

The GRU output $h_i$ is computed from $x'_i$ (which depends on router parameters of layer $i$ only through the projector $\text{Proj}_i$ ) and $h_{i-1}$ (which depends on everything that happened at layer $i-1$ ).
$h_{i-1}$ is computed from $x'_{i-1}$ and $h_{i-2}$ , and so on recursively back to $h_0 = 0$ .
During backpropagation, the gradient of the language modeling loss with respect to $h_i$ (denoted $\partial \mathcal{L}_{\text{LM}} / \partial h_i$ ) flows backward through the GRU update equations (Equations 2-3) to produce gradients for $h_{i-1}$ , the GRU parameters ( $W_s, U_s, W_z, U_z, W_h, U_h$ ), and the projector parameters at layer $i-1$ (through $\partial h_{i-1} / \partial x'_{i-1}$ and then $\partial x'_{i-1} / \partial \text{Proj}_{i-1}$ ).
This creates a chain of gradient flow that connects the router's training signal at deeper layers back to routing decisions made at shallower layers.

The paper validates the importance of this Recurrent Gradient through the "detach $h_{i-1}$ " experiment in Table 5. In the RMoE + detach $h_{i-1}$ setting, the forward pass is identical to RMoE (the GRU still receives $h_{i-1}$ and produces $h_i$ using all the same computations), but the backward pass is cut — the gradient is not allowed to flow from $h_i$ back to $h_{i-1}$ . The result is striking:

"RMoE + detach $h_{i-1}$ performs even worse than RMoE-NP, showing that the Recurrent Gradient is important."

Specifically, on Enwiki8 test, RMoE achieves 1.116 BPC, RMoE-NP (no passing of recurrent states at all) achieves 1.123 BPC, and RMoE + detach $h_{i-1}$ achieves 1.133 BPC — worse than both. This is a key piece of evidence that the GRU's benefit is not purely informational (providing forward-pass context about previous routing decisions) but also optimization-related (providing richer, longer-range gradient signals to the router parameters).

Analogy to residual connections: The paper explicitly draws a parallel to residual networks, noting that "the spirit echoes the principles behind residual network, where residual connection are used to create direct paths for gradient propagation, thereby mitigating gradient vanishing as layers deepen." Just as residual connections provide a direct gradient highway from deeper layers to shallower ones in standard transformers, the GRU provides a gradient highway specifically for the router parameters across layers. This interpretation is supported by the scaling behavior: Figure 2 shows that as model depth increases from 6 to 32 layers, the gap between RMoE and SMoE widens, and the gap between RMoE and RMoE-NP-r0.5 also increases. Deeper models suffer more from gradient vanishing in the router, and the Recurrent Gradient becomes progressively more valuable.

Alternative mechanisms that fail to replicate the Recurrent Gradient: The paper tests a simpler approach to cross-layer information sharing: directly adding the gating logits from the previous layer to the current layer's logits (RMoE + NP + r- $\alpha$ settings in Table 5). The idea, inspired by Realformer's residual attention scores (He et al., 2020), is that the gating score for expert $n$ at layer $i$ becomes:

$g_n(h_i; G_i, k) + \alpha \, g_n(h_{i-1}; G_{i-1}, k)$

where $\alpha$ is a mixing coefficient (tested at 0.5 and 1.0). In the forward pass, this does share routing information across layers. In the backward pass, it provides an additional gradient path through the residual connection (unless explicitly detached). However, the results show that this approach is only partially effective: RMoE+NP+r-0.5 achieves 1.124 BPC (vs. 1.123 for pure NP and 1.116 for full RMoE), and critically, when the gradient through the residual logits is detached (RMoE+NP+r-0.5+detach-r), performance crashes to 1.133 BPC. This confirms that the gradient pathway, not just the forward-pass information, is the active ingredient. Yet even with the gradient pathway intact, the logit residual underperforms the full GRU, and the paper explains why:

"the indexes of experts in layer $i$ are not aligned with those in other layers, directly adding logits can lead to improper constraints and hurt the model performance"

In other words, there is no guarantee that Expert 3 at layer 2 plays the same role as Expert 3 at layer 5. Adding their logits assumes a correspondence that may not exist, effectively constraining the router in a way that can be harmful. The GRU avoids this problem because it does not force alignment — it learns a flexible, non-linear transformation from the previous layer's routing state to the current layer's routing state, without imposing any direct equivalence between expert indices.

The Full RMoE Layer: Integration and Training

Putting all components together, a single RMoE layer (say, the $i$ -th transformer layer) processes an input hidden state $x_i$ as follows:

Projection: The layer-specific projector maps $x_i \in \mathbb{R}^h$ to $x'_i \in \mathbb{R}^p$ via a linear transformation.
GRU Update: The shared GRU takes $(x'_i, h_{i-1})$ and produces $h_i \in \mathbb{R}^p$ , the updated routing memory. At the first layer, $h_0 = \mathbf{0}$ (the zero vector).
Routing: The layer-specific linear router $G_i \in \mathbb{R}^{p \times N}$ maps $h_i$ to expert scores, softmax + Top-k is applied, producing sparse gating weights $g_n(h_i; G_i, k)$ for the top $k$ experts.
Expert Computation: The selected $k$ experts each process the original hidden state $x_i$ through their FFN, producing $E_n(x_i)$ .
Aggregation: The layer output is the weighted sum $y_i = \sum_{n \in \text{Top-}k} g_n(h_i; G_i, k) \, E_n(x_i)$ .
Forward State Passing: $h_i$ is passed to layer $i+1$ for the next GRU step. $y_i$ is the hidden state input to layer $i+1$ 's attention block.

Then the process repeats: $y_i$ goes through the next layer's attention mechanism to produce $x_{i+1}$ , which enters the RMoE block at layer $i+1$ .

Training: RMoE is trained with the same objectives as standard SMoE:

Language modeling loss $\mathcal{L}_{\text{LM}}$ : the standard cross-entropy loss for next-token prediction, applied to the model's final output logits.
Load balancing loss $\mathcal{L}_{\text{LB}}$ : an auxiliary loss that encourages the router to distribute tokens evenly across experts, preventing collapse where all tokens go to a single expert. The paper uses the standard formulation from the MoE literature with a balancing weight of 0.01 (as stated in Section 4.1: "we employ balance loss with weights 0.01 during training").

The total loss is $\mathcal{L} = \mathcal{L}_{\text{LM}} + 0.01 \cdot \mathcal{L}_{\text{LB}}$ . The GRU parameters, projector parameters, and router parameters are all trained end-to-end with this loss via standard backpropagation. No special training procedures, separate optimization stages, or auxiliary losses specific to the GRU are needed — the Recurrent Gradient emerges naturally from the backpropagation through the GRU's recurrent connections across layers.

Initialization: The GRU state at the first layer is initialized to zero ( $h_0 = \mathbf{0}$ ), which is standard for GRU initial states. The paper does not specify any special initialization for the GRU parameters, implying standard random initialization (likely Xavier/Glorot uniform or normal, as is default in most deep learning frameworks).

Computational cost: At each layer, the additional computation consists of (a) one matrix multiply $x_i \cdot W_{\text{proj}, i}$ with cost $O(h \times p)$ ; (b) the GRU forward pass with six matrix multiplies each of size $p \times p$ , costing $O(p^2)$ ; and (c) the router's linear transformation from $h_i$ to expert scores costing $O(p \times N)$ . The dominant cost is the projector at $O(h \times p)$ . For the 0.91B model configuration ( $h = 1280$ , $p = 128$ ), this is $1280 \times 128 \approx 164\text{K}$ multiply-adds per token per layer, compared to the expert FFN computation which is orders of magnitude larger (each expert FFN has hidden dimension 448 with $1280 \times 448 \times 2$ parameters, costing millions of multiply-adds). This is why the paper reports "negligible wall time and memory cost" in Table 3 — the actual measured overhead is 49.07 s/step for RMoE vs. 48.87 s/step for SMoE, a 0.4% increase, and 48.69 GB memory vs. 48.00 GB, a 1.4% increase.

Compositionality: Integrating RMoE with Existing Router Designs

One of the paper's key framing claims is that RMoE is "orthogonal to and compatible with most existing methods" because it introduces a computation stage (the projector + GRU chain) that transforms the input to the router, without modifying the router itself. This means any existing router design that currently takes $x_i$ as input can instead take $h_i$ as input, with the projector and GRU inserted before the router.

The paper demonstrates this explicitly with XMoE (Chi et al., 2022) in Table 2. The standard XMoE router works by:

Projecting the hidden state $x_i$ to a low-dimensional space (default dimension 16) via a layer-specific projection.
Computing cosine similarity between this projected vector and learned low-dimensional expert embeddings.
Applying a learnable temperature to the softmax operation.

The RMoE-augmented version (denoted "XMoE + GRU router" in Table 2) modifies this by:

Using the RMoE projector (dimension 128) and GRU to produce $h_i$ from $x_i$ and $h_{i-1}$ .
Feeding $h_i$ (instead of $x_i$ ) into the XMoE router's projection and cosine similarity computation.

The results in Table 2 show that the GRU router benefits all three tested configurations of XMoE (with lower dimensions of 8, 16, and 32). For example, XMoE(16) achieves 1.125 test BPC on Enwiki8; XMoE(16) + GRU router achieves 1.119 BPC. This demonstrates that the cross-layer information from the GRU is complementary to the anti-collapse benefits of XMoE's cosine-similarity routing — they address different problems (layerwise isolation vs. representation collapse) and can be combined additively.

Parameter accounting for composition: When combining RMoE with XMoE, the additional parameters are exactly the RMoE components (projectors + shared GRU), since the XMoE router replaces the standard linear router that the GRU output would normally feed into. Table 1 reports that standard SMoE has 0.04M router parameters, while RMoE has 0.47M — the difference (0.43M) is the projectors + GRU. When added to XMoE(16), which has 0.09M router parameters, the combined model would have approximately $0.09 + 0.43 = 0.52\text{M}$ routing-related parameters (the exact number depends on whether the 128-dimensional GRU output replaces or supplements the 16-dimensional XMoE projection; the paper's description suggests it replaces it, with $h_i$ directly entering the cosine similarity computation).

Freezing during fine-tuning: The paper notes an interesting practical consideration for downstream fine-tuning. Citing Zoph et al. (2022), the standard practice is to "freeze the router during SMoE tuning" because the router's learned assignments may be brittle and fine-tuning on a small dataset can disrupt them. For RMoE, the paper tests two freezing strategies: (1) freeze only the linear router layer (keeping the GRU and projectors trainable), and (2) freeze both the router and the GRU/projectors. Table 3 reports results for both, with the general pattern being that freezing the router (but not the GRU) preserves most of the pre-training gains during SFT, and freezing both is comparable to freezing neither in some configurations. This suggests the GRU state itself captures stable, transferable information about routing patterns that generalizes to the fine-tuning distribution, which is a practical advantage.

Hyperparameter Configurations Used in Experiments

The paper tests RMoE across two scales with different hyperparameters:

Small-scale language modeling (Section 4.2, Tab. 1):

Model: 8-layer decoder-only transformer, hidden size $h = 352$ , 8 attention heads, sequence length 512.
MoE: 16 experts, Top-2 selection, expert FFN hidden size 352.
RMoE: GRU hidden dimension $p$ is not explicitly stated for this scale, but the ablation in Table 6 suggests 128 (the default used in most settings).
Training: Adam optimizer, learning rate 0.0007, 4,000 warmup steps, batch size 48, 80,000 training steps.
Datasets: Enwiki8 (character-level, BPC metric), WikiText-103 (word-level, PPL metric).

Large-scale pre-training (Section 4.1/Section 4.2, Tab. 3):

Model: 24-layer decoder-only, Rotary Embedding, SwiGLU activations, RMSNorm, hidden size $h = 1280$ , 20 attention heads, sequence length 4,096.
MoE: 16 experts, Top-4 selection, fine-grained with expert FFN hidden size 448 (following DeepSeekMoE design), with shared experts.
RMoE: GRU hidden dimension $p = 128$ (default, with ablations at 256 and 512 in Tab. 7).
Total parameters: approximately 0.91B total, 0.53B activated per token. RMoE adds approximately 3.5M additional parameters relative to SMoE (stated in Table 3 caption).
Training: Global batch size 1,120, warmup 2,000 iterations, learning rate 4.2e-4, minimum learning rate 4.2e-5, cosine decay, Adam optimizer ( $\beta_1=0.9, \beta_2=0.95$ ), weight decay 0.1, gradient clipping 1.0.
Pre-training data: multilingual corpus including Wikipedia, finance, and legal texts, 20B or 40B tokens.
Load balancing: auxiliary loss weight 0.01.
SFT: Alpaca dataset (52K instruction-response pairs), 3 epochs, learning rate 2e-5, batch size 128, cosine schedule, warmup ratio 0.03.
Hardware: 8 NVIDIA A100 GPUs, training for approximately 5 days (40B tokens).

Summary of Design Rationale

The paper's architectural choices are driven by a clear set of principles:

Separation of routing state from language modeling state: The projector + GRU create a narrow, dedicated channel for routing information that does not have to compete with semantic/syntactic features in the main residual stream. This is the core insight that motivates the entire design.
Layerwise recurrence, not sequence recurrence: By unfolding the GRU across model depth rather than token positions, the architecture preserves training parallelism while addressing the specific problem of inter-layer routing isolation.
Lean recurrent mechanism: The GRU with $p=128$ is deliberately small — it adds negligible computational overhead (0.4% step time increase) while providing outsized benefits through better gradient flow and cross-layer information sharing. The ablations show that larger GRU states ( $p=256, 512$ ) do not help, and may hurt, suggesting the routing memory channel needs just enough capacity to carry useful signal without becoming a crutch or introducing optimization difficulties.
Layer-specific projectors, shared GRU: The projectors are per-layer because the hidden state statistics vary across layers (proven in Appendix Figure 11). The GRU is shared because the abstract operation — "given current routing features and routing history, produce updated routing state" — is the same function at every layer. This is the established pattern from weight-tied transformer blocks and shared LayerNorm, applied to the recurrent routing component.
Orthogonality to existing router designs: By inserting the projector + GRU as a preprocessing stage before whatever router mechanism is being used (linear, MLP, cosine-similarity, hypernetwork), RMoE composes with any existing router without requiring changes to that router's internal design. This is what the paper means by "a new computation stage orthogonal to existing methods."

4. Key Insights and Innovations

Innovation 1: Reframing MoE Routing as a Sequential Decision Problem Across Layers

The paper's most fundamental conceptual move is to redefine routing from a set of independent per-layer classification problems into a single sequential decision process unfolding across model depth. This is not a technical tweak — it is a reframing of what routing is and what information it should have access to.

What the field assumed before this paper: The dominant paradigm in MoE design — spanning the original Shazeer et al. (2017) formulation, Switch Transformers (Fedus et al., 2022), GShard (Lepikhin et al., 2021), and essentially all subsequent variants — treated each layer's router as a standalone function $g_i: \mathbb{R}^h \to \Delta^{N-1}$ that maps the current hidden state to a distribution over experts. The implicit assumption was that the hidden state $x_i$ contains all information necessary to make optimal routing decisions at layer $i$ , and that any coordination across layers happens implicitly through the shared representational space of the residual stream. Prior work on improving routers — XMoE's cosine-similarity routing (Chi et al., 2022), HyperMoE's hypernetwork-generated weights (Do et al., 2023), SMoE-MLP's increased capacity (Shen et al., 2023) — all operated within this per-layer independence assumption. They asked "given $x_i$ , what is a better function $g_i$ ?" without questioning whether $x_i$ alone was sufficient input.

What RMoE changes at the conceptual level: The paper argues that routing at layer $i$ should be conditioned not just on $x_i$ but on the history of routing decisions made at layers $1$ through $i-1$ . This transforms routing from a Markovian process (current state is sufficient) to a non-Markovian one (past decisions matter). The key insight is that the residual stream is an unreliable carrier of routing history — it exists to serve language modeling, not routing coordination, and routing-relevant signals must "compete" with semantic and syntactic features for representational space (as the paper explicitly argues in Section 1). This is analogous to the motivation behind LSTM's cell state: the main hidden state is busy representing output-relevant features, so a separate channel is needed for memory that persists across time steps. RMoE applies this same logic to the depth dimension, arguing that routing decisions form a trajectory through the layers and that this trajectory benefits from its own dedicated state vector.

Evidence that this reframing matters: The mutual information analysis in Figure 3 provides direct empirical support. In standard SMoE (Figure 3a), cross-layer mutual information between routing distributions is consistently low — routers at different layers make largely independent decisions. In RMoE (Figure 3d), cross-layer MI is substantially higher, showing that the GRU successfully creates the coordination channel that the reframing demands. Critically, when the GRU's recurrent state passing is disabled (RMoE-NP, Figure 3e), MI drops back toward baseline levels, and performance degrades accordingly (1.123 vs. 1.116 test BPC on Enwiki8, Table 5). This establishes that the reframing is not merely philosophical — it corresponds to a measurable structural property of the model's routing behavior and directly predicts performance.

Significance beyond performance: This reframing changes how researchers should think about router design. Rather than asking "how can we make each router better in isolation?" (the approach of all prior work cited in Section 2), the field should ask "what is the right state representation to carry forward across routing decisions, and how should that state be updated?" This opens a design space orthogonal to existing router improvements — it is compatible with any per-layer router architecture — and suggests that future work should focus on the routing state channel (its dimensionality, its update mechanism, what information it encodes) rather than solely on per-layer routing functions.

The paper's most counterintuitive finding, and the one with the deepest implications for understanding why RMoE works, is that the GRU's benefit comes primarily from the backward-pass gradient pathway it creates, not from the forward-pass routing history it provides. This is a diagnostic contribution rather than a design contribution — the GRU design itself is straightforward — but it fundamentally changes the explanation for why cross-layer recurrence helps.

What the naive explanation would be: The obvious story for why RMoE should outperform independent routers is that the GRU provides forward-pass information: the router at layer 5 knows which experts were selected at layers 1–4, enabling coordinated expert assignments that avoid redundancy, encourage complementary specialization, or maintain consistent processing pathways for each token. This is the story the paper itself leads with in the introduction and Figure 1.

What the evidence actually shows: The ablation in Table 5 systematically dismantles this explanation. RMoE + detach $h_{i-1}$ — which preserves the full forward-pass information flow (the GRU still receives and processes $h_{i-1}$ ) but cuts the backward gradient flow through the recurrent connection — performs worse than RMoE-NP, which removes the forward-pass information entirely (1.133 vs. 1.123 test BPC). This means that forward-pass information without the corresponding gradient pathway is actively harmful. Conversely, RMoE-NP+r-0.5, which adds a residual gating logit connection (providing some forward-pass information with gradient flow), matches but does not exceed RMoE-NP (1.124 vs. 1.123 BPC) — the forward-pass information alone does not drive the gains. And when the gradient through that residual connection is cut (RMoE-NP+r-0.5+detach-r), performance crashes to 1.133 BPC, mirroring the detach- $h_{i-1}$ result.

What this means conceptually: The GRU is not primarily a communication channel for routing decisions. It is primarily a gradient propagation structure that allows optimization signal from the language modeling loss at deep layers to flow back to shallow-layer router parameters through a dedicated, low-impedance pathway. This interpretation connects RMoE to a broader principle in deep learning: deep networks suffer from vanishing gradients, and architectures that provide direct gradient highways (residual connections, LSTM gates, attention shortcuts) succeed in part because they improve optimization, not just because they improve representational capacity. The paper explicitly makes this connection, noting the "spirit echoes the principles behind residual network."

This is not just a curiosity — it explains several otherwise puzzling findings:

Why simply increasing router capacity fails: SMoE-MLP adds parameters to the router but provides no new gradient pathways. The optimization landscape remains difficult, and more parameters in a poorly-optimized regime do not help.
Why the Recurrent Gradient matters more in deeper models: Figure 2 shows the gap between RMoE and SMoE widening from 6 to 32 layers. In deeper models, the gradient vanishing problem for shallow-layer routers (which receive gradient only through the sparse, indirect path of gating weights propagated through many layers of FFN and attention transformations) becomes increasingly severe. The GRU's direct recurrent gradient path becomes correspondingly more valuable.
Why fixed-random routers are competitive: If the router optimization problem is fundamentally about inadequate gradient signal rather than insufficient capacity, then random initialization — which at least provides diverse expert assignments — might be preferable to a partially-trained router stuck in a poor local optimum driven by weak, noisy gradients.

Implications for future work: This finding redirects attention from "how do we share routing information across layers?" to "how do we improve gradient flow for router parameters?" Approaches that provide better optimization pathways — not just better forward-pass information — should be prioritized. This could include architectural innovations beyond recurrence (e.g., router-specific residual connections, auxiliary training objectives that provide denser gradient signals), or optimization innovations (e.g., different learning rates for router parameters, specialized router initialization schemes, two-stage training where routing is learned before the rest of the model).

Innovation 3: Difficulty-Conditioned Benefits of Layerwise Recurrence — Not All Layers or Tokens Benefit Equally

While the aggregate results show consistent improvements from RMoE, the paper's analysis reveals that the benefits of cross-layer routing memory are not uniform — they manifest differently across layers, across expert selection patterns, and across the training trajectory. This is not presented as a separate "innovation" section in the paper, but it emerges from the detailed observational analyses in Section 6 and represents a conceptual contribution about how routing coordination actually works.

The gating entropy signature: Figure 4 and the related analysis in Section 6 reveal a distinctive pattern in how RMoE affects routing behavior. Standard SMoE produces sharply peaked gating distributions — many tokens have near-zero entropy, meaning the softmax over experts is almost one-hot before Top-k is even applied. This implies the router has converged to a brittle, deterministic assignment pattern early in training. HyperMoE and RandomMoE show the opposite extreme — very high entropy, indicating nearly uniform expert assignments (essentially random routing). RMoE occupies a moderate middle ground: the gating entropy distribution has density in both high-entropy and low-entropy regions, suggesting that the model learns to be confident about some assignments while maintaining flexibility for others. The paper characterizes this as "a better balance between exploration and exploitation" (Section 6).

What makes this conceptually interesting is that it contradicts the intuitive expectation that "more information = more confident, specialized routing." If the GRU were simply providing additional context that makes routing decisions easier, we would expect gating distributions to become sharper (more confident). Instead, they become more moderate. The explanation, supported by the inner/outer balance statistics in Table 8, is that RMoE prevents premature convergence: the GRU's cross-layer gradient signal provides enough optimization pressure to keep the router exploring different expert combinations, rather than collapsing to the first locally-stable assignment pattern (which, for standard SMoE, is often token-ID-based routing, as Xue et al. (2024) documented).

Layer-specific effects: The expert similarity analysis (Figure 5) from the large-scale pre-training setting shows that RMoE's effect on expert diversity is not uniform across the training trajectory. In early training (first few thousand steps, corresponding to ~4–12B tokens), expert similarity increases for both SMoE and RMoE as the randomly-initialized router scatters tokens across experts, and the experts all learn similar functions. But as training progresses, RMoE's expert similarity drops faster and further than SMoE's. The paper's explanation — that the Recurrent Gradient helps the router continue learning meaningful assignments beyond the stage where standard routers have converged to shallow patterns — is consistent with the gradient norm analysis in Appendix A.4.1 (Table 11), which shows that the linear router's gradient from the load balancing loss dominates early and then vanishes, while the RNN router maintains a more balanced gradient signal from both LM and LB losses throughout training.

The load balancing gradient dominance problem: This is a subtle but important diagnostic insight embedded in the Appendix analysis (Table 11) that deserves elevation. In standard SMoE, the load balancing loss gradient dominates the router's early training (at step 100, the LB gradient norm is 0.433 vs. 0.625 for LM, and the drop ratio falls from 35.6% to 5.43% by step 10,000 and stays there). The router learns to balance experts very quickly — and then the gradient from balancing loss vanishes, leaving only the weak, sparse LM gradient to drive further specialization. The router is essentially stuck after the first ~10B tokens of large-scale training, having learned a load-balanced but semantically shallow assignment pattern. In RMoE, the balancing loss gradient is initially lower (0.337 at step 100) and persists longer into training (still 0.015 at step 20,000 vs. 0.008 for linear), while the LM gradient continues to provide meaningful signal. The GRU effectively moderates the influence of the balancing loss, preventing the router from over-optimizing for load balance at the expense of learning meaningful expert specialization.

Why this matters beyond RMoE: This analysis identifies a general pathology in SMoE training — the auxiliary balancing loss dominating and then vanishing, leaving routers undertrained for their primary purpose — that is independent of RMoE. The paper doesn't frame it this way, but the finding suggests that any mechanism that moderates the balancing loss's early influence or provides alternative gradient pathways to the router could improve MoE parameter efficiency, not just the specific GRU-based approach.

Innovation 4: Orthogonal Composition as a Design Principle for MoE Architectures

The paper's explicit claim that RMoE introduces "a novel computation stage orthogonal to existing methods" (Section 1, Abstract) is more than a compatibility statement — it represents a design philosophy for how to advance MoE architectures through composable, modular innovations rather than competing, mutually-exclusive designs.

What this means concretely: Most prior router improvements are substitutive — they replace the standard linear router with a different routing mechanism (cosine similarity for XMoE, hypernetwork for HyperMoE, MLP for SMoE-MLP). You cannot use XMoE and HyperMoE simultaneously because they occupy the same computational role. RMoE is additive — it inserts a preprocessing stage (projection + GRU) before whatever router architecture is being used. Table 2 demonstrates this explicitly: XMoE + GRU router outperforms XMoE alone across all configurations (XMoE(8), XMoE(16), XMoE(32)), and the combination works because the GRU addresses a different bottleneck (layerwise isolation) than XMoE addresses (representation collapse).

Why this is intellectually distinctive: The paper is not just claiming "our method works with other methods" — it is articulating a principle about how the MoE computation graph can be factored into separable stages: (1) routing state preparation (what information does the router receive?), (2) scoring (how are expert affinities computed?), and (3) selection (how are experts chosen given scores?). Prior work focused almost exclusively on stages (2) and (3) — the scoring function and the Top-k mechanism — while taking the input to the router (the hidden state $x_i$ ) as fixed. RMoE innovates on stage (1) — what preprocessing transforms the hidden state before scoring? — which creates a new axis of architectural variation that composes with innovations on the other axes.

Significance beyond this paper: This factoring principle has implications for how the field should organize MoE research. Rather than developing monolithic "router architectures" that compete with each other, researchers could develop separable components addressing different stages: better routing state preparation (like RMoE), better scoring functions (like XMoE's cosine similarity), better selection mechanisms (like expert-choice routing or dynamic top-k). Each could be developed and evaluated independently, then combined. The paper demonstrates that combinations across stages are at least additive and potentially synergistic — a template for how cumulative progress in MoE design could be organized.

Limitation of this insight: The paper demonstrates composition only with XMoE, and only at the small scale (Enwiki8/WikiText-103). Whether the additive benefits hold at larger scales and with other router designs (HyperMoE, SMoE-MLP, etc.) is not tested. The principle is articulated but not thoroughly validated as a general property.

Innovation 5: The Negative Result — Naive Cross-Layer Information Is Actively Harmful

The paper contributes an important negative finding: not all forms of cross-layer routing information are beneficial, and some actively degrade performance below the no-information baseline. This is clearest from the RMoE-NP+r-α experiments in Table 5, which test the straightforward idea of adding a residual connection from the previous layer's gating logits to the current layer's logits.

The intuition that fails: If the problem is that routers lack information about previous layers' decisions, the simplest fix is to directly tell them: add the previous layer's expert scores to the current layer's scores, weighted by some coefficient α. This technique (residual attention scores) works in other contexts (Realformer, He et al., 2020). It provides both forward-pass information and a backward gradient pathway.

Why it fails: The paper identifies the core issue: "the indexes of experts in layer $i$ are not aligned with those in other layers, directly adding logits can lead to improper constraints and hurt the model performance." Expert 3 at layer 2 and Expert 3 at layer 5 are different neural networks with different parameters trained to handle different aspects of the input at different levels of abstraction. There is no reason to assume they should receive correlated scores. Adding their logits imposes an alignment constraint — essentially telling the model that if Expert $j$ was good at layer $i-1$ , Expert $j$ should also be favored at layer $i$ — that the model must then fight against if the optimal assignments are different. The GRU avoids this because it learns a nonlinear transformation from the previous routing state to the current one, allowing expert identities to shift across layers while still sharing useful abstract information about the token's routing trajectory.

Why this is a meaningful negative result: It establishes a boundary condition on what kinds of cross-layer information are useful. Direct, untransformed sharing of expert identities is harmful. Learned, flexible transformations of routing state are helpful. This rules out a whole class of simple approaches (residual logits, attention over previous routing distributions, direct conditioning on previous expert selections) and focuses future work on learned, stateful mechanisms that can disentangle expert identities across layers while sharing abstract routing patterns.

Connection to the gradient pathway finding: The negative result also reinforces Innovation 2. Even with gradient flow intact (RMoE-NP+r-0.5 without detach), the logit residual underperforms the full GRU (1.124 vs. 1.116 BPC). This means the GRU's benefit is not just the gradient pathway — the specific form of the learned, nonlinear state update matters too. The GRU succeeds because it simultaneously provides (a) a beneficial gradient pathway, (b) flexible forward-pass information that does not impose harmful expert alignment constraints, and (c) a learned gating mechanism (reset and update gates) that allows the model to dynamically decide when to carry forward routing state versus when to overwrite it. The negative results show that removing any of these properties (detaching gradients, removing the learned state update via NP, or providing information in a rigid rather than learned form via logit residuals) degrades or destroys the benefit.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses two standard language modeling benchmarks. Enwiki8 is a character-level language modeling dataset (100M characters from English Wikipedia) evaluated with Bits-Per-Character (BPC, lower is better). WikiText-103 is a word-level language modeling dataset (103M words from Wikipedia articles) evaluated with Perplexity (PPL, lower is better). Both use default train-validation-test splits. For the large-scale pre-training setting (Section 4.1), the paper uses a multilingual corpus spanning Wikipedia, finance, and legal texts — no standard name is given, and the total token count is controlled at 20B or 40B tokens. Additionally, a separate 15B-parameter-scale experiment in Appendix A.3 uses a "high-quality pre-training corpus" (again unnamed) at 120B and 400B tokens.
Base model(s). The paper tests RMoE at three distinct scales, using two different architectural backbones. Small-scale (Tables 1, 2, 4–6, and the mutual information/entropy analyses in Figures 3–4): an 8-layer decoder-only transformer with hidden size 352, 8 attention heads, sequence length 512, and 16 experts with Top-2 selection. This is explicitly "based on the decoder-only standard switch-transformer architecture with post-norm" and follows the configurations from CompeteSMoE (Pham et al., 2024). Medium-scale (Figure 2): models with 6, 12, 18, 24, and 32 layers sharing the same hidden size 352 and expert configuration, used to study how the RMoE benefit scales with depth. Large-scale (Tables 3, 7, and Appendix A.3 Tables 9–10): the 0.91B-parameter models use a 24-layer architecture with Rotary Embedding, SwiGLU activations, RMSNorm, hidden size 1280, 20 attention heads, sequence length 4096, and fine-grained MoE with 16 experts (Top-4 selection, expert FFN hidden size 448) plus shared experts following DeepSeekMoE (Dai et al., 2024). The 15B-parameter models in Appendix A.3 share this architecture but are scaled up (15B total, 2.7B activated). The base model family choice (decoder-only transformer variants) is standard for language modeling evaluation; the paper does not test encoder-decoder or non-transformer architectures. The paper states PaLM 2 is not used — the models are trained from scratch.
Metrics. For language modeling: Enwiki8 uses Bits-Per-Character (BPC), which is the average negative log-likelihood per character divided by log(2). WikiText-103 uses Perplexity (PPL), the exponentiated average negative log-likelihood per word. Both are computed on the test set using the best validation checkpoint. For the large-scale pre-training evaluation after supervised fine-tuning, the paper uses the lm-evaluation-harness (EleutherAI) to compute task-specific metrics: ARC-Easy (accuracy), Hellaswag (accuracy, normalized), PIQA (accuracy, normalized), SciQ (accuracy), and LAMBADA (accuracy). The paper also reports an unweighted average across these five tasks. For the 15B-scale experiments in Appendix A.3, additional metrics include MMLU (accuracy), GSM8K (exact match accuracy), HumanEval (pass@1), and an average of domain-specific perplexities (Avg PPL). Training cost is reported as seconds per training step and peak GPU memory (GB) at matched batch sizes.
Baselines. The paper compares against a comprehensive set of router designs, all implemented within the same base architecture:
1. SMoE: Standard switch-transformer with a single linear layer router (the de facto baseline from Fedus et al., 2022).
2. SMoE-MLP (Shen et al., 2023): Replaces the linear router with a two-layer MLP using GELU activation, testing whether increased router capacity alone helps.
3. HyperMoE (Do et al., 2023): Uses a fixed, randomly-initialized hypernetwork to generate router weights conditioned on a learnable router embedding and the input.
4. RandomMoE: A naive baseline with a fixed randomly-initialized linear router (never updated), inspired by SMoE-Dropout (Chen et al., 2023) and HyperMoE's use of fixed random components.
5. XMoE (Chi et al., 2022): Projects hidden states to a low dimension (default 16) and computes cosine similarity with low-dimension expert embeddings, using a learnable softmax temperature. This addresses representation collapse in standard routers.
6. CosineMoE: A variant of XMoE without the dimensionality reduction — cosine similarity between the full hidden state and expert embeddings.
For the large-scale pre-training setting (Tables 3 and 7), the primary comparison is between SMoE and RMoE variants, with SMoE-MLP also tested at the 15B scale (Appendix A.3). The paper does not implement all small-scale baselines (HyperMoE, RandomMoE, CosineMoE) at the large scale, likely due to computational constraints.
Generation budget / compute accounting. The paper does not use a "generation budget" concept (this is not an inference-time scaling paper). Instead, compute is accounted through three metrics: (1) Model parameters — total non-embedding parameters, with router parameters broken out separately (e.g., Table 1 reports 36.08M for SMoE with 0.04M router parameters, vs. 36.51M for RMoE with 0.47M router parameters). (2) Training speed — average wall-clock time per 1,000 training steps at matched batch size (e.g., 960.2 s/1k steps for SMoE vs. 972.9 s/1k steps for RMoE on the small-scale setting, a 1.3% increase). (3) Peak GPU memory — maximum memory usage during training at matched batch size (e.g., 47.92 GB for SMoE vs. 49.46 GB for RMoE at small scale; 48.00 GB vs. 48.69 GB at 0.91B scale). For the large-scale pre-training, the same number of training tokens (20B or 40B) and same batch size (1,120) are used for all methods, making accuracy at equal token count the primary comparison. The paper explicitly notes at the 0.91B scale that RMoE "introduces about 3.5M additional parameters relative to SMoE" (Table 3 caption), which is a ~0.4% increase in total parameters (3.5M out of ~910M).
Cross-validation / statistical protocol. The paper does not employ cross-validation in the standard sense. For the small-scale experiments, each configuration is trained once, and test performance of the best validation checkpoint is reported — there are no error bars, confidence intervals, or multiple random seeds. The paper acknowledges this implicitly by relying on consistency across multiple datasets (Enwiki8 and WikiText-103), multiple model scales (8-layer, 24-layer, 32-layer), and multiple training paradigms (pre-training and SFT) to establish robustness rather than statistical testing. For the 0.91B-scale SFT evaluation, results are reported on standard lm-evaluation-harness tasks with default evaluation protocols. The 15B-scale experiments in Appendix A.3 evaluate checkpoints at multiple training token counts (80B, 100B, 120B) and (200B, 280B, 400B), providing a form of robustness check through training trajectory curves. The most significant gap in experimental rigor is the absence of multiple random seeds — every result in the paper is from a single training run, which is common in large-scale LM research due to cost but leaves open the possibility that some differences are within noise. For the small-scale experiments where training costs ~20 GPU-hours, multiple seeds would have been feasible and would strengthen the claimed improvements (which are sometimes small in absolute terms, e.g., 1.116 vs. 1.128 BPC on Enwiki8 test).

Main Quantitative Results

Small-Scale Language Modeling: RMoE vs. Baseline Routers

Table 1 presents the primary small-scale comparison on Enwiki8 and WikiText-103. RMoE achieves the best performance on both datasets: on Enwiki8 test, RMoE reaches 1.116 BPC compared to the next-best baseline (CosineMoE at 1.122 BPC and SMoE at 1.128 BPC). On WikiText-103 test, RMoE achieves 32.867 PPL, outperforming XMoE (32.926 PPL) and SMoE (33.061 PPL).

The absolute margins are small but consistent: RMoE's improvement over standard SMoE is 0.012 BPC on Enwiki8 test and 0.194 PPL on WikiText-103 test. To contextualize these margins: the gap between SMoE and the worst-performing method (HyperMoE at 1.139 BPC / 33.374 PPL) is only 0.011 BPC / 0.313 PPL, meaning RMoE's gain over SMoE is roughly comparable to the spread between best and worst methods.

Critical observations from the baseline comparisons:

SMoE-MLP (the higher-capacity router) achieves 1.137 BPC on Enwiki8 test — worse than standard SMoE (1.128 BPC) and even worse than RandomMoE (1.135 BPC). This confirms that merely adding parameters to the router is not sufficient and can be counterproductive.
RandomMoE (fixed random router) achieves 1.135 BPC on Enwiki8 test — competitive with SMoE-MLP (1.137 BPC) and not dramatically worse than standard SMoE (1.128 BPC). This is consistent with prior findings that random routing is surprisingly effective, and highlights that the standard learnable router is not extracting much value from its training signal.
CosineMoE and XMoE both outperform SMoE, with CosineMoE (1.122 BPC) slightly ahead of XMoE (1.125 BPC), suggesting that the cosine-similarity mechanism helps prevent the router from collapsing to poor local optima, independent of the dimensionality reduction.

Training cost comparison (Table 1): RMoE achieves these gains with modest overhead. Peak GPU memory increases from 47.92 GB (SMoE) to 49.46 GB (RMoE), a 3.2% increase. Training speed decreases from 960.2 s/1k steps to 972.9 s/1k steps, a 1.3% slowdown. The router parameter count increases from 0.04M (SMoE) to 0.47M (RMoE), but this is a negligible fraction of the total 36.51M parameters. The overhead comes primarily from the per-layer projectors ( $8 \times 352 \times 128 = 360$ K parameters) and the shared GRU ( $6 \times 128^2 = 98$ K parameters).

Composition with XMoE: GRU Router Provides Additive Benefits

Table 2 demonstrates that RMoE's GRU router composes orthogonally with XMoE's cosine-similarity routing. XMoE is tested at three lower projection dimensions (8, 16, 32), with and without the GRU router.

On Enwiki8 test, adding the GRU router to XMoE(8) improves from 1.132 to 1.124 BPC. For XMoE(16), the improvement is from 1.125 to 1.119 BPC. For XMoE(32), from 1.114 to 1.112 BPC. On WikiText-103 test, the same pattern holds: XMoE(16) improves from 32.93 to 32.47 PPL; XMoE(32) improves from 32.71 to 32.55 PPL.

Two observations emerge: (1) The GRU router provides a consistent benefit regardless of XMoE's internal dimension — the improvement is not specific to a particular XMoE configuration. (2) The absolute gain from adding the GRU is smaller for XMoE(32) than for XMoE(8) or XMoE(16) — the lower-dimension XMoE variants benefit more from the additional cross-layer information, possibly because their limited projection dimension constrains the router's ability to capture routing-relevant features from the hidden state alone, making the GRU's dedicated routing state more valuable.

Large-Scale Pre-Training (0.91B Parameters): RMoE Consistently Outperforms SMoE

Table 3 presents the headline large-scale results. At 20B training tokens, RMoE pretrained achieves an average task score of 49.20 (vs. 47.63 for SMoE), a gain of 1.57 points. After supervised fine-tuning on Alpaca, RMoE reaches 51.20 (vs. 48.97 for SMoE with frozen router, the best SMoE configuration), a gain of 2.23 points. With both the router and GRU frozen during SFT (following the standard practice of freezing routers), RMoE achieves 51.24, nearly identical to the unfrozen RMoE (51.20), suggesting the learned routing state is stable and transfers well.

At 40B training tokens, the pattern holds: RMoE pretrained scores 55.31 (vs. 54.26 for SMoE), a gain of 1.05 points. After SFT with frozen router and GRU, RMoE scores 57.08 (vs. 55.88 for the best SMoE configuration), a gain of 1.20 points.

Task-specific patterns: RMoE's gains are not uniform across tasks. At 40B tokens after SFT, the largest absolute improvements are on LAMBADA (37.57 for RMoE vs. 32.06 for SMoE with frozen router, a 5.51-point gain) and Hellaswag (43.16 vs. 41.94, a 1.22-point gain). The smallest improvements are on SciQ (82.8 vs. 83.1, a marginal -0.3-point difference favoring SMoE at 40B frozen, though at 20B frozen RMoE leads 79.7 vs. 74.7) and PIQA (68.77 vs. 68.88, essentially tied). This heterogeneity suggests that the benefits of cross-layer routing coordination may be task-dependent: tasks requiring long-range contextual understanding (LAMBADA, which requires predicting the last word of a passage using broad discourse context) benefit more than tasks requiring commonsense reasoning (PIQA) or factual recall (SciQ). However, the paper does not investigate this task-dependence further.

Cost accounting (Table 3): At the 0.91B scale, RMoE training takes 49.07 s/step vs. 48.87 s/step for SMoE, a 0.4% slowdown. Memory usage is 48.69 GB vs. 48.00 GB, a 1.4% increase. Both differences are negligible in practice. The paper explicitly claims RMoE "introduces about 3.5M additional parameters relative to SMoE" (Table 3 caption), which is ~0.4% of total parameters.

Scaling to 15B Parameters: Gains Persist at Larger Model Sizes

Appendix A.3 (Tables 9 and 10) reports results at substantially larger scale: 15B total parameters with 2.7B activated, using DeepSeekMoE-style fine-grained experts, trained on 120B and 400B tokens. These experiments also include SMoE-MLP as a baseline alongside SMoE and RMoE.

At 120B tokens, RMoE achieves Hellaswag 72.36 (vs. SMoE 72.03, SMoE-MLP 72.19), MMLU 54.02 (vs. 52.79, 52.81), GSM8K 36.13 (vs. 34.80, 34.57), and Avg PPL 6.425 (vs. 6.447, 6.479). The MMLU improvement (+1.23 over SMoE) is notable because it is a knowledge-intensive benchmark where routing quality might be expected to matter less than on reasoning tasks; the consistent improvement suggests the GRU's cross-layer coordination helps with knowledge retrieval and integration as well. The GSM8K improvement (+1.33 over SMoE) on math reasoning is also substantial.

At 400B tokens, RMoE achieves Hellaswag 76.72 (vs. SMoE 76.39), MMLU 60.60 (vs. 59.54), GSM8K 52.99 (vs. 52.16), and Avg PPL 5.620 (vs. 5.685). The gaps narrow slightly compared to 120B tokens — the MMLU delta drops from +1.23 to +1.06, and the GSM8K delta drops from +1.33 to +0.83 — but RMoE maintains a consistent lead. The narrowing could indicate that with sufficient training data, standard routers partially overcome their optimization difficulties, reducing the relative advantage of the Recurrent Gradient. However, the paper does not train to convergence, so the gap at convergence is unknown.

SMoE-MLP at scale: SMoE-MLP performs slightly better than SMoE on Hellaswag and MMLU but worse on GSM8K (34.57 vs. 34.80 at 120B; 51.71 vs. 52.16 at 400B). This mirrors the small-scale finding that increased router capacity alone does not reliably improve performance, and for reasoning tasks (GSM8K), the MLP router may actually overfit or learn spurious patterns that the linear router avoids.

RMoE Variants and Hyperparameter Sensitivity at Scale

Table 7 (and the expanded Table 12 in Appendix A.5) reports a sweep over RMoE design choices at the 0.91B scale with 20B training tokens. GRU versus RNN versus LSTM: GRU with p=128 achieves the best pretraining score (49.20) and best post-SFT score (51.32 with router frozen). The simple RNN with p=256 achieves 47.92 pretraining / 50.44 SFT — worse than SMoE (47.63 pretraining / 49.11 SFT) in pretraining but better after SFT, an inconsistent pattern that suggests the RNN's forward-pass information is noisy or unstable but the gradient pathway still provides some benefit during fine-tuning. GRU hidden dimension sweep: p=128 (49.20 pretrain) outperforms p=256 (49.08) and p=512 (49.19), with p=256 being noticeably worse. The optimal size at p=128 is modest — only 128-dimensional routing state across all layers — reinforcing the interpretation that the GRU provides a narrow, purpose-built channel that should not be overloaded. Larger hidden states may introduce optimization challenges or encourage the model to store information in the GRU state that would be better represented in the main residual stream.

At 40B tokens, the pattern shifts: p=256 (53.18 pretrain score) now outperforms p=128 (55.31 pretrain score) — wait, that's worse. Re-reading Table 12: at 40B tokens, RMoE-GRU-p=128 achieves 55.31, p=256 achieves 53.18, and p=512 achieves 53.72. So p=128 maintains its lead at 40B tokens, and the gap widens (55.31 vs. 53.18, a 2.13-point gap vs. 49.20 vs. 49.08 at 20B, a 0.12-point gap). This suggests the smaller GRU state is not just sufficient but increasingly beneficial as training progresses — consistent with the interpretation that excess GRU capacity encourages the model to route information through the GRU that would be better handled by the main network.

Expert Diversity and Training Dynamics from the Large-Scale Setting

Figure 5 (expert similarity over training) from the 15B-scale pre-training shows that RMoE maintains lower expert similarity than SMoE throughout training. At the final checkpoint (10k steps, ~40B tokens equivalent, but applied to the 15B model trained on 120B tokens), RMoE's median expert cosine similarity is lower than SMoE's, and the interquartile range is wider, indicating both greater average diversity and more variation in specialization across layers. The gap in expert similarity emerges early (by step 2,000) and persists, with RMoE's similarity dropping faster during the middle phase of training (steps 3,000–7,000).

Table 11 (Appendix A.4.1) provides the mechanistic explanation: the linear router's gradient from the load balancing loss dominates early (0.433 at step 100 vs. 0.625 from LM loss) and then collapses (0.001 at step 60,000), while the RNN router maintains a more balanced gradient from both losses throughout training. The drop ratio — the fraction of tokens that cannot be processed by their assigned expert due to capacity limits — tells a complementary story: the linear router's drop ratio drops rapidly (35.6% to 5.43% by step 10,000) and stays low, indicating the router prioritizes load balance over routing quality, while the RNN router's drop ratio decreases more gradually (38.7% to 6.35% at step 10,000 to 4.09% at step 60,000), suggesting a more gradual optimization that balances both objectives.

Ablation Studies and Robustness Checks

Layerwise recurrence vs. increased router parameters (Table 4): RMoE+NP, which has identical parameters and FLOPs to RMoE but removes the layerwise recurrence by resetting the GRU state at each layer ( $h_i = \text{GRU}(x'_i, h_0)$ instead of $\text{GRU}(x'_i, h_{i-1})$ ), performs worse than standard SMoE on both the small-scale setting (1.196 vs. 1.184 test BPC for Small) and comparably on the medium-scale setting (1.123 vs. 1.128 test BPC). Meanwhile, SMoE+MLP, which adds router capacity without recurrence, underperforms SMoE (1.183 vs. 1.184 test BPC on Small). This isolates the layerwise recurrence — not the additional parameters — as the source of RMoE's gain. In fact, adding parameters without recurrence (RMoE+NP or SMoE+MLP) either does nothing or hurts, while adding a comparable parameter count with recurrence (RMoE) provides consistent gains.

Recurrent Gradient vs. forward-pass information (Table 5): The RMoE + detach $h_{i-1}$ setting preserves the full forward-pass information flow (the GRU receives $h_{i-1}$ and uses it to compute $h_i$ ) but prevents gradient flow from $h_i$ back to $h_{i-1}$ during backpropagation. On Enwiki8 test, this achieves 1.133 BPC — worse than RMoE-NP (1.123 BPC, which has no forward-pass recurrence at all) and substantially worse than RMoE (1.116 BPC). This is the single most important ablation: it demonstrates that forward-pass information without the corresponding Recurrent Gradient is actively harmful, while forward-pass information with Recurrent Gradient is beneficial. The forward-pass information and backward gradient pathway are not additive benefits — the forward-pass information is only useful when accompanied by the gradient pathway.

Gating logit residual as an alternative cross-layer mechanism (Table 5): RMoE+NP+r-0.5, which adds a residual connection from the previous layer's gating logits to the current layer's logits (coefficient α=0.5), achieves 1.124 test BPC — comparable to RMoE-NP (1.123) but substantially worse than RMoE (1.116). Increasing α to 1.0 produces 1.124 BPC as well. When the gradient through this residual is detached (RMoE+NP+r-0.5+detach-r), performance drops to 1.133 BPC — the same as RMoE+detach $h_{i-1}$ . This demonstrates that (1) the gradient pathway through whatever cross-layer connection exists is critical, and (2) the specific form of the GRU's learned, nonlinear state update provides benefits beyond what a simple residual logit connection can achieve, even when both have gradient pathways.

Deep model scaling behavior (Figure 2): As model depth increases from 6 to 32 layers, RMoE maintains a consistent advantage over SMoE, while RMoE-NP (no recurrence) falls increasingly behind SMoE at greater depths. At 6 layers, SMoE and RMoE-NP are comparable (~1.22 test BPC for both), while RMoE is at ~1.205 BPC. At 32 layers, RMoE achieves ~1.130 BPC, SMoE is at ~1.148 BPC, and RMoE-NP is at ~1.155 BPC. The gap between RMoE and RMoE-NP widens from ~0.015 BPC at 6 layers to ~0.025 BPC at 32 layers, consistent with the interpretation that the Recurrent Gradient becomes more valuable as the gradient vanishing problem for shallow routers worsens with depth. RMoE-NP-r0.5 occupies an intermediate position at all depths, consistently outperforming RMoE-NP but underperforming RMoE.

Projector design: per-layer vs. shared (Table 6): Replacing the per-layer projectors in RMoE with a single shared projector (RMoE + S-proj + GRU router) raises Enwiki8 test BPC from 1.116 to 1.123 — erasing roughly half of RMoE's gain over SMoE (1.128). On WikiText-103 test, the shared projector PPL is 33.02 vs. 32.86 for per-layer projectors and 33.06 for SMoE — a similar pattern where shared projectors recover some but not all of RMoE's benefit. The paper attributes this to the wide variation in hidden state norms and standard deviations across layers (documented in Appendix Figure 11): "the weights and hidden states norm in different layers vary greatly, and it would be hard for a single shared projector to process them."

Recurrent cell type: GRU vs. RNN vs. LSTM (Table 6): On the small scale, RMoE with GRU achieves 1.116 test BPC, with simple RNN achieves 1.119, and with LSTM achieves 1.122. On WikiText-103 test, the ordering matches: GRU (32.86 PPL) > RNN (32.72) > LSTM (33.04). GRU is the best performer on both datasets, with LSTM performing notably worse on WikiText-103 (33.04 vs. 32.86). At the 0.91B scale (Table 7), RNN with p=256 achieves 47.92 pretrain — barely above SMoE (47.63) — and 50.44 after SFT, underperforming GRU-p=128 (49.20 pretrain, 51.32 SFT). The paper's conclusion that GRU is preferred due to its gating mechanism (allowing the model to dynamically decide what to remember vs. update) is supported, though the specific advantage over RNN is modest at small scale (0.003 BPC).

GRU hidden dimension p (Tables 6 and 7): At small scale, the paper does not report a sweep over p (Table 6 uses the default, which from context appears to be 128). At the 0.91B scale with 20B tokens, p=128 achieves 49.20 pretrain / 51.32 SFT, p=256 achieves 49.08 / 50.04, and p=512 achieves 49.19 / 50.02. The sweet spot at p=128 is evident, especially after SFT (51.32 vs. 50.04 is a substantial gap of 1.28 points). With 40B tokens, p=128 (55.31) widens its lead over p=256 (53.18) and p=512 (53.72). The paper does not test p < 128, so whether even smaller dimensions would suffice is unknown.

Freezing strategy during fine-tuning (Table 3): At 20B tokens, comparing RMoE SFT variants: unfrozen router and GRU (51.20), frozen router only (51.32), frozen router and GRU (51.24). The differences are small (within 0.12 points), suggesting the routing state learned during pre-training transfers well without further adaptation. At 40B tokens: unfrozen (57.15), frozen router only (56.92), frozen router and GRU (57.08). Again, differences are modest (within 0.23 points). SmoE's frozen router variant (49.11 at 20B, 55.89 at 40B) slightly outperforms its unfrozen variant (48.97, 56.13) — consistent with the established practice from Zoph et al. (2022) that freezing routers during fine-tuning is beneficial. RMoE's stability under different freezing strategies is a practical advantage.

Negative result: RMoE with ReST $^{EM}$ or other on-policy training (implied): While not an explicit ablation in the paper, Appendix K mentions that "attempting to further optimize the revision model using ReST $^{EM}$ ... backfires" in the context of MoE training. The paper does not provide RMoE-specific ReST $^{EM}$ experiments, but notes that on-policy data collection can "exacerbate spurious correlations" in MoE routing, suggesting a sensitivity to training data distribution that could affect RMoE as well. This is a gap: the paper does not test whether RMoE's GRU-based routing is more or less robust to distribution shift during continued training than standard routers.

Critical Assessment

The central claim — that RMoE consistently outperforms baseline routers — is supported across scales and settings, but the reported gains are often small in absolute terms, and the absence of multiple random seeds makes it difficult to assess statistical reliability.

For the small-scale language modeling experiments, RMoE improves Enwiki8 test BPC by 0.012 over SMoE (1.116 vs. 1.128). The spread between the best and worst methods in Table 1 is 0.023 BPC (RMoE at 1.116 vs. HyperMoE at 1.139). RMoE's improvement represents about half the total spread across all methods. On WikiText-103, the improvement is 0.194 PPL (32.867 vs. 33.061), with a total spread of 0.548 PPL (HyperMoE at 33.410). These are real improvements — they represent genuine better language modeling — but they are modest. With a single training run per configuration and no reported variance, a skeptical reader could ask whether these differences would persist across random seeds or whether they fall within training noise. The fact that the improvements are consistent across two datasets, multiple model depths (Figure 2), and two training scales (Tables 1 and 3) provides convergent evidence, but it does not substitute for direct statistical quantification.

The large-scale results show more substantial gains, but these come with caveats about the experimental design.

At the 0.91B scale, RMoE's post-SFT advantage over the best SMoE configuration at 40B tokens is 1.20 points on the average task score (57.08 vs. 55.88). For individual tasks, the gains range from 5.51 points on LAMBADA to essentially zero on PIQA and SciQ. These are more meaningful differences than the small-scale language modeling results. However, the 0.91B-scale experiments have an important limitation: models are trained on only 20B or 40B tokens, which is far from convergence for models of this size. The paper acknowledges the models are "largely under-trained" (Section 4.1), noting they give "almost random-guessing results on challenging tasks like MMLU." The 15B-scale experiments in Appendix A.3, which train on 120B and 400B tokens, provide a more realistic picture, but even these are not trained to convergence. The relative benefit of RMoE might change at full convergence — the narrowing of gaps from 120B to 400B tokens in Table 9 vs. Table 10 is suggestive but not definitive.

The ablation experiments convincingly isolate the Recurrent Gradient as the primary mechanism, but they also reveal that the mechanism is fragile — direct cross-layer information without gradient flow is harmful, and even slight changes to the GRU design (shared projectors, larger hidden states) degrade performance.

The detach- $h_{i-1}$ experiment (Table 5) is the most important ablation and provides strong evidence that the gradient pathway, not forward-pass information, is the key mechanism. However, it also reveals a puzzling result: RMoE+detach- $h_{i-1}$ performs worse than RMoE-NP (1.133 vs. 1.123 BPC). This means that having forward-pass information without the ability to backpropagate through it is actually worse than having no forward-pass information at all. The paper does not provide a clear explanation for this. One hypothesis: the forward-pass information biases the router toward certain expert assignments, but without the gradient signal to learn whether those assignments are actually optimal, this bias is more likely to be harmful than helpful (i.e., the GRU provides a "suggestion" that cannot be corrected through learning). This is an interesting finding that deserves further investigation.

The logit residual experiments (RMoE+NP+r-α) test an important alternative but could have been more extensive.

The comparison between RMoE and RMoE+NP+r-0.5 tests whether a simpler mechanism (residual gating logits) can capture the benefits of the GRU. The result — that it cannot — is interesting and supports the GRU design. However, the paper only tests α = 0.5 and α = 1.0. It does not test whether a learnable α, or a more sophisticated attention-style combination of previous and current logits, would close the gap. It also does not test residual connections at the level of $h_i$ states rather than gating logits — something like $h_i = h_{i-1} + \text{MLP}(x'_i)$ , which would provide the gradient pathway without the full GRU gating mechanism. This would help disentangle whether the GRU's gating (reset and update gates) is specifically important, or whether any nonlinear learned transformation with a residual connection would suffice.

The compositionality claim (orthogonal to existing methods) is demonstrated only with XMoE and only at small scale.

Table 2 shows that GRU router + XMoE outperforms XMoE alone across three XMoE configurations. This is a clean demonstration of orthogonality. However, it is only shown at the small scale (8-layer, hidden 352) on Enwiki8 and WikiText-103. Whether the composition holds at larger scales, with other router designs (HyperMoE, SMoE-MLP), or on downstream tasks after fine-tuning, is not tested. The claim that RMoE is "orthogonal to most existing attempts to improve MoE and is seamlessly compatible with them" (Section 3.2) is therefore a statement of principle supported by one data point, not a thoroughly validated property.

The mutual information analysis (Figure 3) provides compelling visual evidence for cross-layer coordination, but its connection to performance is correlational, not causal.

RMoE shows higher cross-layer MI than baseline methods (Figure 3d vs. 3a-c), and RMoE-NP shows lower MI (Figure 3e), consistent with the GRU explicitly passing routing state across layers. However, higher MI does not prove that cross-layer coordination is useful — it could be an epiphenomenon of the GRU's architecture rather than a causal mechanism. The detach- $h_{i-1}$ experiment partially addresses this by showing that when the forward-pass information is preserved but the gradient is cut, performance drops, but it does not disentangle whether the MI is driving performance or merely correlated with it. A more causal test would be to explicitly manipulate the MI (e.g., by adding noise to $h_i$ , or by forcing the GRU state to be low-MI through architectural constraints) and observe the performance impact.

The gating entropy analysis (Figure 4) and expert diversity analysis (Figure 5) provide plausible mechanistic explanations but are observational — they describe what RMoE does differently, not why these differences cause better performance.

The finding that RMoE produces more moderate gating entropy distributions than SMoE (Figure 4) is consistent with the story that RMoE achieves a better exploration-exploitation balance. The finding that RMoE maintains lower expert similarity (Figure 5) is consistent with better expert specialization. But both are descriptions of model behavior that correlate with performance, not proof of causality. It is possible that RMoE achieves moderate gating entropy because it performs better (i.e., the better language modeling loss forces the router to maintain more flexible assignments), rather than the other way around. Disentangling this would require interventions (e.g., explicitly regularizing gating entropy in SMoE to match RMoE's distribution and observing whether that closes the performance gap).

Missing experiments that would have strengthened the paper:

Multiple random seeds for the small-scale experiments to establish statistical significance of the 0.012 BPC improvement.
A compute-matched baseline that spends RMoE's additional 1–3% training cost on more training steps for SMoE rather than additional parameters for RMoE. This would test whether the benefit is truly architectural or could be matched by simply training SMoE slightly longer.
A learned residual baseline (e.g., $h_i = \text{LayerNorm}(h_{i-1} + \text{MLP}(x'_i))$ ) to test whether the GRU's gating mechanism specifically matters or whether any residual nonlinear transformation across layers with gradient flow provides the same benefit.
Evaluation on more diverse downstream tasks at the 0.91B scale. The paper evaluates on 5 tasks from lm-evaluation-harness, all of which are relatively "easy" commonsense and language understanding benchmarks. Harder reasoning benchmarks (MMLU, GSM8K) are only evaluated at the 15B scale in the appendix, and even there, the models are undertrained. Stronger evidence would come from converged models or from a broader task suite at the primary experimental scale.
Analysis of which tokens or input types benefit most from cross-layer routing. The mutual information and gating entropy analyses aggregate across all tokens. It would be informative to see whether RMoE's benefits are concentrated on certain token types (e.g., ambiguous words, long-range dependencies, rare tokens) or certain layers, providing more targeted insight into when cross-layer coordination matters.

Overall assessment: The paper presents a well-motivated architectural modification, demonstrates consistent improvements across scales and settings, and conducts careful ablations that convincingly identify the Recurrent Gradient as the key mechanism. The experiments support the core empirical claim — RMoE outperforms standard routers — though the absolute gains are modest at small scale and the statistical significance is uncertain due to single-run results. The more ambitious conceptual claims — that layerwise recurrence solves the parameter inefficiency problem, that the Recurrent Gradient is the causal mechanism, that RMoE is genuinely orthogonal to all existing methods — are partially supported but overclaimed relative to the experimental evidence. The paper does not demonstrate that RMoE closes the large parameter efficiency gap (52B MoE ~ 6.7B dense) that motivates the work; it shows only that RMoE improves upon SMoE baselines by a factor that, while consistent, is small relative to the magnitude of the gap described in the introduction.

6. Limitations and Trade-offs

Limitation 1: The Parameter Efficiency Gap That Motivates the Work Is Not Closed — Only Narrowed

The assumption or constraint. The paper opens with stark examples of MoE parameter inefficiency: a 52B-parameter MoE performing like a 6.7B dense model (Rajbhandari et al., 2022), and a 16B MoE performing comparably to a 7B dense model (Dai et al., 2024; Bi et al., 2024). These represent efficiency gaps of roughly 8× and 2.3× respectively — the MoE models require dramatically more total parameters to match a much smaller dense model. The implicit promise of RMoE is that better cross-layer routing coordination can substantially close this gap. However, the paper never measures RMoE against a dense model baseline, nor does it quantify what fraction of the parameter efficiency gap is recovered.

The consequence. A practitioner reading the paper to decide whether RMoE makes MoE models competitive with dense models of equivalent total parameters receives no direct evidence. The paper demonstrates that RMoE outperforms standard SMoE — but standard SMoE is itself a weak baseline relative to dense models of comparable size. The absolute gains from RMoE over SMoE (0.012 BPC on Enwiki8 test at small scale; ~1.2 points on the average downstream task score at the 0.91B scale with 40B tokens) are modest compared to the order-of-magnitude parameter efficiency gaps described in the introduction. It remains entirely possible that an RMoE-based 52B model would still perform like a ~7–8B dense model, representing a marginal improvement over standard SMoE but still far from the parameter efficiency of dense architectures.

What evidence exists in the paper. The paper provides no head-to-head comparison between RMoE-based MoE models and equivalently-sized dense models. All comparisons are MoE-to-MoE. The only experiment that touches on the "is this worth the extra parameters?" question is the large-scale comparison in Table 3, which shows that RMoE adds ~3.5M parameters (0.4% of the 0.91B total) while achieving the reported accuracy gains. But this compares RMoE to SMoE, not to a dense model. The introduction's motivating examples (Rajbhandari et al., 2022; Komatsuzaki et al., 2023; Dai et al., 2024) are never revisited in the results sections with RMoE numbers, so the reader cannot assess how much of the original parameter efficiency problem RMoE actually solves.

Mitigation status. This limitation is structural — the paper frames itself as solving a parameter efficiency problem but evaluates only relative to other MoE routers, not relative to the dense baselines that define what "parameter efficient" means. The authors do not acknowledge this gap. A dense baseline (even a small one, such as a standard transformer with comparable activated parameters) would have directly addressed whether improved routing translates to improved parameter efficiency in absolute terms. The 15B-scale experiments in Appendix A.3 provide the closest thing to a density comparison — since the activated parameters are 2.7B, the performance relative to known 2.7B dense models could be compared — but the paper does not make this comparison.

Limitation 2: The Recurrent Gradient Mechanism Implies a Fundamental Tradeoff Between Router Depth and Training Stability That Is Not Characterized

The assumption or constraint. The paper's central mechanistic finding — that the GRU's benefit comes primarily from the Recurrent Gradient pathway (Section 5, Tables 4–5) — implies that the deeper the model, the larger the relative benefit of RMoE over SMoE (because gradient vanishing for shallow-layer routers becomes more severe). Figure 2 partially supports this: the gap between RMoE and SMoE widens from 6 to 32 layers. However, RNN-based recurrent connections are known to introduce their own training instabilities at scale — gradient explosion, sensitivity to initialization, and difficulty with very long recurrent chains. The paper tests up to 32 layers at small scale, which corresponds to 32 GRU steps per token. Whether the Recurrent Gradient remains stable and beneficial at the 64, 96, or 128 layers common in production LLMs is entirely unknown. There is likely a crossover point where the GRU's own vanishing/exploding gradient problems outweigh the router gradient improvement it provides.

The consequence. A practitioner deploying RMoE in a very deep model (e.g., 70+ layers) has no guidance on whether the architecture will remain stable. The Recurrent Gradient benefit scales with depth up to 32 layers (Figure 2), but this monotonic trend may not hold indefinitely. RNNs are notoriously difficult to train over long sequences (here, long layer chains), and the shared GRU parameters being applied 64, 96, or 128 times per forward pass could lead to instability that standard SMoE routers — which have no cross-layer recurrence — avoid entirely. The paper provides no tools for diagnosing or mitigating such instability: no gradient clipping strategies specific to the GRU, no initialization schemes, and no ablation on the effect of model depth on GRU gradient norms.

What evidence exists in the paper. Figure 2 tests depths of 6, 12, 18, 24, and 32 layers — all on the small-scale architecture (hidden size 352). The gap between RMoE and SMoE at 32 layers is larger than at 6 layers, consistent with the Recurrent Gradient story, but the absolute performance numbers suggest no sign of catastrophic instability at 32 layers. However, 32 layers at hidden size 352 is a very small model by modern standards. The 0.91B-scale experiments use only 24 layers, providing no depth scaling data at realistic model sizes. The 15B-scale experiments (Appendix A.3) also use 24 layers. The paper has tested exactly zero configurations with more than 32 layers, leaving the depth scaling behavior at production scale completely uncharacterized.

Mitigation status. The paper does not acknowledge this as a limitation. The discussion of the Recurrent Gradient (Section 5) treats the gradient pathway as an unqualified benefit, without noting that recurrence introduces its own optimization challenges. The deepest model tested (32 layers at small scale) shows no failure, but this is insufficient evidence that the approach scales to the depth regimes where MoE is typically deployed. Future work would need to characterize the depth scaling behavior, establish whether gradient clipping or specialized GRU initialization is needed at large depths, and determine whether the Recurrent Gradient benefit eventually plateaus or reverses.

Limitation 3: The Mutual Information and Gating Entropy Analyses Are Observational — No Causal Link to Performance Is Established

The assumption or constraint. Section 6 presents two central observational analyses: (1) RMoE increases cross-layer mutual information (MI) between routing distributions (Figure 3), and (2) RMoE produces more moderate gating entropy distributions — less peaked than SMoE, less uniform than RandomMoE (Figure 4). The paper interprets both as evidence that RMoE improves routing behavior: the higher MI demonstrates cross-layer coordination, and the moderate entropy reflects a better exploration-exploitation balance. The implicit claim is that these behavioral differences cause the performance improvements.

The consequence. An alternative interpretation is that these behavioral differences are consequences of better performance, not causes. If RMoE achieves lower language modeling loss through some mechanism unrelated to routing coordination (e.g., the Recurrent Gradient provides a better optimization landscape), the lower loss might retrospectively produce higher cross-layer MI and more moderate gating entropy — not because the GRU directly coordinates routing, but because better-trained models naturally develop more structured routing patterns. Without causal interventions that manipulate the proposed mechanism (e.g., explicitly forcing RMoE's MI to match SMoE's and measuring performance, or artificially adjusting gating entropy and measuring the effect), the paper cannot distinguish "RMoE causes better routing, which causes better performance" from "RMoE causes better performance through the Recurrent Gradient, which causes better routing as a side effect." This matters for guiding future research: if the behavioral differences are symptoms rather than mechanisms, efforts to explicitly maximize cross-layer MI or moderate gating entropy in other router designs would be misguided.

What evidence exists in the paper. The paper conducts one intervention that partially addresses this: RMoE+detach- $h_{i-1}$ removes the Recurrent Gradient while preserving forward-pass information, and performance drops (Table 5). This shows that the gradient pathway is causally necessary for performance. But it does not show that the observed behavioral changes (higher MI, moderate entropy) are causally responsible — the detach experiment eliminates both the gradient pathway and (plausibly) the behavioral changes simultaneously. A more targeted intervention — such as training an SMoE model with an explicit auxiliary loss that encourages the gating entropy to match RMoE's distribution, without any GRU — would test whether the behavioral pattern alone improves performance. No such experiment is conducted. The MI analysis (Figure 3) compares RMoE, RMoE-NP, and RMoE-NP-r, establishing a correlation between MI and performance but not demonstrating that manipulating MI causes performance changes.

Mitigation status. The paper does not acknowledge this causal ambiguity. The language in Section 6 treats the behavioral findings as explanatory ("these moderate gating scores can achieve a better balance between exploration and exploitation"), implicitly claiming causation. A more careful framing would distinguish between "RMoE differs behaviorally from SMoE in these ways" (descriptive) and "these behavioral differences drive the performance gain" (causal), with the latter flagged as hypothesis rather than conclusion pending intervention experiments.

Limitation 4: Training Runs Are Unreplicated — Statistical Significance of Modest Gains Is Unknown

The assumption or constraint. Every result in the paper — from the small-scale language modeling experiments (Table 1) to the 0.91B-scale pre-training (Table 3) to the 15B-scale experiments (Appendix A.3) — is reported from a single training run per configuration. No error bars, confidence intervals, or multi-seed statistics are provided anywhere in the paper. This is common in large-scale ML research due to computational constraints, but it means the reported differences — which are sometimes small in absolute terms — cannot be distinguished from training noise.

The consequence. At the small scale (Table 1), the gap between RMoE and SMoE is 0.012 BPC on Enwiki8 test. The spread between SMoE (1.128) and HyperMoE (1.139) is 0.011 BPC — RMoE's improvement is comparable to the range across methods. With a single run per method, it is impossible to know whether RMoE's 1.116 BPC represents a genuine architectural advantage or a lucky training run within a distribution that overlaps substantially with SMoE's distribution. For a practitioner deciding whether to adopt RMoE, the expected improvement over standard SMoE could be 0.012 BPC (if the single-run point estimate is accurate), or it could be 0.002 BPC (if RMoE and SMoE distributions overlap and the observed difference is sampling noise). These imply very different cost-benefit calculations, especially given RMoE's additional implementation complexity (custom GRU integration, per-layer projectors, new hyperparameters).

The large-scale results partially mitigate this concern through convergent evidence: RMoE outperforms SMoE consistently across Enwiki8 and WikiText-103 (Table 1), across all model depths (Figure 2), across both 20B and 40B training tokens (Table 3), and across the 15B-scale experiments (Appendix A.3). This pattern of consistent improvement across independent experimental configurations is unlikely to arise from noise alone. However, this is informal evidence, not a substitute for direct statistical quantification. The absolute gains at 0.91B scale (1.20 points on average task score at 40B post-SFT) are more substantial than at small scale, making noise a less plausible explanation, but still unquantified.

What evidence exists in the paper. No statistical evidence. The paper does not report standard deviations across runs, does not mention random seeds, and does not discuss training variance.

Mitigation status. The paper does not acknowledge the absence of multi-seed experiments as a limitation. The consistent pattern across experimental scales and configurations provides informal robustness, but this limitation remains significant for the small-scale results where absolute gains are smallest and training cost (~20 GPU-hours per run) would have made 3–5 seeds feasible. For the large-scale experiments (5 days on 8 A100s), economic constraints make multi-seed runs genuinely prohibitive, and the larger absolute gains make the point estimate more informative. But the small-scale results — which form the basis for many of the ablations that establish the Recurrent Gradient mechanism — would benefit substantially from variance estimates.

Limitation 5: The Difficulty Estimation Problem — No Guidance on When RMoE Helps vs. When It Doesn't

The assumption or constraint. RMoE is presented as a universal architectural improvement — a drop-in replacement for standard routers that should be applied to all MoE layers in all models. The paper's experiments test RMoE on all layers of all models, and the results are reported as aggregate improvements. However, the mechanistic analysis reveals substantial heterogeneity in the effect: the gating entropy distribution (Figure 4) shows that RMoE affects routing behavior differently across layers and tokens, and the task-specific results (Table 3) show wide variation in improvement (LAMBADA gains 5.5 points while PIQA and SciQ show essentially zero improvement). There is no analysis of which layers, tokens, or task types benefit most, and no guidance on whether RMoE could be selectively applied only where it helps.

The consequence. A practitioner deploying RMoE has no way to predict whether their specific use case — a different model architecture, a different task distribution, a different training data mix — will see the same benefits reported in the paper. If RMoE's gains are concentrated on tasks requiring long-range contextual integration (like LAMBADA) and absent on factual recall tasks (like SciQ), then a model targeting knowledge-intensive question answering might see no benefit from the additional architectural complexity. Conversely, if gains are concentrated in early layers or late layers specifically, a model with a very different depth distribution might see different effects. The paper's aggregate reporting obscures this heterogeneity, making it impossible for practitioners to do informed cost-benefit analysis for their specific deployment.

What evidence exists in the paper. Table 3 provides the only task-level breakdown, at the 0.91B scale with 40B tokens. RMoE's post-SFT gains over SMoE (frozen router) are: LAMBADA +5.51, Hellaswag +1.22, PIQA -0.11, SciQ -0.3 (SMoE better), ARC-Easy -0.34 (SMoE better). The fact that RMoE underperforms SMoE on three of five tasks at 40B post-SFT is buried in the average improvement (+1.20) and not discussed. A practitioner running only SciQ and PIQA would conclude RMoE is slightly worse than SMoE, contradicting the paper's universal-improvement narrative.

Mitigation status. The paper does not analyze task-level heterogeneity or layer-level heterogeneity in performance impact. The per-task breakdown is present in the data (Table 3) but not discussed in the text — the paper only reports and discusses the average across tasks. The layer-specific analyses (gating entropy per layer in Appendix Figures 8–10, expert frequency in Appendix Figure 12) show substantial variation but are not connected to performance outcomes. A simple analysis — e.g., correlating per-layer gating entropy with per-layer contribution to the final loss, or ablating RMoE on subsets of layers — would provide actionable guidance but is absent.

Limitation 6: The Training Cost Overhead Is Modest but the Inference Overhead Is Unexplored

The assumption or constraint. The paper reports training-time overhead: at the 0.91B scale, RMoE training is 49.07 s/step vs. 48.87 s/step for SMoE, a 0.4% increase (Table 3). At small scale, the overhead is 972.9 vs. 960.2 s/1k steps, a 1.3% increase (Table 1). The paper treats this as negligible and does not further analyze inference costs. However, training and inference have fundamentally different cost structures: during training, the GRU computation (projectors + shared GRU) is a tiny fraction of total FLOPs relative to the expert FFNs and attention mechanisms. During inference, particularly in deployment scenarios where expert parallelism and memory bandwidth constraints dominate, the additional GRU forward pass — which must be computed sequentially across layers — could become a non-trivial latency contributor, especially for small-batch or single-query inference where the expert computation is already fast.

The consequence. A practitioner deploying an RMoE model in a latency-critical application (e.g., real-time chatbots, on-device inference) cannot estimate the inference-time overhead from the paper's training-time measurements. The GRU introduces a new sequential dependency in the forward pass: $h_i$ cannot be computed until $h_{i-1}$ is available, which means the GRU chain forces layer-by-layer sequential execution even if the rest of the transformer could be partially parallelized across layers (e.g., through pipeline parallelism). For large-batch inference where expert computation is the bottleneck, this sequential GRU cost may remain negligible. For single-query, latency-bound inference, it could become measurable. Additionally, the GRU state $h_i$ may need to be stored per-token if caching strategies are used for autoregressive generation, increasing the KV-cache-equivalent memory footprint by $p = 128$ floats per token per layer.

What evidence exists in the paper. The paper provides no inference-specific measurements: no latency benchmarks, no throughput measurements at various batch sizes, no memory footprint analysis for autoregressive decoding, and no comparison of time-to-first-token or tokens-per-second between RMoE and SMoE at inference time. All reported costs are training-time (s/step, peak GPU memory during training).

Mitigation status. The paper does not acknowledge inference overhead as a concern. The training overhead numbers (0.4–1.3%) are reported matter-of-factly, implying the method is essentially free. But training cost is a poor proxy for inference cost, particularly for MoE models where inference deployments often face different bottlenecks (expert load imbalance, communication overhead for distributed experts, memory bandwidth for loading expert weights). A thorough cost analysis would need to measure end-to-end inference latency and throughput at representative batch sizes, with and without the GRU, to determine whether the architectural benefit justifies the added complexity in deployment. The paper's claim that RMoE "introduces negotiable costs" (Abstract) is based exclusively on training and memory measurements, not inference measurements.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a new axis for MoE improvement that the field had largely overlooked: the routing state channel as a separable architectural component, distinct from the scoring function and the selection mechanism. Prior work on MoE routers — from XMoE's cosine-similarity scoring (Chi et al., 2022) to HyperMoE's hypernetwork-generated weights (Do et al., 2023) to SMoE-MLP's increased capacity (Shen et al., 2023) — all operated within the assumption that each layer's router should take the current hidden state $x_i$ as input and produce expert scores. The implicit design space was "given $x_i$ , what function $g_i$ maximizes routing quality?" RMoE questions that premise by inserting a preprocessing stage — a learned, recurrent transformation of $x_i$ that incorporates routing history — before the scoring function, effectively arguing that the input to the router is as important a design choice as the router architecture itself.

The magnitude of this shift is incremental but generative. It is not a paradigm shift — RMoE does not replace the standard routing formulation, it augments it with an additional computation stage that any existing router design can adopt. But it is generative because it factors the MoE computation graph into three separable stages — (1) routing state preparation, (2) expert scoring, (3) expert selection — that can be innovated on independently and combined. This factoring, demonstrated explicitly through RMoE's composition with XMoE (Table 2), provides a template for cumulative progress: a researcher developing a better scoring function (Stage 2) need not also solve the cross-layer coordination problem (Stage 1), and vice versa. This is analogous to how the transformer architecture's factoring into attention, feed-forward, and normalization sublayers enabled independent innovation on each component rather than requiring monolithic redesigns.

The most significant conceptual contribution is the diagnostic finding, not the architecture itself. The paper's systematic ablation of the Recurrent Gradient (Tables 4–5) demonstrates that the GRU's benefit comes primarily from the backward-pass gradient pathway it creates for router parameters, not from the forward-pass routing history it provides. This reframes the problem: the central challenge in MoE routing is not "how do we share information across layers?" but rather "how do we provide adequate optimization signal to router parameters, especially in early layers of deep models?" This reorientation has immediate implications for research priorities. Approaches that focus on richer forward-pass representations (larger routers, more expressive scoring functions) without addressing the gradient starvation problem are unlikely to help — consistent with the paper's finding that SMoE-MLP fails to improve over standard SMoE. Approaches that provide better gradient pathways — through recurrence, through explicit residual connections designed for routers, through auxiliary training objectives that densify the router gradient signal, through staged training where routing is learned separately — become the high-priority research direction.

The paper reconciles a genuine puzzle in the MoE literature. Multiple prior works found that non-learned routers — hash functions (Roller et al., 2021), stochastic policies (Zuo et al., 2021), fixed-random initialization (Chen et al., 2023) — are competitive with learned routers. This was deeply puzzling: if routing is important (as evidenced by the large parameter efficiency gap between MoE and dense models), why doesn't the learned router extract value from its training signal? The paper's Recurrent Gradient analysis provides a coherent explanation: the router's gradient signal from the language modeling loss is weak, sparse (only the top-k experts receive non-zero gradients), and indirect (flowing through many layers of FFN and attention transformations), while the gradient from the load balancing loss — which pushes toward uniform expert assignment, not semantically meaningful routing — dominates early training and then vanishes (Table 11). The learned router is essentially optimization-starved: it receives just enough signal to balance expert loads, but not enough to discover semantically meaningful token-expert assignments. A fixed-random router avoids this problem by not trying to learn at all — it provides diverse (if random) expert assignments that prevent the worst failure mode (all tokens going to one expert) while incurring no optimization cost. The Recurrent Gradient partially remedies this by providing a direct, low-impedance pathway for language modeling gradients to reach shallow-layer routers, giving them enough signal to learn beyond load balancing. This explanation is consistent with why deeper models benefit more from RMoE (Figure 2) — the gradient vanishing problem is more severe when the router is farther from the loss.

The paper also suggests a research direction that becomes less attractive: naively scaling router capacity. SMoE-MLP consistently underperforms standard SMoE at small scale (Table 1: 1.137 vs. 1.128 BPC), and at the 15B scale, SMoE-MLP is mixed — better on Hellaswag and MMLU, worse on GSM8K (Appendix A.3, Table 9). If the router's core problem is inadequate gradient signal rather than insufficient representational capacity, adding more parameters to the router is like widening a pipe without increasing the water pressure — it adds capacity that cannot be utilized. This is a negative result with practical implications: the common intuition that "a more expressive router should be better" is likely wrong given current MoE training paradigms. Research effort should shift from router capacity to router optimization.

Follow-Up Research This Work Enables

Router-specific gradient pathways without recurrence. The Recurrent Gradient finding — that the GRU helps primarily through improved gradient flow, not forward-pass information — suggests an immediate follow-up: test whether simpler gradient highways can achieve the same or better benefit without the sequential GRU computation. A natural experiment is a residual router connection: at each layer $i$ , the router input is $[x_i; h_{i-1}]$ where $h_{i-1} = \text{MLP}(x_{i-1})$ is a learned projection of the previous layer's hidden state, and the gradient flows freely through this residual path. This would provide the gradient pathway without the GRU's gating mechanism, disentangling whether the GRU's specific update rules (reset gate, update gate) matter or whether any direct gradient connection across layers suffices. A stronger version would test whether Skip-Router connections — where the language modeling loss gradient is explicitly backpropagated to router parameters at all previous layers through direct residual connections that bypass attention and FFN transformations — outperform RMoE. If simple residuals match or exceed the GRU, the field should pursue gradient pathway design rather than recurrent state design for router improvement, which is architecturally simpler and more compatible with existing training infrastructure.

Selective RMoE: applying layerwise recurrence only where it helps. The paper's task-level breakdown (Table 3) reveals substantial heterogeneity: at 40B tokens post-SFT, RMoE outperforms SMoE by 5.51 points on LAMBADA but underperforms SMoE on SciQ (-0.3) and ARC-Easy (-0.34). This suggests RMoE's benefit is not universal across all token types or reasoning demands. A targeted follow-up would train a difficulty or token-type classifier (e.g., using the GRU's own update gate activations as a signal of whether cross-layer coordination is being used) and selectively apply the GRU only to layers or tokens where historical routing information is predictive of downstream performance. Concretely: measure, per token, whether the GRU's update gate $z_i$ tends to be close to 1 (carrying forward state) or close to 0 (overwriting), and correlate this with the token's contribution to the final loss. If tokens where the GRU carries forward state show larger performance improvements, this identifies a subset of tokens that benefit from cross-layer coordination, and the GRU computation can be skipped for tokens where it provides no benefit, reducing overhead. The experiment would measure the performance-overhead Pareto frontier — what fraction of RMoE's gain is preserved when the GRU is applied to only the top-X% of tokens or layers by predicted benefit?

Causal intervention on gating entropy to test the exploration-exploitation story. The paper's gating entropy analysis (Figure 4, Table 8) argues that RMoE achieves a better balance between exploration (high entropy, diverse expert assignments) and exploitation (low entropy, confident assignments), and implies this balance causally contributes to performance. The causal direction is ambiguous: does moderate gating entropy cause better training, or is it a symptom of better optimization through the Recurrent Gradient? To test this, train a standard SMoE model with an explicit entropy regularization term added to the training loss: $\mathcal{L} = \mathcal{L}_{\text{LM}} + \lambda_{\text{LB}} \mathcal{L}_{\text{LB}} + \lambda_{\text{entropy}} \mathcal{L}_{\text{entropy}}$ , where $\mathcal{L}_{\text{entropy}}$ penalizes the gating distribution's deviation from RMoE's observed entropy distribution (measured from a pre-trained RMoE checkpoint). If entropy regularization alone recovers some or all of RMoE's gain, the behavioral difference is causal. If it does not, then moderate entropy is a byproduct of the Recurrent Gradient and efforts to directly manipulate routing entropy are misguided. This experiment would clarify whether future work should target "better exploration" in routing (which could be achieved through simpler means than recurrence) or "better gradient flow" (which requires architectural changes).

Characterization of the depth scaling limit for recurrent routing. Figure 2 shows the RMoE advantage growing from 6 to 32 layers, but the paper tests no models deeper than 32 layers and no model with the hidden size and width of production MoE systems (which may reach 64–128 layers). RNNs are known to suffer from vanishing and exploding gradients over long sequences, and the GRU in RMoE processes a "sequence" of length equal to the number of transformer layers — potentially 64, 96, or 128 steps. A stress-test experiment would train RMoE and SMoE at depths of 48, 64, and 96 layers (with comparable total parameters by adjusting hidden size) and measure three quantities at each depth: (1) the performance gap between RMoE and SMoE (does it continue to widen, plateau, or reverse?); (2) the GRU gradient norm across layers (does it explode or vanish at very deep settings, and does gradient clipping stabilize it?); and (3) whether standard RNN stabilization techniques — layer normalization on the GRU state, orthogonal initialization of GRU weight matrices, or replacing the GRU with a simpler residual MLP that has no multiplicative gates — preserve or improve RMoE's benefit at depth. This would establish the depth regime where RMoE is safe to deploy and identify what modifications are needed to extend it to arbitrarily deep models.

Cross-architecture generalization: testing RMoE in encoder-decoder and vision MoE. The paper exclusively tests decoder-only language models. The Recurrent Gradient mechanism — providing a gradient pathway for router parameters across layers — should, in principle, benefit any deep MoE architecture regardless of modality or encoder/decoder structure. A straightforward generalization experiment would apply RMoE to MoE-based vision transformers (e.g., MoE-ViT for image classification) and to encoder-decoder MoE models (e.g., MoE-based T5 for translation or summarization). The experiment would test whether the Recurrent Gradient benefit is specific to autoregressive language modeling (where the loss is computed only at the final token) or generalizes to tasks where the loss is computed at every position (masked language modeling, image classification, sequence-to-sequence tasks). If RMoE benefits autoregressive LMs more than encoder-based models, this would suggest the gradient starvation problem is particularly acute when the loss signal must propagate through many autoregressive steps before reaching early-layer routers, providing a more precise diagnosis of when recurrent routing is most valuable. If the benefit is uniform across architectures, RMoE becomes a general-purpose MoE component rather than an LM-specific optimization.

Load balancing loss scheduling informed by RMoE's gradient dynamics. Table 11 in Appendix A.4.1 reveals a specific pathology: the standard linear router's gradient from the load balancing loss dominates early training (gradient norm 0.433 vs. 0.625 from LM loss at step 100) and then vanishes (0.001 at step 60,000), while the RNN router maintains a more balanced gradient throughout. This suggests that load balancing loss decay scheduling — starting with a high balancing coefficient and annealing it to zero over training — could replicate RMoE's effect of preventing early over-optimization for load balance while allowing the LM loss to drive later specialization. A direct experiment would sweep over decay schedules for $\lambda_{\text{LB}}$ in standard SMoE (e.g., constant, linear decay, cosine decay, exponential decay) and measure whether an appropriately tuned schedule matches or exceeds RMoE's performance. If a simple schedule change in standard SMoE recovers the RMoE benefit, the architectural complexity of the GRU becomes unnecessary, and the paper's contribution shifts from "use GRU for routing" to "the diagnostic that load balance loss scheduling is a critical but overlooked hyperparameter in MoE training." This experiment would cost substantially less than training full RMoE models at scale, since it only requires changing a loss coefficient, not modifying the model architecture.

Practical Applications and Downstream Use Cases

Cost-efficient MoE training for organizations with fixed compute budgets. The paper's central result — that RMoE consistently improves MoE performance with negligible training overhead (0.4% slower step time at 0.91B scale, Table 3) — translates directly to better model quality per GPU-hour. For a team training a 15B-parameter MoE model on 400B tokens (the scale of Appendix A.3), RMoE delivers a 1.06-point MMLU improvement (60.60 vs. 59.54) and a 0.065 improvement in domain-average perplexity (5.620 vs. 5.685) with no measurable wall-clock penalty beyond the 3.5M additional parameters (a 0.02% increase at 15B scale). In a production setting where model quality directly impacts downstream revenue (e.g., an API provider whose models are benchmarked by customers), this is a free improvement. The implementation cost is a one-time engineering investment to add the per-layer projectors and shared GRU to the training code; the inference overhead, while not measured in the paper, is likely minimal since the GRU state computation is tiny (128-dimensional operations per layer) relative to expert FFN computation (thousands of dimensions). For organizations already using MoE architectures (e.g., Mixtral-style models in open-source deployments, or proprietary MoE systems), adopting RMoE requires no changes to the expert architecture, load balancing strategy, or training data pipeline.

Improved knowledge integration for long-context and retrieval-augmented models. The task-level breakdown at 40B tokens post-SFT (Table 3) shows RMoE's largest single-task gain on LAMBADA (+5.51 points over SMoE with frozen router), a benchmark specifically designed to require broad discourse context for word prediction. This suggests RMoE's cross-layer routing coordination is particularly beneficial when the model must integrate information across long token sequences — exactly the regime that retrieval-augmented generation (RAG) and long-context LLMs target. A system builder deploying an MoE-based LLM for document question-answering or multi-hop reasoning could specifically benefit from RMoE, since these tasks demand coordinated processing across multiple pieces of evidence that are spread across the input context. The GRU's routing memory — which tracks which experts processed earlier tokens in the sequence — may help the model maintain coherent "processing pathways" for related pieces of information even when they are far apart in the input. The paper provides no direct evidence for this (it does not test on long-context or multi-hop tasks), but the LAMBADA result is suggestive and consistent with the mechanism: if cross-layer coordination helps with long-range dependency resolution in language modeling, it should also help with cross-document information integration in RAG settings.

Expert specialization auditing through GRU state analysis. For teams deploying MoE models where interpretability matters — e.g., models used for educational assessment, medical decision support, or legal document analysis — the GRU state provides a new diagnostic tool. Because $h_i$ explicitly encodes routing history, it can be analyzed to understand why the model routes certain tokens to certain experts. A deployer could record the GRU state trajectory for specific input types and visualize which previous routing decisions most influence current expert selection, potentially identifying biases (e.g., certain demographic terms consistently routing to experts that produce lower-quality outputs) or failure modes (e.g., the GRU state failing to update on critical disambiguation cues, leading to persistent expert misassignment). This is a speculative application — the paper does not develop GRU-state interpretability tools — but it follows naturally from the architecture: the GRU state is a low-dimensional (128-dim), purpose-built representation of routing decisions that is cleaner to analyze than the full hidden state, which mixes routing-relevant and semantic features. A practical deployment could include a monitoring dashboard that tracks GRU state statistics (e.g., update gate activation distributions per input domain) and flags anomalies that correlate with known quality issues.

When to Prefer This Method

The paper does not explicitly position RMoE against named alternatives with a clear tradeoff matrix. It presents RMoE as a universal architectural improvement — "a novel computation stage orthogonal to existing methods" (Abstract) — to be applied to all MoE layers, not as a method to be chosen in some circumstances and avoided in others. The paper identifies no regime where RMoE underperforms standard SMoE in aggregate (though individual task results in Table 3 show SMoE outperforming RMoE on SciQ and ARC-Easy at 40B tokens). The decision rule implied by the paper is straightforward: if you are training an MoE model, you should use RMoE, because it adds negligible cost and consistently improves performance. The paper provides no conditional guidance (e.g., "prefer RMoE for deep models but not shallow ones," though Figure 2 suggests the benefit grows with depth) and no identified failure mode that would contraindicate its use. Given the absence of explicit tradeoffs in the paper, a forced decision matrix would impose distinctions the authors did not make. The one caution that emerges from the analysis (though not from the paper's explicit guidance) is that the inference overhead is unmeasured — until inference latency and throughput are benchmarked, practitioners deploying in strict latency-bound settings should profile RMoE's GRU cost at their target batch size before adopting.