Mixture-of-Experts

Holding many feed-forward networks but running only a few per token.

A dense model runs every weight for every token. A Mixture-of-Experts (MoE) layer changes one thing — it replaces the single feed-forward network with a bank of many experts (each an ordinary feed-forward network) plus a small router that, for each token, picks only the top-k experts to actually run. Attention usually stays dense; it's the feed-forward that goes sparse.

That single change decouples two numbers a dense model keeps equal: capacity (all the experts you store) from compute (the few that fire per token). It's how a model can be "30B total, 3B active" — large in memory, small in per-token cost. The pages below build the idea up one piece at a time, each standalone:

The router — the gating network that scores experts and picks the top-k.
An expert — why an "expert" is just the feed-forward network you already know.
Sparse activation — running only k / E of the experts, and the two parameter counts it creates.
Load balancing — keeping the router from overusing a few experts, and the auxiliary loss that helps.

For how MoE compares to the dense base model, and how its parameter counts are named, see Dense vs Mixture-of-Experts and Active vs effective parameters.

The real routing here is illustrated with seeded weights (the decoupled demo in src/variants/), so the figures show the mechanism — route, fire a sparse few, renormalize, balance — not trained quality.

Load balancing

Keeping the router from overusing a few experts and starving the rest.

Routing is decided per token, and nothing in the bare mechanism forces the work to spread evenly. Left alone, a router tends to develop favourites: some experts run noticeably more often than others. Even a few tens of percent of imbalance is a problem from both ends. Under-used experts are wasted capacity — parameters you stored and never exercise. Over-used experts are a bottleneck: in real MoE systems each expert has a fixed capacity per batch, and tokens routed to an already-full expert get dropped (passed through unprocessed), which hurts quality. And the skew compounds during training — a busy expert improves faster, so the router favours it even more. So an uneven load is not just untidy: it costs you both capacity and accuracy.

The fix is a load-balancing auxiliary loss added during training. Alongside the normal language-modelling loss, the model is penalized when the router's traffic is lopsided — roughly, the product of how often each expert is picked and how much probability the router assigned it. Minimizing that term pushes the router toward an even split, so all experts learn and all stay useful. It's a gentle pressure, not a hard constraint: the router still specializes, it just can't collapse onto a few experts.

The figure routes many tokens through a seeded router and tallies how often each expert fires. The dashed line is the perfectly balanced share (total activations ÷ experts); bars above it (amber) are over-loaded, below it (teal) under-used. Change the token count and the shape persists — imbalance is structural, not noise — and reseed to see a different draw show the same lopsidedness. Raise top-k and the load evens out a little (more experts share each token), but the gap between busiest and idlest rarely vanishes on its own. That stubborn spread is exactly what the auxiliary loss exists to flatten.

Per-expert firing counts over many routed tokens. The dashed line is the balanced share; the busiest-vs-idlest spread is what a load-balancing auxiliary loss is trained to shrink.

So load balancing is the piece that makes the capacity-vs-compute bargain actually pay off: it's no use holding a large pool of experts if the router only ever calls on a few of them.