Mixture of Experts Lab

About

A mixture of experts (MoE) replaces one big network with several smaller ones, the experts, plus a small router that decides which experts should handle each input. With top-k routing, only k experts run for any given input, so the model can hold far more parameters than it spends on each prediction.

This lab trains a tiny MoE classifier on 2D points. Every expert is a small neural network that outputs a vote for class A or class B. The router looks at the same point and scores every expert. The left map shows what the whole model predicts; the right map shows which expert the router picks at each spot, coloured by expert.

The interesting part is the right map. Nobody tells the router how to divide the plane; the division emerges from training. You will see it carve the data into territories, sometimes neatly, sometimes by handing almost everything to one or two favourite experts. That second behaviour, called expert collapse, is one of the central engineering problems in real MoE language models, and the Balancing controls are the tools researchers invented to prevent it.

Frontier open-weight models are built this way: Mixtral 8x7B routes each token to 2 of 8 experts per layer, DeepSeek-V3 to 8 of 256 plus one shared expert, gpt-oss-120b to 4 of 128. The Mixture of Experts lesson explains why.

How it works

1. Score every expert

For a point x the router produces one logit per expert, z = router(x), and turns them into probabilities with a softmax, p = softmax(z). In an LLM the router is usually a single linear layer applied to the token's hidden state. Here it has one small hidden layer, standing in for the rich features a transformer's router would already have.

2. Pick the top k

Only the k experts with the highest scores run. With k = 1 the gate is the chosen expert's raw probability, exactly as in the Switch Transformer, so the router still receives a gradient telling it whether its choice helped. With k of 2 or more the probabilities of the chosen experts are renormalised to sum to 1, as in Mixtral and gpt-oss. “Dense” runs every expert and weights them all: a soft mixture, the 1991 original, with no compute saving.

3. Combine the votes

Each chosen expert outputs a number (its vote). The model's output is sigmoid(Σ gateᵢ · expertᵢ(x)) over the chosen experts, and the loss is binary cross-entropy. Backpropagation updates the chosen experts (weighted by their gates) and the router (through the gates). Experts that were not chosen for a point receive no gradient from it at all. That is the root of collapse: an expert that gets few points learns little, becomes less useful, and gets picked even less.

Look closely at the decision map with top-1 routing and you will often see a kink or a small step in the boundary exactly where the routing map changes colour. That is real: on one side of the line one expert decides, on the other side another, and nothing forces their answers to agree at the seam. Top-k routing makes the model's function discontinuous, one reason top-2 and dense mixtures tend to give smoother boundaries.

4. Train

Each step runs all 300 points through the model (full-batch), accumulates exact gradients and applies one Adam update with learning rate 0.03. Several steps run per animation frame; the maps are re-evaluated on a 72 × 72 grid.

Load balancing

The problem: rich get richer

Early in training some expert is, by chance, slightly better for a region. The router sends it more points, so it trains faster, so it becomes even better, so the router sends it still more. Left alone, many MoE models end up with a few overworked experts and many idle ones, wasting both parameters and hardware. Set Balancing to Off with top-1 routing and watch the load bars.

Fix 1: an auxiliary loss

The Switch Transformer adds α · N · Σᵢ fᵢ · Pᵢ to the training loss, where fᵢ is the fraction of points routed to expert i and Pᵢ is the average probability the router gave it. The published default is α = 0.01; this lab uses larger values because it trains for only a few hundred steps. The cost is that this extra loss competes with the real objective: too large an α and the router spreads points evenly even where that hurts accuracy.

Fix 2: a bias that only affects selection

DeepSeek-V3 dropped the auxiliary loss (keeping only a tiny per-sequence one). Each expert gets a bias that is added to its score when choosing the top k, but not when computing the gate weights. After every step the bias of an overloaded expert is lowered and that of an underloaded expert raised by a fixed amount γ. Because the bias never enters the loss, balancing no longer fights accuracy. In this lab the scores in Bias mode are sigmoid affinities, as in DeepSeek-V3, and γ = 0.005 per step.

Expect Bias mode to be jumpy here, especially with top-1. The update moves every expert's bias by the same fixed step, and with only 300 points and 8 experts that step can flip a whole region from one expert to another at once, so the load oscillates instead of settling. It behaves much better with top-2. DeepSeek-V3 applies the same rule while routing each token to 8 of 256 experts over batches of millions of tokens, where each nudge moves only a sliver of the traffic.

Capacity and dropped tokens

Hardware wants fixed-size work: GShard and the Switch Transformer give each expert a buffer of capacity factor × (points × k) / N slots. Points routed to a full expert are dropped from that expert. In a transformer a dropped token still passes through the residual connection, so it is not lost, just unprocessed by that layer. Here a dropped assignment simply contributes nothing, and dropped points are ringed in red on the routing map. A higher factor wastes memory on empty slots; a lower one drops more. Many recent systems, including DeepSeek-V3, balance well enough to avoid dropping entirely.

Try this

1. 01Cause a collapseChoose Moons, top-1, 8 experts, Balancing Off, then Reset a few times and train each for about 400 steps. Count how many experts are in use. Often it is 2 or 3, sometimes just 1.
2. 02Cure it liveWith a collapsed model still running, switch Balancing to Aux loss. Watch idle experts get recruited: new colours appear on the routing map and the busiest bar falls towards the fair-share line.
3. 03Loss-free balancingSwitch to top-2 and repeat with Bias instead. The bias moves by a fixed step each update, so the load wobbles around the fair share rather than settling on it, but the task loss never has to compete with a balancing term. With top-1 the wobble can grow into whole regions flipping between experts.
4. 04Linear expertsSet expert size to Linear. Each expert can now only draw a straight line, so the model can solve Spirals only if the router cuts the plane into many pieces. Compare how many experts are in use with and without balancing.
5. 05Top-1 vs top-2 vs denseTrain the same dataset with each routing mode. Dense uses every expert for every point: look at Params per point. Top-2 usually trains a little more smoothly than top-1 because each point gets a second opinion.
6. 06Squeeze the capacitySet capacity to 1.0 with Balancing Off. Red rings mark dropped points. Turn balancing on and the drops mostly disappear, because no expert is oversubscribed.
7. 07Ask the routerHover over a boundary between two colours on the routing map. The inspector shows two experts with similar scores: this is where top-k is making a close call.
8. 08Do experts specialise by class?On Rings, check whether each expert owns a class or a region. Usually it is a region: experts specialise in where a point is, not what it is, a small echo of the Mixtral finding that experts track syntax more than topic.