Mixture-of-Experts

Scaling laws say loss keeps falling as a power law in parameter count $N$ , but the Transformer accounting ties every one of those $N$ parameters to per-token FLOPs: a dense forward pass touches all of them. A mixture-of-experts layer breaks that tie. It holds many feed-forward networks but routes each token through only a few, so the model can carry an order of magnitude more parameters at the same arithmetic cost, the design behind the largest language models in deployment.¹

Conditional computation

A dense network applies the same weights to every input. Conditional computation instead activates an input-dependent subset of the network, so the parameters a token sees are a function of the token itself.²

The two costs decouple. Write $N$ for the total parameter count and $k$ for the number of expert sub-networks each token activates out of $E$ total. Dense computation is the special case $k = E$ ; sparse routing keeps $k$ fixed as $E$ grows.

dense FLOPs per token \propto N, sparse FLOPs per token \propto \frac{k}{E} N ≪ N when k ≪ E .

A model with $E = 64$ experts and $k = 1$ active per token carries $64 \times$ the feed-forward parameters of its dense twin while doing the same per-token arithmetic.

Dense (left) routes every token through all weights, so FLOPs track total parameters; sparse (right) activates

k

E

experts, holding FLOPs flat as

E

grows.

The mixture-of-experts layer

The layer that realizes this replaces a single feed-forward block with $E$ of them, the experts, plus a small gating network (the router) that decides how much each expert contributes.³

The experts are ordinary position-wise FFNs: each is $E_{i} (x) = W_{2}^{(i)} σ (W_{1}^{(i)} x)$ with its own weights $W_{1}^{(i)} \in R^{d_{ff} \times d}$ and $W_{2}^{(i)} \in R^{d \times d_{ff}}$ , exactly the two-layer FFN the Transformer already uses, replicated $E$ times. The router is a single linear map $W_{g} \in R^{E \times d}$ followed by a softmax. In a Transformer, the MoE layer drops in where the FFN sublayer sat, leaving attention untouched.

Dimension walk-through

Trace one token $x \in R^{d}$ through the layer. With model width $d = 4096$ , expert hidden width $d_{ff} = 4 d = 16384$ , and $E = 8$ experts routed top- $2$ :

x h = W_{g} x T = 2 arg top (h) g_{i} = softmax (h_{T})_{i} E_{i} (x) = W_{2}^{(i)} σ (W_{1}^{(i)} x) y = \sum_{i \in T} g_{i} E_{i} (x) \in R^{d} \in R^{E} \subset {1, \dots, E} \in [0, 1], i \in T \in R^{d} \in R^{d} . (d = 4096) (E = 8 router logits) (∣ T ∣ = 2) (\sum_{i \in T} g_{i} = 1) (two matmuls per active i)

The router costs $d E = 4096 \cdot 8 \approx 33 k$ multiply-adds, negligible against the $2 \cdot 2 d d_{ff} = 2 \cdot 2 \cdot 4096 \cdot 16384 \approx 0.27 B$ per active expert. Routing top- $2$ runs two experts, so per-token FFN cost is $\approx 0.54 B$ MACs, independent of how many experts the layer holds. Adding experts grows the parameter store, not this number.

An MoE layer: the router scores the token

x \in R^{d}

, two experts fire, and their outputs combine by the gate weights

g_{2}, g_{4}

into the output.

Sparse top- $k$ gating

A dense gate that runs every expert defeats the purpose. The fix is to keep only the $k$ highest-scoring experts and zero the rest before the softmax, so the combination has at most $k$ nonzero terms.¹

A worked routing example

Take $E = 4$ experts and top- $2$ routing. Suppose the router produces logits $h = (2.0, 0.5, 2.6, - 0.3)$ for one token. The two largest are $h_{3} = 2.6$ and $h_{1} = 2.0$ , so $T = {1, 3}$ . Softmax over just those two:

g_{3} = \frac{e ^{2.6}}{e ^{2.6} + e ^{2.0}} = \frac{13.46}{13.46 + 7.39} = 0.646, g_{1} = \frac{e ^{2.0}}{e ^{2.6} + e ^{2.0}} = 0.354,

with $g_{2} = g_{4} = 0$ . The token's output is $y = 0.646 E_{3} (x) + 0.354 E_{1} (x)$ . Renormalizing over the survivors (rather than the full softmax over all four) matters: the full softmax would put $\approx 0.10$ on the dropped experts, and folding that mass back onto the top two is what keeps the gate a valid convex combination.

Why gate on the softmax, not a hard argmax

The gate values multiply the expert outputs, so gradients reach the router through the $g_{i}$ . If the router just picked an expert and passed the output through unweighted, $\partial y / \partial W_{g} = 0$ and the routing decision would never learn. Weighting by $g_{i}$ makes the selection differentiable in the magnitude even though the discrete choice of $T$ is not. The original sparsely-gated layer adds tunable Gaussian noise to the logits before the top- $k$ ,

h_{i} (x) = (W_{g} x)_{i} + ε_{i} \cdot softplus ((W_{noise} x)_{i}), ε_{i} \sim N (0, 1),

which spreads tokens across experts during training (a token near a decision boundary lands on either side across steps) and gives the top- $k$ selection a smooth expected behaviour.¹

Top-

k

routing (

k = 2

): the router scores all four experts, keeps the two highest logits, and renormalizes the gate over just those before combining.

Setting $k = 1$ routes each token to a single expert, the Switch simplification; $k = 2$ is the common choice in GShard and Mixtral, giving the gate at least two gradients to weigh against each other per token.⁴⁵ The figure below contrasts the two: top- $1$ passes the whole token to one expert; top- $2$ splits it across two by the renormalized gate.

Top-

1

(left) sends the token to a single expert with gate

1

; top-

2

(right) splits it across two experts weighted by

g_{a}, g_{b}

with

g_{a} + g_{b} = 1

The load-balancing problem

Nothing in the gated objective forces the router to use all the experts. Early in training a few experts get slightly more gradient, the router prefers them, they improve faster, and the preference compounds.

The standard remedy is an auxiliary loss that rewards balanced assignment. For a batch $B$ of $T$ tokens and $E$ experts, let $f_{i}$ be the fraction of tokens routed to expert $i$ and $P_{i}$ the average gate probability mass it received.

The product $f_{i} P_{i}$ pairs the (non-differentiable) routing count $f_{i}$ with the differentiable probability $P_{i}$ , so gradients flow into the router through $P_{i}$ while $f_{i}$ scales the penalty toward whichever experts are overloaded.⁴ The $P_{i}$ term nudges the router to shave probability off experts that are already crowded; the $f_{i}$ multiplier makes that push proportional to how crowded each is. At the uniform optimum the sum equals $E \cdot E \cdot (1/ E)^{2} = 1$ , scaled by $α$ .

For example, with $E = 4$ , a fully collapsed router that sends every token to expert $1$ has $f = (1, 0, 0, 0)$ and $P \approx (1, 0, 0, 0)$ , so $L_{aux} = α \cdot 4 \cdot 1 = 4 α$ . A uniform router has $f = P = (\frac{1}{4}, \dots)$ , giving $α \cdot 4 \cdot 4 \cdot \frac{1}{16} = α$ . The collapsed state costs $4 \times$ more, and the gradient points away from it.

The coefficient $α$ is small, $1 0^{- 2}$ in Switch, so the balancing pressure never dominates the task loss; it only breaks the symmetry that would otherwise let the router collapse. Too large an $α$ forces uniform routing regardless of the token distribution, which discards the specialization the experts could learn.

A second mechanism bounds the damage when balance still fails. Each expert is given a fixed capacity, and tokens beyond it are dropped past the layer through the residual connection.

Capacity has to be fixed ahead of time because the expert weight tensors are allocated as dense $[E, C, d]$ buffers for the batched matmul, so the buffer size cannot depend on the data-dependent routing counts. With $T = 4096$ tokens, $E = 8$ experts, and $c = 1.25$ , each expert holds $C = ⌈ 1.25 \cdot 512 ⌉ = 640$ slots. A perfectly uniform batch fills $512$ of them and wastes $128$ per expert to padding; a skewed batch overflows the crowded experts and drops their excess. The capacity factor provides headroom against skew, at the cost of padded compute.

Load balancing: tokens spread evenly under the auxiliary loss (left); under collapse (right) one expert overflows its capacity and the excess tokens are dropped.

An alternative inverts the decision. Token-choice routing has each token pick its experts; expert-choice routing has each expert pick its top- $C$ tokens, which guarantees perfect balance by construction (every expert fills exactly its capacity) at the price that some tokens may be chosen by many experts and others by none.⁶

Architectures

The same gated layer appears in a line of progressively simpler and larger models.

Sparsely-gated MoE (2017). The first practical instance interleaved an MoE layer between stacked LSTM layers in a language and translation model, with up to $131, 072$ experts and noisy top- $k$ gating ( $k = 4$ ), reaching $137$ B parameters at the compute of a far smaller dense model.¹

GShard (2021). Ported the MoE layer into the Transformer, replacing every other FFN sublayer with a top- $2$ MoE, and added the sharding and capacity machinery to train a $600$ B-parameter multilingual translation model across thousands of devices.⁵

Switch Transformer (2022). Simplified routing to $k = 1$ : each token goes to exactly one expert. This halves the routing and communication cost of top- $2$ , and the paper shows that with the right capacity factor and auxiliary loss, top- $1$ matches or beats top- $2$ quality, scaling to $1.6$ T parameters.⁴

The two all-to-all shuffles are the crux of the distributed cost: the first (dispatch) sends each token to the device holding its expert; the expert runs locally; the second (combine) returns the weighted output to the token's home device. Because routing is data-dependent, the volume each device sends is uneven, and a single overloaded expert stalls the collective while the others wait. This is the communication reason to care about balance, on top of the wasted-capacity reason.

Expert parallelism across two devices: dispatch all-to-all sends each token to its expert's device, experts run locally, and combine all-to-all returns outputs.

The forward pass of one MoE layer makes the routing, dispatch, and combine explicit.

Algorithm:

\textsc{MoELayer}(X, k)

— sparse top-

k

mixture-of-experts forward pass

1
input batch $X = (x_1, \dots, x_T)$ ; experts $E_1, \dots, E_E$ ; capacity $C$
2
$\mathit{load}_i \gets 0$ for each expert $i$
tokens assigned so far
3
for each token $x_t$ in $X$ do
4
$h \gets W_g\, x_t$
router logits
5
$\mathcal{T} \gets \argtop_k(h)$
the $k$ highest-scoring experts
6
$g \gets \softmax(h_{\mathcal{T}})$
renormalize over the survivors
7
$y_t \gets 0$
8
for each expert $i$ in $\mathcal{T}$ do
9
if $\mathit{load}_i < C$ then
expert still has a free slot
10
$y_t \gets y_t + g_i \cdot E_i(x_t)$
11
$\mathit{load}_i \gets \mathit{load}_i + 1$
12
else
13
$y_t \gets y_t + x_t / k$
dropped: pass through residual
14
return $(y_1, \dots, y_T)$ and the batch loads for $\mathcal{L}_{\text{aux}}$

In practice the per-token loop is a single batched gather/scatter: tokens are sorted by their chosen expert, dispatched in one all-to-all, processed as dense matrix multiplies inside each expert, and scattered back. The dense expert matmul operates on the padded $[E, C, d]$ buffer, so its shape and FLOP count are fixed at compile time regardless of the realized routing.

Practical issues

Sparse models trade their FLOP savings for a set of engineering and statistics problems the dense model never had.

Training instability. The router's argmax is discontinuous, so small logit changes flip a token between experts and produce loss spikes; the ST-MoE recipe adds a router $z$ -loss, $L_{z} = \frac{1}{T} \sum_{x} (log \sum_{i} e^{h_{i} (x)})^{2}$ , which penalizes large logits and keeps the softmax in a numerically stable range, the single most effective stabilizer reported.⁷ The $log \sum e^{h_{i}}$ term is the log-partition of the router softmax; squaring and minimizing it pulls the logits toward a smaller magnitude, so a single flip moves the gate values less and the loss surface stays smoother.

Fine-tuning. Sparse models overfit downstream tasks faster than dense ones of equal quality, because each expert sees only a fraction of the (already small) fine-tuning set; ST-MoE recommends fine-tuning a subset of parameters and a smaller auxiliary-loss weight.⁷

Expert specialization. Trained experts do specialize, though rarely along human-legible lines; in Mixtral the routing is largely uniform across experts and only weakly correlated with topic or syntax, so the gain comes from added capacity more than from clean division of labor.⁸

Sparse upcycling. Rather than train an MoE from scratch, copy a trained dense FFN into $E$ identical experts, add a fresh router, and continue training; the copies diverge under the load-balancing pressure, recovering most of the gain at a fraction of the from-scratch cost.

Inference routing cost. At serve time the FLOP savings are real but the full parameter set must be resident in memory, and the all-to-all dispatch adds latency and makes throughput sensitive to how evenly a batch happens to route. The binding constraint shifts from compute to memory and communication.

Mixtral (2024). A modern decoder-only LLM, Mixtral 8×7B, places $E = 8$ experts with top- $2$ routing in every layer. It holds $47$ B total parameters but activates only $13$ B per token, matching or beating a $70$ B dense model at the inference cost of a $13$ B one.⁸

Dense vs sparse: the accounting

The three routing regimes differ only in $k$ , but that one number sets active parameters, communication, and the headline parameter count.

Model	Routing	Active params / token	Total params	Key idea
Dense Transformer	all FFNs	$N$	$N$	one FFN per layer, every weight used
Switch Transformer	top- $1$	$\approx N / E$	up to $1.6$ T	one expert per token; cheapest routing
GShard / Mixtral	top- $2$	$\approx 2 N / E$	$600$ B / $47$ B	two experts per token; more gate signal

The pattern is constant per-token cost with growing capacity: each row adds experts (total params) without adding much active compute, exactly the decoupling conditional computation promised. The figure below sizes the trade for one FFN sublayer: the dense block uses all its parameters every token; the sparse block stores $E$ times as many but touches only $k$ of them.

Dense FFN uses all its parameters per token (params = FLOPs); the MoE stores

E

experts but activates

k

, so params scale with

E

while FLOPs scale with

k

Fine-grained and shared experts

Goodfellow describes the classic gated mixture; the design that ships in the largest open models since makes two modifications to the top- $k$ layer, both aimed at getting more specialization out of the same active-parameter budget.⁹

Fine-grained experts. Instead of $E$ experts each of hidden width $d_{ff}$ routed top- $2$ , split each expert into $m$ slices of width $d_{ff} / m$ , giving $m E$ small experts routed top- $2 m$ . The active-parameter count is unchanged — the same total hidden width fires per token — but the router now chooses from a much larger menu, so each token can assemble a more precise combination of specialists. The intuition is combinatorial: routing $8$ of $64$ small experts gives far more distinct expert-subsets than routing $2$ of $8$ large ones, and finer partitions let experts specialize more sharply.

Shared experts. A fine-grained MoE also keeps one or two shared experts that every token passes through, alongside the routed ones. The shared expert absorbs the common computation that all tokens need — the syntax and general structure that would otherwise be redundantly relearned by every routed expert — so the routed experts are freed to specialize on what actually differs between tokens. This removes the redundancy that makes plain MoE experts specialize only weakly.

Fine-grained plus shared experts. Every token passes through a shared expert (always on); the router additionally picks several fine-grained experts from a large pool, so active compute is unchanged but the routed menu is far larger.

These two changes, on top of the auxiliary-loss and capacity machinery derived above, are why modern MoE LLMs push the active-to-total ratio far lower than Mixtral's $13/47$ while keeping quality: a large pool of finely-sliced routed experts, a shared expert absorbing common work, and the same $z$ -loss and load-balancing machinery. The layer is the same $y = \sum_{i} g_{i} (x) E_{i} (x)$ ; only the granularity and a fixed always-on term changed.

Takeaways

Conditional computation activates an input-dependent subset of weights, so total capacity $N$ and per-token FLOPs ( $\propto \frac{k}{E} N$ ) become separate dials.
A mixture-of-experts layer replaces one FFN with $E$ experts and a router, outputting $y = \sum_{i} g_{i} (x) E_{i} (x)$ ; it drops in where the Transformer FFN sat.
Sparse top- $k$ gating keeps the $k$ largest router logits, renormalizes the softmax over them, and zeroes the rest, so only $k$ experts run per token.
Without intervention the router collapses onto a few experts; the load-balancing loss $α E \sum_{i} f_{i} P_{i}$ is minimized at the uniform split $f_{i} = P_{i} = 1/ E$ .
A capacity factor $c$ caps tokens per expert and drops the overflow through the residual; expert-choice routing balances by construction.
The architectures simplify and grow: sparsely-gated LSTM MoE, GShard (top- $2$ Transformer, $600$ B), Switch (top- $1$ , $1.6$ T); experts are sharded by expert parallelism with two all-to-all shuffles per layer.
Sparse models need a router $z$ -loss for stability, need careful fine-tuning, and shift the serving bottleneck from compute to memory and communication.
Mixtral 8×7B ( $47$ B total, $13$ B active, top- $2$ ) is the canonical modern MoE, matching a $70$ B dense model at a fraction of the inference compute.
Modern designs: fine-grained experts (many thin experts, larger routed menu, same active budget) and shared experts (one always-on expert absorbing common computation) let modern MoE LLMs push the active-to-total ratio far below Mixtral's while keeping quality.

Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR 2017 — sparse noisy top- $k$ gating between LSTM layers, the load-balancing loss, and the first MoE at hundreds of billions of parameters. ↩ ↩² ↩³ ↩⁴
Goodfellow, Deep Learning, §12.4.3 — conditional computation and mixtures of experts: gating networks that activate input-dependent sub-networks to decouple capacity from per-example cost. ↩
Jacobs, Jordan, Nowlan & Hinton, Adaptive Mixtures of Local Experts, Neural Computation 1991 — the original mixture-of-experts: a gating network softly partitions the input space among specialist sub-networks trained jointly. ↩
Fedus, Zoph & Shazeer, Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR 2022 — top- $1$ routing, the differentiable load-balancing loss, and scaling to $1.6$ T parameters. ↩ ↩² ↩³
Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, ICLR 2021 — top- $2$ MoE Transformer with capacity, automatic sharding, and all-to-all dispatch for a $600$ B-parameter translation model. ↩ ↩²
Zhou et al., Mixture-of-Experts with Expert Choice Routing, NeurIPS 2022 — inverts routing so each expert selects its top tokens, guaranteeing perfect load balance without an auxiliary loss. ↩
Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022 — the router $z$ -loss for training stability and a fine-tuning recipe for sparse models. ↩ ↩²
Jiang et al., Mixtral of Experts, 2024 — a decoder-only LLM with $8$ experts and top- $2$ routing per layer, $47$ B total / $13$ B active parameters, matching a $70$ B dense model. ↩ ↩²
Dai et al., DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, 2024 — fine-grained expert segmentation plus isolated shared experts, raising specialization at a fixed active-parameter budget. ↩