II · THE IDEA · ARTIFICIAL INTELLIGENCE
MoE Routing Strategies
▶ Listen · narrated
Making a neural network wider eventually costs more than anyone can pay. Sparse routing takes another path: keep adding expert sub-networks, but let each input wake only a few of them.
At a glance
- What it is
- A learnable gate selects a sparse set of expert sub-networks per input
- Typical gate
- Noisy scores, top-k keep, then softmax over the survivors
- Compute pattern
- Parameters grow with expert count; work stays near top-k width
- Balance
- Auxiliary losses discourage collapse onto a few experts
- At device scale
- Capacity limits and automatic sharding keep routing feasible
A hospital has many specialist wards and one triage desk. Every patient meets the desk; the desk sends them to one or two wards; the other wards do nothing for that patient. The building can hold a great many specialists without every specialist seeing every case.
A mixture-of-experts layer works the same way inside a neural network. A small gating network looks at each input and gives every expert sub-network a score. It keeps only the best one or two, converts those scores into weights, and blends the chosen experts' outputs. The experts it passed over do no arithmetic at all for that input, which is why the model can store far more than it runs.
Left alone, the gate would keep picking the same few experts, because an expert that gets used gets better, and an expert that gets better gets used. Extra penalties during training push the gate to spread the traffic instead. And when experts sit on different machines, each expert is capped at a set number of inputs per batch, so no one machine is swamped. The point is controlled sparsity: many more stored specialists, roughly the same work per input.
The sparsely-gated MoE layer (Shazeer et al.) computes y = Σ_i g(x)_i · E_i(x) with g(x) sparse. Logits are H(x) = x · W_g + StandardNormal · Softplus(x · W_noise). KeepTopK zeros all but the k largest logits, and a softmax over the remainder gives g(x). Gradients reach the gate through the soft weights on the selected experts. Two auxiliary losses are added to the objective with a scalar coefficient: importance, the batch-averaged gate probability per expert, and load, the number of selections or a smooth surrogate for it. Both penalise uneven utilisation, and both are cheap relative to expert compute.
GShard puts such layers in the position-wise feed-forward slots of a Transformer and typically sets k = 2. Tokens are dispatched independently. With experts sharded across devices, each expert holds capacity ≈ (tokens_per_batch / num_experts) · capacity_factor; assignments beyond that are dropped or refused, which bounds buffer sizes and all-to-all traffic. Automatic sharding compiles the conditional pattern across accelerators instead of requiring a hand-written partition per configuration.
The limitations follow directly from the mechanism. Top-k is a non-smooth decision, softened only through the retained weights. The balance coefficient trades sharpness of specialisation against utilisation. The capacity factor trades hardware predictability against the fraction of gate decisions actually honoured. Per-token compute scales with k and expert width, not with total expert count — the intended scaling property, and one that holds only while the auxiliaries and capacity rules keep the router clear of both collapse and chronic overflow.
Look closer
The gate is itself a trained network
Routing is not a fixed hash or a hand-written rule. A small network maps each incoming hidden state to a vector of expert scores. Those scores are sparsified — commonly by keeping only the top one or two — and the survivors become combination weights. Because the gate is differentiable through the soft weights on the chosen experts, the same back-propagation pass that updates the experts also reshapes who gets chosen next time.
Noise and auxiliaries keep experts alive
A pure greedy gate tends to reinforce early winners: busy experts improve faster, so the gate prefers them still more. The sparsely-gated MoE layer adds tunable noise to the logits before top-k selection, which forces occasional trials of quieter experts during training. Separate auxiliary losses then penalise imbalance in how often each expert is selected and in the load implied by the soft gate probabilities, so utilisation stays nearer to even across the batch.
Capacity is a hard ceiling on the soft choice
When experts live on different devices, as in GShard-style layouts, each expert can accept only a bounded number of tokens per batch — its capacity, set by a capacity factor relative to a perfectly balanced share. Tokens that the gate assigns to an already-full expert may be dropped or deferred. The learned router therefore operates inside a systems constraint: soft preference first, then a hard quota that keeps communication and memory predictable.
The story
A sparsely-gated mixture-of-experts layer pulls apart two things that normally rise together: how much a model stores, and how much work it does on each example. The layer holds a pool of experts — small sub-networks of identical shape, any one of which could process the input — plus a much lighter gating network, whose only job is to choose. For each input the gate produces one score per expert. Only the few highest scores are kept; the rest are set to zero. The layer's output is a weighted sum of the chosen experts' outputs alone. An expert that was not chosen computes nothing at all for that example. That is the whole bargain. Adding experts adds parameters, because the weights must be stored. It does not add work, because the extra experts sit idle.
In the design from Shazeer and colleagues, the gate does more than score. A second projection produces a noise scale for each expert — a number saying how much random jitter to apply — passed through a softplus function so that the scale is always positive. Random noise of that size is added to the scores before the top few are kept, and a softmax over the survivors turns them into mixing weights. The noise matters only during training. It makes the ranking wobble, so an expert that is currently second-best is sometimes picked anyway. Without that wobble the gate settles on a small clique while its judgements are still close to random.
Even with noise, sparse routing tends to collapse onto a handful of favourites, and the reason is a plain feedback loop. An expert that is chosen often is trained often, so it improves, so the gate scores it higher and chooses it more. An expert that is rarely chosen barely changes, stays weak, and keeps being skipped. The same paper therefore adds auxiliary losses: extra penalty terms bolted onto the training objective. One measures, across a batch, what share of examples each expert actually received. The other measures the load implied by the gate's soft probabilities before the cut to top-k. A lopsided spread raises the penalty, so gradient descent pushes traffic back outwards. Computing those statistics is cheap next to running the experts themselves, and it keeps more of the pool learning.
GShard carries the same idea into deep Transformer stacks built for translating many languages at once. Mixture-of-experts layers replace some of the ordinary feed-forward sublayers, and each token is routed on its own. The design usually keeps 2 experts per token rather than 1, so every token gets a second opinion instead of a single hard verdict. Because different experts live on different accelerator chips, each expert is given a capacity: a fixed ceiling on how many tokens it may accept in one batch. Automatic sharding then describes the model and its routing pattern in a form a compiler can split across chips, so the layout need not be rewritten by hand every time the expert count changes.
The scaling claim is modest in mechanism and large in effect. More experts means more stored parameters. The arithmetic and the activation memory spent per token follow only the width of the few experts that actually run, plus the small gate. Balance losses and capacity ceilings are what make that bargain survive contact with real hardware. Without balance losses, most experts starve. Without capacity ceilings, the popular experts overflow the chips holding them.
Routing strategy is a bundle of choices, not one switch. How many experts a token may keep, whether noise is added before the cut, how hard balance is enforced, and how tight the capacity ceiling sits together decide whether stored capacity becomes usable quality or merely idle weights and discarded tokens. The 2017 sparsely-gated layer and the later GShard stack fixed workable settings for language modelling and large-scale translation. The template they left behind — a learned gate, a top-k cut, auxiliary balance losses, capacity-aware dispatch — is what later systems still rearrange.
Why it mattered then
By the mid-to-late 2010s, the easy route to better quality — making dense networks wider — was running into memory and compute budgets. Conditional computation, the idea of running only part of a model on any given input, had long looked attractive on paper. Training it was the hard part. The router makes a discrete choice, which gradients do not pass through cleanly; gates collapsed onto a few experts, most experts idled, and spreading experts across machines meant writing device layouts by hand. The sparsely-gated MoE layer answered each of those: noise in the scores to stop early lock-in, gradients carried by the soft weights on the chosen experts, and auxiliary penalties to force even use. That was enough to train enormous expert pools on ordinary language-modelling and translation objectives. GShard then showed the same pattern could be stacked through a giant multilingual translator, with per-expert capacity limits and automatic sharding doing the work that bespoke device code used to.
Why it matters now
The squeeze those papers addressed has not eased. Parameter count is still a reliable lever on quality, while training and serving budgets are still bounded by arithmetic throughput, memory bandwidth, and the links between chips. Sparse mixture-of-experts layers remain one of the few widely used ways to grow what a model stores faster than what it computes per token. The machinery in current systems is the direct descendant of this work: a learned gate, a cut to the top few experts, load-balancing penalties, and capacity ceilings once experts are spread across devices. Reading a modern MoE design means reading rearrangements of that template — a different number of experts kept, different balance coefficients, a different capacity ceiling — resting on the same bet that a small trained router can put a large idle pool to good use.
The surprising detail
What the gate wants is only half the story once experts sit on separate chips. In GShard-style training each expert may accept only a fixed number of tokens per batch. If the gate sends a token to an expert that has already filled its quota, that token can simply be dropped rather than sent elsewhere. So a hardware quota sits in the middle of an otherwise end-to-end trained router: the model can prefer an expert and still be refused. That collision between a soft learned preference and a hard fixed ceiling is easy to miss if one treats MoE as a neat module rather than as a dispatch problem across machines.
What is disputed
These papers demonstrate that sparse MoE layers can train at very large parameter counts and improve translation and language-modelling quality under fixed compute budgets. They do not settle how far expert specialisation is linguistic or semantic in any strong sense, versus a more opaque partition of the representation space. Claims about what individual experts “mean” should be treated as interpretive, not as results established here.
Remember this
MoE grows capacity by learning which few experts run per input. Noise, balance penalties and capacity ceilings are what stop that sparsity collapsing onto a favoured handful.
Test yourself
A team trains an MoE layer with top-1 gating and no auxiliary balance loss. Early in training a handful of experts receive most tokens. Explain two distinct ways this harms the final model, one about unused parameters and one about the gate’s later behaviour.
First, the neglected experts barely update, so the parameters spent on them contribute little useful function — capacity on paper, not in the trained network. Second, the gate’s preference is self-reinforcing: stronger experts produce better local loss reduction, so gradient updates push the gate to choose them still more often, locking in the early imbalance. Load-balancing auxiliaries and, in the 2017 design, gating noise exist specifically to break that loop while the experts are still plastic.
Go deeper
- [1701.06538] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer · arxiv.org
- [2006.16668] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.