Architecture

Mixture-of-experts routing

How a model can have a trillion parameters and run like a much smaller one. Route tokens to experts, watch the router collapse onto favourites, and see tokens dropped at capacity.

Experts
Top-k 2
Capacity factor 1.25x
TOKENSEXPERTSt0t1t2t3t4t5t6t7t8t9t10t11t12t13t14t15t16t17t18t19t20t21t22t23E0 6/8E1 7/8E2 6/8E3 7/8E4 6/8E5 6/8E6 6/8E7 4/8red line is capacity. dashed links were dropped.synthetic router affinities, not a trained gate.
25% of experts
Expert FFN active
0%
Tokens dropped
1.17x
Load imbalance
8tokens
Capacity per expert

Each token runs 2 of 8 expert blocks. Attention and embeddings still run for every token, so this is a large saving on the feed-forward parameters rather than on the whole model.

A dense model runs every parameter for every token. A mixture-of-experts model replaces the feed-forward block with many parallel experts and a small router that picks a couple of them per token. Total capacity goes up, and the compute per token does not, because only the chosen experts run.

That is the headline and it is true. The interesting content is what has to be arranged for it to hold.

The router is learned, and a learned router left to itself collapses. Early in training a few experts are marginally better, so they get more tokens, so they train faster, so they get more tokens. Turn off the balancing loss here and watch it happen: the busiest expert takes several times the average load while others sit almost idle, which wastes most of the capacity you added the experts for. Every MoE therefore has some balancing mechanism, whether an auxiliary loss or, in more recent designs, a per-expert bias adjusted during training specifically because the auxiliary loss itself costs quality.

Then there is capacity. When experts run as fixed-size batched matrix multiplications, each has a maximum number of tokens it can accept per batch, and tokens routed to an expert that is already full are dropped. The capacity factor is the buffer, and it is a direct trade, because raising it means padding every expert with wasted slots. This is a consequence of fixed-shape batching rather than of the idea: dropless implementations avoid it using variable-size grouped matrix multiplications, and most inference stacks now do.

Dropping sounds fatal and usually is not, which is worth understanding. With top-2 routing a dropped token still gets its other expert, and the residual stream carries it forward regardless, so the system degrades rather than breaks. At top-1 there is no other expert, which is why top-1 routing is far more sensitive to capacity. Push the capacity factor below one and watch the difference.

One number to read carefully: the share of parameters this saves you is the share of EXPERT parameters. Attention, embeddings and norms run for every token no matter what the router decides, so a model that activates two of eight experts is not running a quarter of itself. Mixtral is top-2 of 8 and runs about 28% of its parameters per token.

The maths

The gate
gi  =  exp(ri)jtop-kexp(rj),y  =  itop-kgiEi(x)g_i \;=\; \frac{\exp(r_i)}{\sum_{j \in \text{top-}k} \exp(r_j)}, \qquad \mathbf{y} \;=\; \sum_{i \in \text{top-}k} g_i \, E_i(\mathbf{x})

The router scores every expert, keeps the top k, and softmaxes over only those. The output is the gate-weighted sum of the chosen experts, which is why the link opacity in the diagram follows the gate weight.

Capacity
C  =  ϕTkEC \;=\; \Big\lceil \phi \cdot \frac{T \cdot k}{E} \Big\rceil

With T tokens, E experts and capacity factor phi. Perfect balance gives each expert Tk/E tokens; phi is the buffer over that. Any routing that arrives at a full expert is dropped, which is the red line in each bar.

Load imbalance
I  =  maxeDeDˉ,De=demand before capacityI \;=\; \frac{\max_e D_e}{\bar{D}}, \qquad D_e = \text{demand before capacity}

Measured on demand, not on post-capacity load: clipping caps the busiest expert at C, so measuring afterwards makes a collapsed router look healthy. The theoretical worst case is E/k, and at the default eight experts with top-2 the collapsed router here reaches it exactly.

Related terms

More visualisations

Building with language models?

These explainers come out of the work. If you want the same thinking applied to your own system, that is what I do.

See how I can help