Mixture of Experts

Dive into Deep Learning · §11.6

Mixture of experts
parameters without FLOPs · a token-choice router · collapse and two repairs · MoE in our GPT

Conditional computation

A dense model spends parameters and compute in lockstep: every token pays for every weight.
MoE breaks the lockstep at the FFN (two thirds of each block): store E experts, route each token to k of them.
Stored parameters are cheap (memory); per-token FLOPs are the expensive resource — capacity and serving cost decouple.

The accounting, computed not asserted

        ours:   16.8M in experts,   2.1M active per token, ratio  8.0x
Mixtral 8x7B:   45.1B in experts,  11.3B active per token, ratio  4.0x
     with attention and embeddings: 46.7B total, 12.9B active
 DeepSeek-V3:  656.5B in experts,  23.1B active per token, ratio 28.4x
(656463101952, 23095410688)

Mixtral 8x7B: our census + attention + embeddings = 46.7B / 12.9B — the model card, reproduced.
DeepSeek-V3: 656B of 671B sits in experts; a token touches 23B — a 28x capacity-to-compute gap.

A token-choice top-k layer

\mathrm{MoE}(\mathbf{x}) = \sum_{i \in \mathcal{E}_k(\mathbf{x})} p_i(\mathbf{x})\, \mathrm{FFN}_i(\mathbf{x}), \qquad \mathbf{p} = \mathrm{softmax}(\mathbf{W}_r \mathbf{x})

Selection (argmax): no gradient. The weight p_i: gradient — that is how the router learns (Switch convention).
Teaching implementation: compute all experts, mask by the gate — exact math, none of the FLOPs savings (real systems: scatter/gather with capacity limits).

class MoELayer(nnx.Module):
    """Mixture-of-experts FFN: a token-choice top-k router over E experts."""
    def __init__(self, num_hiddens, num_experts, num_active, rngs=None):
        rngs = nnx.Rngs(0) if rngs is None else rngs
        self.num_experts, self.num_active = num_experts, num_active
        self.router = nnx.Linear(num_hiddens, num_experts, use_bias=False,
                                 rngs=rngs)
        self.experts = nnx.List([d2l.FeedForward(num_hiddens, rngs=rngs)
                                 for _ in range(num_experts)])
        self.expert_bias = nnx.Variable(jnp.zeros(num_experts))
        self.usage = nnx.Variable(jnp.zeros(num_experts))
        self.aux_loss = nnx.Variable(jnp.zeros(()))

    def __call__(self, X):
        probs = jax.nn.softmax(self.router(X), -1)        # (B, T, E)
        scores = probs + self.expert_bias[...]            # selection only
        _, idx = jax.lax.top_k(scores, self.num_active)   # (B, T, k)
        mask = jax.nn.one_hot(idx, self.num_experts).sum(-2)
        gates = probs * mask                              # weight = p_i
        Y = jnp.stack([e(X) for e in self.experts], -1)   # (B, T, d, E)
        out = (Y * gates[..., None, :]).sum(-1)
        frac = mask.sum((0, 1)) / mask.sum()              # realized load
        self.usage[...] = self.usage[...] + mask.sum((0, 1))
        self.aux_loss[...] = self.num_experts * (
            frac * probs.mean((0, 1))).sum()
        return out

It drops into the block

Experts are d2l.FeedForward; the layer maps (n, d) \to (n, d), so it enters d2l.TransformerBlock through ffn_factory — the seam the block was built with:

router 2048, per expert 524544, total 4198400, active 1051136
usage at initialization: 0.13 0.11 0.13 0.13 0.15 0.12 0.11 0.12

At initialization: usage nearly uniform. Training will not keep it so.

Routing collapse: the rich get richer

A slightly-lucky expert improves on its tokens and gains routing probability — more tokens, more gradient, more probability.
Losers starve, learn slowly, sink. End state: a couple of experts serve everything; the advertised capacity is dead weight.
Observed since the first sparse MoE (Shazeer et al., 2017); every deployed MoE ships a countermeasure.

Two repairs

Auxiliary loss (GShard, Switch): penalize load-probability correlation, \mathcal{L}_{\mathrm{balance}} = E \sum_i f_i\, \bar{p}_i — differentiable, but its gradient argues with the LM objective on every token.

Aux-free bias (Wang et al., 2024; DeepSeek-V3): a thermostat, \mathcal{E}_k = \operatorname{argtop}_k(p_i + b_i), \qquad b_i \leftarrow b_i + u\,\mathrm{sign}(\bar{f} - f_i) — steers selection only; the loss contains no balancing term at all.

Three runs, one budget

Same init, same data, same 800 steps; only the balancing differs:

  no balancing: training loss 0.74, experts under 2% usage: 6 of 16
auxiliary loss: training loss 0.53, experts under 2% usage: 0 of 16
          bias: training loss 0.46, experts under 2% usage: 0 of 16

Reading the triptych

<matplotlib.legend.Legend at 0x765ee00d91f0>

No balancing: collapse in every seed, both frameworks — from two dead experts up to a layer served by a single one.
Both repairs: usage flat on the uniform line, zero dead experts, every balanced run beats every collapsed one.
Aux vs. bias: loss differences within seed noise at this scale; DeepSeek’s preference rests on the interference argument, a frontier-scale margin invisible from here.

Fine-grained and shared experts

DeepSeek: same budget, narrower experts, larger E and k — \binom{64}{8} \approx 4 \times 10^9 combinations vs. \binom{8}{2} = 28: specialists that compose.
Plus one shared expert, always on: common processing lives once; routed width spent purely on specialization.

Model	Experts	k	width/d	shared
Mixtral 8x7B	8	2	3.5	none
DeepSeek-V3	256	8	0.29	1
Qwen3-235B	128	8	0.38	none
gpt-oss-120b	128	4	1.0	none

MoE in our GPT

Swap every block’s FFN (moe_gpt), balance with the thermostat, match active parameters against the dense GPT of the previous sections:

dense:  3.16M parameters (3.16M active), best validation 1.49, final training loss 0.20
  MoE: 17.85M parameters (3.16M active), best validation 1.49, final training loss 0.17

5-6x the parameters at the same per-token compute.
Best validation: indistinguishable. Training loss: MoE slightly lower — spare capacity memorizes faster on 180 KB.
The mechanism transfers to the frontier; the payoff needs frontier data.

Recap

MoE = E FFNs + a router; parameters scale with E, FLOPs with k — Mixtral 46.7B/12.9B, DeepSeek-V3 a 28x gap.
Routers collapse without help: rich-get-richer dynamics kill most experts.
Repair by auxiliary loss (gradient, interferes) or by a bias thermostat (no gradient, selection only) — both flatten usage; at our scale, indistinguishable losses.
Fine-grained + shared experts: the same budget, sliced for composability.
In our GPT: several times the parameters at matched active compute, comparable training — the economics behind much of the 2026 frontier.