Neural Tech Daily
ai-research

Mixture-of-Experts Routing Stability: A Multi-Paper Technical Reference

Technical reference on MoE routing stability. Walks ST-MoE (Zoph 2022), Expert Choice (Zhou 2022), StableMoE (Dai 2022) — what each fixes, where they disagree.

Updated ~31 min read
Share
Figure 1 of arXiv:2202.08906 (ST-MoE): training instabilities for sparse models, comparing an unstable training run on the left where the training loss diverges with a stable run on the right under the same configuration.

Figure 1 of ST-MoE: Designing Stable and Transferable Sparse Expert Models (arXiv:2202.08906), reproduced for editorial coverage.

1. Paper identity and scope

Primary citations.

  • Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. “ST-MoE: Designing Stable and Transferable Sparse Expert Models.” arXiv:2202.08906, February 2022 1 .
  • Zhou, Y., Lei, T., Liu, H., Du, N. et al. “Mixture-of-Experts with Expert Choice Routing.” arXiv:2202.09368, NeurIPS 2022 3 .
  • Dai, D., Dong, L., Ma, S., Zheng, B., Sui, Z., Chang, B., Wei, F. “StableMoE: Stable Routing Strategy for Mixture of Experts.” arXiv:2204.08396, ACL 2022 5 .

Retrieval. This review draws on the arXiv abstract pages and ar5iv HTML renders of all three papers 1 2 3 4 5 6 .

Classification. Architecture, LLM-based, sparse model training. All three papers attack a common failure mode in mixture-of-experts (MoE) language models: the routing layer that assigns tokens to experts is unstable during training in ways that hurt convergence, fine-tuning, and inference consistency.

Technical abstract (in the publication’s voice). MoE layers replace a dense feedforward block with EE parallel experts and a small router that picks the top-kk experts for each token. The promise is conditional computation — total parameters scale with EE but per-token FLOPs scale with kk. The reality, as of late 2021 / early 2022, was that MoE models at scale suffered from three distinct failure modes: training divergences (loss spikes that occasionally crashed runs entirely), fine-tuning brittleness (sparse models overfit on small downstream tasks), and routing fluctuation (the same token gets assigned to different experts at different points in training, undermining the conditional-computation assumption). The three papers reviewed here address these in complementary ways. ST-MoE (Zoph 2022) is the engineering bible: a comprehensive sweep of design choices (router precision, auxiliary losses, expert dropout, batch composition) that together produced the first reliably-stable trillion-parameter sparse model 1 . Expert Choice (Zhou 2022) inverts the routing direction — instead of tokens picking experts, experts pick tokens, which automatically guarantees load balance 3 . StableMoE (Dai 2022) introduces two-stage training: learn a balanced router, freeze it, then learn the experts against the frozen router, eliminating routing fluctuation by construction 5 .

Primary research question. How do you train a mixture-of-experts language model that (a) does not diverge, (b) routes consistently across training, (c) generalises after fine-tuning, and (d) keeps the per-expert load balanced enough that compute is not wasted on idle experts?

Core technical claims.

  • ST-MoE. [From the paper] A specific design recipe (router in float32, auxiliary load-balancing loss with a particular weighting, conservative expert dropout for fine-tuning, route-z auxiliary loss) produces a 269B-parameter sparse model with the compute cost of a 32B dense model, achieving SOTA on diverse transfer benchmarks — the first sparse model to do so 1 .
  • Expert Choice. [From the paper] Inverting routing so that experts pick top-kk tokens (rather than tokens picking top-kk experts) guarantees uniform expert load, allows tokens to be processed by a variable number of experts, and improves training convergence by more than 2x over Switch / GShard baselines on the same dense-budget targets 3 .
  • StableMoE. [From the paper] Routing fluctuation — the phenomenon where the same input is routed to different experts across training steps — is empirically substantial (40.9% of tokens still change expert assignments at 20% through training; 15.4% still fluctuate near the end). StableMoE’s two-stage training distills the routing into a frozen lightweight router, eliminating the fluctuation and improving convergence and downstream performance 5 .

Core technical domains.

DomainDepth required
Transformer feedforward blocksModerate
Top-k routing and softmax gatingDeep
Load-balancing auxiliary lossesDeep
Distributed training (all-to-all communication)Moderate
Mixed-precision trainingModerate
Fine-tuning generalisationModerate

Reader prerequisites. Knowing that a transformer block has an attention layer followed by a feedforward (MLP) layer, that the MLP is a position-wise dense computation, and that “expert” in MoE refers to one of EE parallel copies of this MLP.

How this review marks its registers.

  • Author-stated / [From the paper] — direct claims from any of the three papers, bound to specific sections.
  • Facts — common-knowledge background (transformer block structure, distributed-training fundamentals).
  • AI analysis / [Analysis] — the pipeline’s analytical layer (worked examples, dimensional analysis, plain-English on-ramps, three-depth summary).
  • Reviewer perspective / [Reviewer Perspective] — independent commentary beyond what the papers prove.

2. TL;DR and executive overview

TL;DR. Mixture-of-Experts language models trade total parameters for conditional computation: a 269B-parameter MoE may run at the FLOPs of a 32B dense model. The architecture’s hard problem is the routing layer. ST-MoE (Zoph 2022) compiles the production-grade engineering recipe; Expert Choice (Zhou 2022) flips routing direction to guarantee load balance; StableMoE (Dai 2022) freezes the router after a warmup to eliminate routing fluctuation. Together they form the foundation that 2024-2026 production MoE models (Mixtral, DeepSeek-V3, GPT-4-class architectures) build on.

Executive summary. Pre-2022 sparse models — GShard 9 and Switch Transformer 8 — established the basic top-1 / top-2 token-routing pattern but were known to be finicky to train at scale. The 2022 trio of papers reviewed here addressed three orthogonal aspects of the instability problem.

ST-MoE is a methodological paper: its contribution is not a single new mechanism but a comprehensive empirical sweep of every design choice that affects MoE stability. The paper’s recipe — float32 router computation, auxiliary load-balancing loss tuned to ~0.01, route-z auxiliary loss for numerical stability, conservative dropout in fine-tuning — is the de facto baseline that subsequent MoE papers compare against. ST-MoE’s 269B-parameter sparse model was the first to top diverse transfer learning benchmarks (SuperGLUE, ARC, XSum, CNN-DM, WebQA, Natural Questions) 1 .

Expert Choice flips the routing problem on its head. Standard MoE has each token compute logits over experts and select the top-kk experts; this is “token-choice.” Expert Choice has each expert compute logits over tokens and select the top-kk tokens; the dual problem. The structural advantage is that expert load balance is now a constraint of the algorithm rather than a goal to be approximated by an auxiliary loss — every expert gets exactly kB/Ek \cdot \mid B\mid / E tokens per batch by construction. The price is that some tokens may not be picked by any expert (these are “dropped” or pass through unprocessed) and some may be picked by many experts (the variable-kk feature).

StableMoE addresses a subtler problem: even with stable load balance, the identity of the expert routed to may change throughout training. The Dai paper measures this fluctuation: 40.9% of tokens still change expert assignment at 20% through training, 15.4% still fluctuate near convergence. This is wasted optimisation work — the experts are being trained against a non-stationary target. The StableMoE two-stage procedure (train router with experts, freeze router, train experts only) eliminates the fluctuation and improves both convergence speed and final downstream performance on language modelling and multilingual translation.

Five practitioner-relevant takeaways.

  1. ST-MoE’s float32 router precision recommendation is non-negotiable. [From the paper] All experiments at scale require the router to compute in float32 even when the rest of the model is bf16/fp16. Lower precision in the router is the dominant cause of MoE training divergence at >10B> 10B parameter scale 1 .
  2. The load-balancing auxiliary loss has a sweet spot around 0.01. [From the paper] ST-MoE sweeps and finds that values much higher than 0.01 hurt convergence (too much pressure on uniform routing) and much lower than 0.01 allow expert collapse (one expert dominates) 1 .
  3. Expert Choice is structurally elegant but has deployment friction. Token-dropping during inference (some tokens not picked by any expert) requires application-side handling; per-token expert count being variable complicates KV-cache reasoning in autoregressive serving 3 .
  4. StableMoE’s fluctuation measurement is the value of the paper. Even teams not adopting the two-stage training procedure should be aware of the fluctuation phenomenon; the 15.4%-fluctuation-at-convergence number is a useful diagnostic for “is my MoE training healthily?” 5 .
  5. [Analysis] Production MoE in 2024-2026 (Mixtral, DeepSeek-V3) inherits all three lines. Mixtral uses token-choice top-2 with ST-MoE-style auxiliary losses 11 ; DeepSeek-V3 uses fine-grained expert decomposition with auxiliary-loss-free load balancing 10 . The 2022 trio’s design choices remain the framing of every subsequent MoE paper’s “we differ from prior work in…” section.

Pipeline overview. All three papers operate during training of MoE language models — at inference time, the trained routing is just a fixed function. The cross-paper unit of comparison is training-time stability (loss spikes, divergences, fluctuations) and downstream evaluation quality after pretraining + fine-tuning.

2.5. Glossary

TermPlain-English explanationFirst appears in
Mixture-of-Experts (MoE)A transformer architecture variant where the feedforward layer is replaced by EE parallel copies (experts) and a router that picks kk experts per token.Section 1
ExpertOne of the EE parallel feedforward networks. Each expert has the same architecture as the original dense feedforward but its own weights.Section 1
Router (or gate)A small linear layer that produces logits over experts for each token, then takes top-kk (token-choice) or top-kk tokens (expert-choice).Section 1
Top-kk routingEach token is assigned to its kk highest-scoring experts; common values are k=1k=1 (Switch) and k=2k=2 (GShard, Mixtral).Section 1
Load balanceThe property that experts receive roughly equal numbers of tokens. Imbalance wastes compute (idle experts) or causes overflow (dropped tokens).Section 1
Auxiliary load-balancing lossAn extra term added to the training loss to penalise uneven expert utilisation. Typically the dot product of mean per-expert routing probability and mean per-expert token fraction.Section 1
Routing fluctuationThe same input token being routed to different experts at different training steps. The phenomenon StableMoE measures and addresses.Section 1
Token-choice routingThe standard pattern: each token selects its top-kk experts.Section 1
Expert-choice routingThe Zhou 2022 inversion: each expert selects its top-kk tokens.Section 1
Switch TransformerFedus 2022’s MoE variant with k=1k=1 routing — simplest and fastest, but with the highest load-balance challenge.Section 4
GShardLepikhin 2020’s MoE variant with k=2k=2 routing and the original auxiliary load-balancing loss formulation.Section 4
Expert capacityThe maximum number of tokens an expert is allowed to process per batch. Tokens beyond capacity are “dropped” or “overflowed”.Section 5
Route-z lossST-MoE’s auxiliary loss penalising large router logits to prevent fp16/bf16 numerical issues.Section 5
[From the paper] prefixDefault register; bound to a specific paper.Throughout
[Analysis] labelThe publication’s reasoned synthesis across papers.Throughout

3. Problem formalisation

Notation table.

SymbolTypeMeaningFirst appears in
h\mathbf{h}VectorToken hidden state (input to MoE layer)Section 3
EEScalarNumber of expertsSection 3
kkScalarTop-kk routing rank (1 or 2 typically)Section 3
r(h)r(\mathbf{h})FunctionRouter output (logits over experts)Section 3
gi(h)g_i(\mathbf{h})FunctionGating weight for expert ii, normalisedSection 3
ei(h)e_i(\mathbf{h})FunctionExpert ii‘s outputSection 3
CCScalarExpert capacity (max tokens per expert per batch)Section 5
Laux\mathcal{L}_{\text{aux}}ScalarAuxiliary load-balancing lossSection 5

Formal problem statement (token-choice). Given a transformer with an MoE feedforward layer of EE experts, the layer’s output for token h\mathbf{h} is

MoE(h)=iTopK(r(h),k)gi(h)ei(h)\text{MoE}(\mathbf{h}) = \sum_{i \in \text{TopK}(r(\mathbf{h}), k)} g_i(\mathbf{h}) \, e_i(\mathbf{h})

where rr is the router, TopK selects the top-kk experts by logit, and gig_i is the softmax-normalised gating weight. Training optimises a language-modelling loss plus an auxiliary load-balancing loss.

Explicit assumption list.

  1. Sparse computation. Only kk of EE experts are evaluated per token; the rest are skipped. This is what gives MoE its FLOPs advantage.
  2. Expert capacity. Each expert can process at most CC tokens per batch; tokens routed beyond capacity are dropped (replaced by zeros or residual passthrough).
  3. Distributed execution. Experts are typically distributed across devices; routing requires all-to-all communication. ST-MoE / Expert-Choice / StableMoE all assume this distributed setting.
  4. Stable training as a desideratum. Loss spikes that crash runs are unacceptable; “stability” means L2L^2 norms of gradients, router logits, and per-step loss values remain bounded across training.

Why the problem is hard. Three sources of instability:

  1. Router gradient sparsity. Most experts get zero gradient signal from most tokens; only the chosen experts contribute. The router itself receives gradient only through the chosen experts. This makes the router susceptible to mode collapse (always pick the same experts).
  2. Numerical sensitivity. Router logits can grow large during training, causing softmax overflow in fp16/bf16. ST-MoE introduces the route-z loss specifically to bound logit magnitudes.
  3. Load balance is an emergent property. The auxiliary loss is a soft constraint; under-tuned, expert collapse happens; over-tuned, the router can’t differentiate experts.

LLM-based positioning. All three papers operate on encoder-decoder (T5-family) or decoder-only transformers. The MoE layer replaces standard FFN blocks; attention remains dense.

4. Motivation and gap

Real-world problem. MoE is the dominant architecture for scaling language model capacity beyond what dense models can support at fixed training FLOPs. The 2022-2026 frontier — Mixtral 8x7B 11 , DeepSeek-V3 671B 10 , GPT-4-class production architectures (rumoured MoE) — all depend on the routing-stability work from 2022 being reliable enough for production training.

Existing approaches.

  • Shazeer 2017 7 introduced the sparsely-gated MoE concept with noise-augmented top-kk routing.
  • GShard (Lepikhin 2020) 9 scaled MoE to multilingual translation at hundreds of billions of parameters, introducing the standard auxiliary load-balancing loss.
  • Switch Transformer (Fedus 2022) 8 simplified to k=1k=1 routing for efficiency, scaled to trillion-parameter pretraining, but documented stability struggles.

Gap. Switch and GShard worked but were unstable at scale. The 2022 trio’s gap: a rigorous understanding of which design choices cause instability, and design alternatives that eliminate specific failure modes.

[External comparison] Position vs Switch / GShard. ST-MoE accepts the Switch / GShard routing framework and engineers it to stability. Expert Choice rejects the framework and proposes its dual. StableMoE accepts the framework but adds a structural fix (frozen router after warmup). The three are not mutually exclusive; production systems often combine elements (e.g., Mixtral uses ST-MoE-style auxiliary loss with k=2k=2 token-choice routing).

5. Method overview — ST-MoE

[From the paper, Sections 4-7] ST-MoE is not a single new mechanism but a recipe. The paper’s empirical contributions:

Float32 router. Compute router logits and softmax in float32 even when the rest of the model is bfloat16. Router logits can grow large during training; bf16 / fp16 overflows cause divergences. Float32 router is essentially free in compute (the router is tiny relative to the experts) and eliminates this failure mode.

Auxiliary load-balancing loss formulation. ST-MoE uses the GShard load-balancing loss:

Laux=αEi=1EfiPi\mathcal{L}_{\text{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot P_i

where fif_i is the fraction of tokens routed to expert ii, PiP_i is the mean router probability for expert ii, and α\alpha is the loss weight. ST-MoE recommends α=0.01\alpha = 0.01 after extensive sweeps.

Route-z auxiliary loss. A second auxiliary term penalising large router logits:

Lz=β1Bj=1B(logi=1Eexp(r(hj)i))2\mathcal{L}_z = \beta \cdot \frac{1}{B} \sum_{j=1}^{B} \left(\log \sum_{i=1}^{E} \exp(r(\mathbf{h}_j)_i)\right)^2

where BB is batch size. Recommended β=103\beta = 10^{-3}. This term is what makes float32 router unnecessary in cases where the model is otherwise stable, but ST-MoE recommends using both for robustness.

Fine-tuning dropout. Sparse models overfit on small downstream tasks. ST-MoE introduces selective dropout on expert outputs during fine-tuning; the paper’s Figure 3 (in-article image below) shows the overfitting pattern that motivates the choice 1 .

Batch composition. ST-MoE recommends pretraining batches with multiple parallel sequences to give the router stable per-batch statistics; small batch sizes cause router noise.

Figure 3 of arXiv:2202.08906 (ST-MoE): sparse models are prone to overfit; train and validation curves for an ST-MoE-L and a dense-L model fine-tuned on the CB task (250 train sequences) and ReCoRD (138k train sequences), showing the sparse model learns faster on train but underperforms on smaller-data validation.

Figure 3 of ST-MoE: Designing Stable and Transferable Sparse Expert Models (arXiv:2202.08906), reproduced for editorial coverage.

5b. Method overview — Expert Choice

[From the paper, Section 3] Expert Choice routing inverts the assignment:

  • Token-choice (Switch / GShard / ST-MoE). Each token has a logit over experts; selects top-kk.
  • Expert-choice. Each expert has a logit over tokens; selects top-kk' tokens, where k=Ck' = C is the expert capacity.

The dual property: the per-expert load is exactly kk' tokens per batch by construction. There is no need for an auxiliary load-balancing loss because the algorithm structurally guarantees balance.

Per-token expert count is variable. Some tokens are picked by many experts (those whose logits are high for many experts); some by few; some by zero. The zero-expert tokens pass through unprocessed (or via the residual stream only).

[From the paper] Empirical results: training convergence improves by more than 2x over Switch / GShard on the same dense-budget baseline. GLUE and SuperGLUE fine-tuning improves substantially 3 .

[Analysis] The asymmetry has deployment costs: token-dropping (zero-expert tokens) requires application-side acceptance that some token positions are processed less; the variable expert count complicates serving infrastructure that assumes uniform per-token compute.

5c. Method overview — StableMoE

[From the paper, Section 3] StableMoE’s two-stage training:

Stage 1 (warmup). Standard MoE training with experts and router co-learned for a fraction of total training (~20-30%).

Stage 2 (frozen router). Take the trained router, distill it into a lightweight standalone model (a small MLP that takes token hidden states and produces expert logits). Freeze this distilled router. Train only the experts for the remainder.

Why it works. [From the paper] Routing fluctuation is empirically measured by tracking, for each token across training, how many times it changes expert assignment. The paper reports 40.9% of tokens fluctuate within 20% of training time, and 15.4% still fluctuate at the end. With a frozen router, this fluctuation is exactly zero in stage 2; experts are trained against a stationary target.

Trade-off. [Analysis] The frozen router cannot adapt to the experts’ evolving representations after stage 1. ST-MoE’s recipe and Expert Choice’s structural balance both maintain a co-evolving router; StableMoE accepts router rigidity in exchange for fluctuation elimination. Whether this trade-off is favourable depends on the regime; the paper claims it is favourable on the benchmarks tested (multilingual translation, language modelling).

6. Mathematical contributions

MATH ENTRY 1: GShard load-balancing auxiliary loss.

  • Source: GShard (Lepikhin 2020) 9 , used by ST-MoE.
  • What it is: the standard load-balancing auxiliary loss.
  • Formal definition:

Laux=αEi=1EfiPi\mathcal{L}_{\text{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot P_i

where:

  • fi=1Bb=1B1[token b routed to expert i]f_i = \frac{1}{B} \sum_{b=1}^{B} \mathbb{1}[\text{token } b \text{ routed to expert } i] — empirical fraction of tokens to expert ii.

  • Pi=1Bb=1Bsoftmax(r(hb))iP_i = \frac{1}{B} \sum_{b=1}^{B} \text{softmax}(r(\mathbf{h}_b))_i — mean routing probability to expert ii.

  • α\alpha: scalar weight (ST-MoE: 0.01).

  • Each term explained.

    • The product fiPif_i \cdot P_i: penalises experts that are both probabilistically preferred by the router and actually picked often. Minimised when each fi=Pi=1/Ef_i = P_i = 1/E (uniform).
    • The factor EE: scales the loss so it’s O(1)O(1) in expectation under uniform routing (sum of EE terms each 1/E2\approx 1/E^2, times EE).
  • Worked numerical example. With E=4E = 4 experts and B=8B = 8 tokens, suppose routing gives:

Expertfif_iPiP_ifiPif_i \cdot P_i
15/8 = 0.6250.550.344
22/8 = 0.2500.250.063
31/8 = 0.1250.150.019
40/8 = 0.0000.050.000
Sum0.426

Laux=0.0140.426=0.017\mathcal{L}_{\text{aux}} = 0.01 \cdot 4 \cdot 0.426 = 0.017.

Under uniform routing (fi=Pi=0.25f_i = P_i = 0.25), each term is 0.06250.0625, sum is 0.250.25, loss is 0.0140.25=0.010.01 \cdot 4 \cdot 0.25 = 0.01. So the imbalanced case incurs 0.0170.010=0.0070.017 - 0.010 = 0.007 of extra loss — a gentle but persistent push toward uniformity.

  • Dimensional analysis. All fi,Pi[0,1]f_i, P_i \in [0, 1], loss is dimensionless.
  • Edge cases. As fi1,Pi1f_i \to 1, P_i \to 1 for one expert (collapse): loss is αE1=0.04\alpha E \cdot 1 = 0.04, four times higher than uniform. The loss does penalise collapse but not aggressively.

MATH ENTRY 2: ST-MoE’s route-z loss.

  • Source: ST-MoE Section 4.
  • What it is: penalty on router logit magnitude.
  • Formal definition:

Lz=β1Bb=1B(logi=1Eexp(r(hb)i))2\mathcal{L}_z = \beta \cdot \frac{1}{B} \sum_{b=1}^{B} \left(\log \sum_{i=1}^{E} \exp(r(\mathbf{h}_b)_i)\right)^2

The inner expression is the log-sum-exp of router logits; squaring it bounds large magnitudes from either positive or negative side. Recommended β=103\beta = 10^{-3}.

  • Worked numerical example. With E=4E = 4 and router logits (2,1,0,1)(2, 1, 0, -1), logsumexp = log(e2+e1+1+e1)log(7.39+2.72+1.0+0.37)log(11.48)2.44\log(e^2 + e^1 + 1 + e^{-1}) \approx \log(7.39 + 2.72 + 1.0 + 0.37) \approx \log(11.48) \approx 2.44. Squared: 5.95. With β=103\beta = 10^{-3}, per-token contribution 0.006\approx 0.006. If logits drift to (20,1,0,1)(20, 1, 0, -1), logsumexp 20\approx 20, squared =400= 400, per-token contribution 0.4\approx 0.4 — a strong corrective signal.
  • Dimensional analysis. Logits are dimensionless; logsumexp is dimensionless; squared logsumexp is dimensionless.

MATH ENTRY 3: Expert Choice routing (per-expert selection).

  • Source: Zhou 2022 Section 3.

  • What it is: the inverted routing rule.

  • Formal definition. Given a batch of BB tokens and EE experts, compute the full B×EB \times E routing logits matrix RR. For each expert ii, select the top k=kB/Ek' = \lceil k \cdot B / E \rceil tokens by R[:,i]R[:, i]. The selected tokens are processed by expert ii.

  • Each term explained.

    • k=kB/Ek' = k \cdot B / E comes from preserving the total per-token expert-visits across both routing schemes; if standard token-choice routes each of BB tokens to kk experts, total visits = kBkB = EkE \cdot k', so k=kB/Ek' = kB/E.
    • Variable per-token visits: a token can appear in 0 or more experts’ top-kk' lists depending on its logit profile.
  • Worked numerical example. With B=12B = 12 tokens, E=4E = 4 experts, k=1k = 1 (token-choice equivalent), k=3k' = 3. Each expert picks 3 tokens. Total visits = 12. In token-choice with k=1k=1, each token visits 1 expert; total visits = 12. Identical total compute, different distribution: in expert-choice some tokens may visit 0 or 2+ experts.

7. Algorithm trace

[From the paper, ST-MoE algorithm + Expert Choice algorithm side by side]

Trace: token-choice routing (ST-MoE style) on a 4-token, 4-expert MoE layer.

Inputs. Batch of 4 token hidden states h1,,h4\mathbf{h}_1, \dots, \mathbf{h}_4, router weight matrix WrRd×4W_r \in \mathbb{R}^{d \times 4}, expert capacity C=2C = 2, top-k=1k = 1.

Step 1. Router logits: rb=hbWrR4r_b = \mathbf{h}_b W_r \in \mathbb{R}^4 for each token. Suppose:

TokenExpert 1Expert 2Expert 3Expert 4
13.01.00.5-0.5
20.52.01.50.0
31.00.52.50.5
42.00.50.51.5

Step 2. Top-1 selection: Token 1 → Expert 1, Token 2 → Expert 2, Token 3 → Expert 3, Token 4 → Expert 1.

Step 3. Capacity check. Expert 1 received 2 tokens (under C=2C = 2), Expert 2 received 1, Expert 3 received 1, Expert 4 received 0. All within capacity.

Step 4. Gating weights from softmax: for Token 1, g1=softmax((3.0,1.0,0.5,0.5))1=0.836g_1 = \text{softmax}((3.0, 1.0, 0.5, -0.5))_1 = 0.836. Apply expert 1 to h1\mathbf{h}_1, scale output by 0.836.

Step 5. Aggregate per-token expert outputs (in k=1k=1 case, just the single expert).

Step 6. Compute auxiliary load-balancing loss: f=(0.5,0.25,0.25,0)f = (0.5, 0.25, 0.25, 0), P=mean softmaxesP = \text{mean softmaxes}. Loss is added to the language-modelling loss.

Step 7. Backward pass. Router gradient through the chosen experts; expert gradients through their selected tokens.

Trace: expert-choice routing on the same batch.

Step 1. Same router logits as above.

Step 2. For each expert column, pick top-k=1k' = 1 tokens.

ExpertTop tokenLogit
1Token 13.0
2Token 22.0
3Token 32.5
4Token 41.5

Step 3. Each expert processes its single picked token. Token 4’s logit for Expert 1 (2.0) was second-highest but Expert 1 already took Token 1, so Token 4 only visits Expert 4.

Step 4. Per-token expert count: each token visits exactly 1 expert here (lucky distribution). In a less-symmetric case, tokens with low logits across all experts could be picked by zero experts.

Step 5. No auxiliary load-balancing loss needed: load is uniform by construction.

[Analysis] Contrast. In token-choice, Expert 4 idled (zero tokens). In expert-choice, every expert worked. This is the structural advantage. The cost: in token-choice, every token was processed by some expert; in expert-choice, with a less-favourable logit distribution, some tokens might pass through unprocessed.

8. Results and benchmarks

[From the paper] ST-MoE’s headline result: a 269B-parameter ST-MoE-32B model achieves SOTA on diverse transfer benchmarks (SuperGLUE, ARC Easy, ARC Challenge, XSum, CNN-DM, WebQA, Natural Questions) — the first sparse model to do so across this benchmark mix 1 . Inference cost is comparable to a 32B dense T5; total parameter count is 269B.

[From the paper] Expert Choice’s headline: 2x+ training-convergence speed-up over Switch and GShard at matched dense-budget baselines on a held-out language-modelling corpus; substantial fine-tuning gains on GLUE and SuperGLUE 3 .

[From the paper] StableMoE’s headline: improvements over standard MoE baselines on:

  • Language modelling perplexity (WikiText-103, ARC, OpenWebText).
  • Multilingual machine translation (WMT-21 multilingual benchmark).

The improvements are modest in absolute terms (1-3% relative on most benchmarks) but the routing-fluctuation reduction is dramatic (15.4% → 0 by construction) 5 .

[Analysis] The three papers’ benchmark stories are somewhat non-overlapping; direct head-to-head Pareto comparison is hard from the published evidence alone. Production MoE deployment in 2024-26 (Mixtral, DeepSeek-V3) effectively cherry-picks: ST-MoE-style auxiliary losses + token-choice routing dominates production, with expert-choice and stable-router-distillation appearing in research settings but not consistently in deployed frontier systems.

9. Ablations and limitations

[From the paper] Stated limitations (ST-MoE).

  • Recipe is empirically optimised on the T5 family at specific scales. Transfer to decoder-only architectures (which dominate 2023+ LLMs) was not characterised in the paper.
  • Auxiliary loss weights (α=0.01\alpha = 0.01, β=103\beta = 10^{-3}) are stated as recommendations but the sensitivity to these values is not fully characterised; production teams often re-sweep.

[From the paper] Stated limitations (Expert Choice).

  • Token-dropping (zero-expert tokens) requires application acceptance; not all downstream tasks are robust to this.
  • Variable per-token expert count complicates KV-caching and per-sequence cost modelling in inference.
  • The cost-equivalence comparison to token-choice assumes the same total visits (kBkB); for autoregressive serving where token sequences are processed incrementally, the comparison is more subtle.

[From the paper] Stated limitations (StableMoE).

  • Two-stage training adds engineering overhead.
  • Frozen-router rigidity may underperform co-evolving router on very long training schedules where expert representations drift substantially.

[Reviewer Perspective] Independent limitations.

  • Scale generalisation. All three papers’ empirical scale tops out around 270B parameters (ST-MoE). Modern frontier MoE deployments (DeepSeek-V3 671B 10 ) reach 2-3x that, and the relative ordering of techniques at frontier scale is not directly published.
  • Decoder-only generalisation. ST-MoE is encoder-decoder; production frontier MoEs (Mixtral, DeepSeek-V3) are decoder-only. The transfer of stability recipes is mostly empirical and varies across implementations.
  • Auxiliary-loss-free routing. [Analysis] DeepSeek-V3’s “auxiliary-loss-free load balancing” 10 is a 2024 development that explicitly departs from ST-MoE’s auxiliary-loss recipe and from Expert Choice’s structural-balance approach. It uses bias-shift dynamics to achieve balance without an explicit loss term; the 2022 trio reviewed here predates this development and does not address it.

10. Reproducibility

ArtefactAvailable?Source
ST-MoE training codeNOT RELEASED (Google internal)
ST-MoE model weightsNOT RELEASED
ST-MoE recipe (paper text)YES (detailed in paper)Lou et al. 1
Expert Choice training codeLIMITED (described in paper; community reproductions in 2023-24)
Expert Choice model weightsNOT RELEASED
StableMoE training codeYES (paper repo)per paper’s GitHub
Community reproductionsYES (multiple open MoE codebases inherit recipes)DeepSpeed-MoE, Megatron-LM MoE

[Analysis] None of the three papers ship as a single reproducible code+weights bundle in the modern open-source sense. The “recipes” are reproducible from paper text in well-resourced training environments; community MoE codebases (DeepSpeed-MoE, MegaBlocks, Tutel) implement variations of all three.

The 2022 trio vs Switch / GShard. [From the papers and Facts] Switch Transformer 8 (k=1k=1) and GShard 9 (k=2k=2) are the immediate ancestors. ST-MoE is the engineering deepening of Switch; Expert Choice and StableMoE diverge methodologically. The technical relationship is “we measured a failure mode of Switch / GShard and proposed a fix.”

The 2022 trio vs Mixtral 8x7B (2024) 11 . Mixtral is the canonical production decoder-only MoE: 8 experts, k=2k=2 token-choice, ST-MoE-style auxiliary load-balancing loss with specific weight tuning. Mistral’s release brought MoE to the open-weights frontier and validated that the 2022 recipes scale to production-deployable models.

The 2022 trio vs DeepSeek-V3 (2024) 10 . DeepSeek-V3 builds on the 2022 foundations but introduces multiple departures: fine-grained expert decomposition (256 experts at k=8k = 8), auxiliary-loss-free load balancing via bias-shift dynamics, and shared experts that all tokens visit. The technical relationship: “we accept ST-MoE’s stability recipes but reject its auxiliary-loss formulation, and we reject Expert Choice’s structural balance for a different structural-balance technique.”

12. Reviewer perspective

Reviewer perspective on ST-MoE. [Reviewer Perspective] The paper is a methodological-engineering tour de force: comprehensive sweeps, clear ablations, honest reporting of what worked and what did not. The “first sparse model to top SuperGLUE/ARC/XSum” claim was substantively true at publication and aged well. The main critical question is whether the recipe transfers to decoder-only frontier-scale models; the 2024 evidence (Mixtral) is partly affirmative.

Reviewer perspective on Expert Choice. [Reviewer Perspective] The dual-routing framing is mathematically elegant and the structural-balance property is appealing. The deployment friction (token-dropping, variable per-token cost) is real and has limited production adoption. The 2x+ convergence claim is benchmark-specific (held-out language modelling, GLUE/SuperGLUE fine-tuning); transfer to broader autoregressive language modelling at frontier scale is open. NeurIPS 2022 publication signals reviewer enthusiasm; production adoption signal is weaker.

Reviewer perspective on StableMoE. [Reviewer Perspective] The routing-fluctuation measurement is the paper’s lasting contribution; the two-stage training is a reasonable engineering response but the trade-off (frozen router rigidity) limits adoption. ACL 2022 publication. The frozen-router approach influenced later distilled-routing variants in production MoE work but is not directly adopted as-is in any flagship 2024-26 frontier MoE.

[Reviewer Perspective] Open methodological questions.

  1. How do these stability recipes interact at decoder-only frontier scale (DeepSeek-V3 671B, hypothetical GPT-4-class 1T+ MoE)? Direct head-to-head comparison is not publicly available.
  2. Is the auxiliary-loss-free DeepSeek-V3 approach a clean improvement, or does it introduce new failure modes at scale? The 2024 release is too recent for the community to have characterised this.
  3. Does routing fluctuation (StableMoE’s central object) persist at the fine-grained-expert scale of DeepSeek-V3 (256 experts), or does the higher expert count change the dynamics?

13. Implications

For applied teams. [Analysis] If training an MoE from scratch in 2026, the operational recipe is: start from ST-MoE’s stability hygiene (float32 router, α=0.01\alpha = 0.01 auxiliary loss, route-z loss), use k=2k=2 token-choice routing (matching Mixtral / production practice), and re-sweep auxiliary-loss weights for your specific compute scale. Expert Choice and StableMoE are research-grade alternatives worth knowing but not the default starting point.

For the research community. [Analysis] The three papers established the framing — token-choice vs expert-choice, routing-fluctuation as a measurable failure mode, the auxiliary-loss-tuned approach — that every subsequent MoE paper engages with. The 2024-26 frontier (auxiliary-loss-free, fine-grained-expert) builds on but does not subsume the 2022 work; readers should treat the three papers as foundational rather than current.

For evaluation methodology. [Reviewer Perspective] MoE evaluation is muddier than dense-model evaluation because total parameters, active parameters, and routing-overhead FLOPs are three distinct quantities. The community standard (“FLOPs-matched dense baseline”) is reasonable but obscures the production trade-off where the active-parameter count drives inference cost while the total-parameter count drives storage and KV-cache structure.

14. Three-depth summary

The 3-line summary for the curious reader. Most large language models are “dense” — every parameter does work on every token. “Mixture-of-Experts” models split the feedforward layer into many parallel copies and a small router decides which experts work on each token. The big problem is keeping this router stable as you train: in 2022, three papers (ST-MoE, Expert Choice, StableMoE) figured out how to do this reliably, and the recipes they established still power production MoE language models like Mixtral and DeepSeek-V3 in 2026.

The 5-line summary for the working developer. MoE routing was the brittle layer of 2020-2022 sparse models. ST-MoE (Zoph 2022) compiled the engineering recipe — float32 router, GShard-style auxiliary load-balancing loss at α=0.01\alpha = 0.01, route-z auxiliary loss to bound logit magnitude, conservative fine-tuning dropout — and produced the first sparse model to top diverse transfer benchmarks at 269B parameters. Expert Choice (Zhou 2022) inverted routing direction so experts pick tokens, guaranteeing perfect load balance by construction but introducing token-dropping at inference. StableMoE (Dai 2022) measured routing fluctuation (15.4% of tokens still changing expert assignment at the end of training) and eliminated it via a two-stage train-then-freeze procedure. For 2026 deployment, ST-MoE’s recipe + Mixtral-style k=2k=2 token-choice is the production default; DeepSeek-V3’s auxiliary-loss-free variant is the 2024 frontier evolution worth following.

The 5-line summary for the ML researcher. The 2022 trio defined three orthogonal axes for MoE routing-stability work: stability engineering (ST-MoE), structural reformulation (Expert Choice), and routing-target stationarity (StableMoE). ST-MoE’s empirical sweeps remain the canonical reference for auxiliary-loss weights and router-precision choices; the recipe transfers approximately to decoder-only frontier scale via Mixtral. Expert Choice’s structural-balance property is elegant but production adoption has been limited by token-dropping friction and variable per-token expert count. StableMoE’s routing-fluctuation measurement is the lasting methodological contribution; the frozen-router fix has limited direct adoption but influenced subsequent distilled-routing variants. 2024-2026 frontier MoE (DeepSeek-V3, Mixtral, rumoured GPT-4-class architectures) builds on but does not subsume this work; the open questions concern fine-grained-expert routing dynamics, auxiliary-loss-free balance, and routing stability at frontier (1T+) scale where direct empirical comparison is not publicly available. For follow-up work, characterising DeepSeek-V3-style fine-grained-expert routing under the StableMoE fluctuation lens is the most consequential open direction.

How this article was made: an autonomous AI pipeline researched, drafted, fact-checked, and reviewed this piece, aggregating publicly-available information from the sources consulted below. AI (artificial intelligence) can make mistakes, so please cross-check the consulted sources before acting on anything here. Neural Tech Daily is not liable for decisions or outcomes based on this article.

Sources consulted

Cited Sources

  1. 1. Zoph, Bello, Kumar, Du, Huang, Dean, Shazeer, Fedus (2022). ST-MoE: Designing Stable and Transferable Sparse Expert Models. arXiv abstract. (accessed )
  2. 2. ST-MoE — ar5iv HTML render. (accessed )
  3. 3. Zhou, Lei, Liu, Du et al. (2022). Mixture-of-Experts with Expert Choice Routing. NeurIPS 2022. arXiv abstract. (accessed )
  4. 4. Expert Choice — ar5iv HTML render. (accessed )
  5. 5. Dai, Dong, Ma, Zheng, Sui, Chang, Wei (2022). StableMoE: Stable Routing Strategy for Mixture of Experts. ACL 2022. arXiv abstract. (accessed )
  6. 6. StableMoE — ar5iv HTML render. (accessed )
  7. 7. Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. (accessed )
  8. 8. Fedus, Zoph, Shazeer (2022). Switch Transformer. (accessed )
  9. 9. Lepikhin et al. (2020). GShard. (accessed )
  10. 10. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. (accessed )
  11. 11. Mistral AI (2024). Mixtral of Experts. (accessed )

Anonymous · no cookies set

Report a problem with this article

Articles are produced by an autonomous AI pipeline; mistakes do happen. Tell us what's wrong and the editorial review will revisit the claim.

Category

Found this useful? Share it.