Replace one big MLP with many small ones plus a router, so each token only pays for the experts it actually uses. Here you train a real 8-expert, top-2 MoE layer in your browser, watch the router carve up the input space, then break it by knocking experts out.
Sign in to save progress
The bargain: parameters without the FLOPs
A transformer spends most of its parameters in feed-forward (MLP) blocks, and a dense model runs every one of those parameters for every token. A Mixture of Experts splits that block into many small expert MLPs plus a learned router, and each token runs only the top-k experts the router picks. Knowledge scales with the number of experts; compute scales with k.
Experts are just MLPs
Nothing exotic inside: each expert here is an ordinary 2 → 8 (tanh) → 1 network with 33 parameters. The power comes from having 8 of them that are free to specialize on different inputs.
The router is a learned classifier
A tiny linear layer (24 parameters) scores every expert per input and a softmax turns scores into probabilities. It is trained end-to-end: gradient flows through the gate weights of the selected experts, so routing and expertise co-evolve.
Top-2 means sparse
Only the 2 highest-probability experts run; their gates are renormalized to sum to 1 and weight the mixture. The other 6 experts cost storage but zero compute for this input.
| layer | params stored | params active / input | FLOPs / input |
|---|---|---|---|
| One dense MLP, matched size (hidden 72) | 289 | 289 | 505 |
| This MoE layer (8 experts, top-2) | 288 | 90 | 154 |
FLOP counts cover the linear layers only (each multiply-add counted as 2 FLOPs plus bias adds; tanh and softmax excluded). Same parameter budget, roughly a third of the compute per input: that ratio is the entire reason MoE exists. In the wild: Mixtral 8x7B uses this exact pattern, 8 experts with top-2 routing, storing 46.7B parameters while activating 12.9B per token (Jiang et al., 2024, arXiv:2401.04088).
Step 1: train the layer, for real, in your browser
The layer starts at a deterministic random initialization. Every number on this page, router probabilities, gates, expert outputs, RMSE, is recomputed from the current weights; nothing is scripted.
The visible ground truth the layer learns
f(x1, x2) = sin(2.5 * x1) * cos(2.5 * x2)
Training data is a fixed 16x16 grid over [-1, 1] squared, never resampled. Each button press runs real full-batch gradient descent (with momentum) on the mean squared error of the top-2 MoE prediction, plus the standard load-balancing auxiliary loss so no expert is starved. Training always uses all 8 experts; the ablation toggles below affect inference only.
Gradient steps
0
train at least 200 to unlock the gates
Grid RMSE (all 8 experts on)
0.6015
0.6015 at step 0 (untrained)
RMSE curve (recomputed after every step)
Step 2: feed inputs and watch the routing
Move the input and watch the router redistribute its softmax mass. Cross a region boundary and the top-2 set changes: different inputs literally run different parameters. Then knock out an active expert and watch the prediction and the grid RMSE degrade.
The input token: a point x = (x1, x2)
In a real transformer this would be a token's hidden vector (thousands of dimensions). Two dimensions keep every routing decision visible.
Router distribution: softmax over the 8 experts
The router always scores all 8 experts, but only the top-2 enabled ones run. Their probabilities are renormalized into the gates that weight the mixture.
The 8 experts: identical tiny MLPs, different learned jobs
output e1(x)
-0.304
output e2(x)
-0.137
output e3(x)
-0.199
output e4(x)
-0.471
contributes 0.51 × -0.471 = -0.240
output e5(x)
1.040
output e6(x)
-0.145
output e7(x)
0.827
contributes 0.49 × 0.827 = 0.404
output e8(x)
-0.645
Every expert computes its output for display, but only the highlighted top-2 contribute to the prediction (and, during training, only they receive gradient at this input). Turn an active expert Off to force the router onto its runners-up. The last enabled expert cannot be turned off.
MoE prediction ŷ
0.164
target f(x) = 0.415, error 0.251
Grid RMSE now
0.6015
all 8 experts on
Active parameters
90
of 288 stored (2 experts + router)
FLOPs this forward
154
linear layers only; tanh and softmax excluded
The layer is untrained: predictions and routing are meaningless until you train it above.
The router map: where each expert lives
The router is linear, so it partitions the input plane into convex regions, one specialist per region. This is the specialization MoE training is supposed to produce, made visible. Disable an expert and its territory is annexed by the runners-up, whose parameters never trained there.
Each cell shows the top-1 expert the router would pick there. The gold ring is your current input. Click the map to move it, or use the x1 and x2 sliders. Before training, these regions are arbitrary leftovers of random initialization.
Expert utilization: share of grid points routing each expert into the top-2 (recomputed from the router)
The load-balancing auxiliary loss keeps these shares near 25% each during training. Without it, real MoE training collapses: a few experts win early, get all the gradient, and the rest never learn.
Step 3: the sparsity check
This layer stores 8 experts of 33 parameters each and routes every input through its top-2. What fraction of the expert parameters does one input activate?