Mixture of Experts

Dive into Deep Learning · §11.6

Mixture of experts
stored and active parameters · token-choice routing · load balancing · MoE in a GPT

Conditional computation

In a dense model, every token uses every layer parameter.
MoE breaks the lockstep at the FFN (two thirds of each block): store E experts, route each token to k of them.
MoE separates stored capacity from active parameters and per-token FLOPs.

Stored and Active Parameter Counts

        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

Replacing the Block FFN

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.

Positive feedback in routing

A slightly-lucky expert improves on its tokens and gains routing probability — more tokens, more gradient, more probability.
Less-used experts receive fewer updates, which can amplify load imbalance. In routing collapse, a few experts receive most assignments.
Sparse MoE systems commonly add capacity constraints or balancing methods (Shazeer et al., 2017).

Two repairs

Auxiliary loss (GShard, Switch): penalize load-probability correlation, \mathcal{L}_{\mathrm{balance}} = E \sum_i f_i\, \bar{p}_i — differentiable, but it adds a gradient term to the language-model objective.

Auxiliary-loss-free bias (Wang et al., 2024; DeepSeek-V3): a controller, \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.

Balancing methods at fixed compute

The runs share their initialization, data, and 800-step budget; only the balancing method differs:

  no balancing: training loss 0.81, experts under 2% usage: 7 of 16
auxiliary loss: training loss 0.42, experts under 2% usage: 0 of 16
          bias: training loss 0.52, experts under 2% usage: 0 of 16

Experimental results

<matplotlib.legend.Legend at 0x776f5210f080>

No balancing: in the displayed seed, several experts receive less than 2% of assignments.
Both repairs: the displayed loads are closer to the uniform reference and training loss is lower than in the unbalanced run.
One run per method gives no uncertainty interval and does not distinguish auxiliary loss from bias control.

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 bias controller, match active parameters against the dense GPT of the previous sections:

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

5–6 times the stored parameters at the same active parameter count.
In the fixed-seed run, best validation losses are both near 1.5 and the MoE ends with slightly lower training loss; neither equivalence nor cause is established.
Published large-scale systems use the same conditional-computation mechanism; this small experiment does not estimate their quality or efficiency gains.

Recap

MoE = E FFNs + a router; parameters scale with E, FLOPs with k — Mixtral 46.7B/12.9B, DeepSeek-V3 a 28x gap.
Positive feedback can concentrate routing on a small subset of experts.
An auxiliary loss changes the training objective; a bias controller changes selection without adding a gradient term. Both reduce imbalance in the fixed-seed experiment.
Fine-grained experts provide more routing combinations at a fixed active budget; a shared expert handles processing common to all tokens.
In our GPT: several times the stored parameters at matched active parameter count; quality comparisons remain local to the teaching experiment.