Learn machine learning · part 10 · session 1
The router
Eight small networks, a router that picks two of them per input, and two parameter counts that are usually reported as one number.
After this session you will be able to compute the parameters a top-k layer stores and the parameters it runs, and say why those are different numbers.
The gate: the hand computation returned exactly, without the third-attempt assist, and all three predictions about raising k correct in one submission.
01 · The question
Only two desks wake up
illustrativea drawing, with this layer's own two countsAll 8 desks are stored, and one piece of text arrives.
Large models got cheap to run by not running all of themselves. Only a slice of the machine wakes up for each piece of text; this session is that trick, at eight desks.
A mixture-of-experts layer holds several small networks, the , and a . The router produces one score per expert for each input, the two best scores are kept and rescaled so they add to 1, and only those experts run. Nothing in that arrangement assigns an expert to a human topic. The router is trained with everything else, and the assignment it ends up with is whatever lowered the loss.
Here, eight experts of 99 parameters each sit behind a router of 24 on a two-dimensional task with six regions and three classes. The course script trained that layer once, for 1,500 steps at seed 7, and this page plays the run back: press play and watch the assignment map settle and the inputs stream to the experts that end up taking them.
Two counts matter and get quoted as one: how many numbers the layer stores, and how many actually run for one input. The second is what a report means by . Then you set k and read both counters, along with the operations they cost. Those are arithmetic, and they answer while you type.
What was recorded
recorded runrecord-runs.mjs · seed 7 · k = 2One run of 1,500 steps at the top 2, with the training loss at every step, ten checkpoints carrying the eight routing shares and a 16 by 16 assignment map, and the 816 trained parameters as int8. The script is in the repository at scripts/ml/grids/the-router.mjs. Nothing on this page trains.
The data is 800 made-up points in a square, split into six regions, with 2 in 100 labels deliberately wrong. Each point is given to the nearest of six fixed centres, the square runs from −1 to 1 on both axes, the first 600 train and the last 200 are held out, and the generator is multiRegion(800, seed 7).
02 · Predict
Before anything runs
Rank these three by the number of parameters that run for one input, largest first.
Put these in the order you expect
Focus an item, press Enter to lift it, move it with the arrow keys, then press Enter again.
03 · Try it
1, 2, 4 and 5 · play the run, and watch where the inputs go
The curve, the shares, the map and the traffic all come from the same recording. The transport chooses a step; it does not run anything.
Loading the recorded run.
Play is off until the ranking above is recorded. That order is the whole mechanic.
3 and 6 · set k, and read the counters
Moving k re-evaluates the recorded weights over the 200 held-out points. It changes how many experts run for each input, not what the map shows: the map is the expert each point is routed to first, and the first choice is the same at every k.
Parameters stored, parameters active
computed in this browserarithmetic on k, E, R and nA value outside 1 to 8 is clamped to the range.
- stored
- 816
- active for one input
- 222
- operations per input
- 1,776
At k = 2: active = k × E + R = 2 × 99 + 24 = 222, and stored = n × E + R = 8 × 99 + 24 = 816. Operations are 8 × 222 = 1,776.
The held-out score at this k is still being worked out.
The held-out score is highest at the k the layer was trained at. Running more experts for each input than the training used costs operations and gives nothing back here, because the gate weights and the experts were fitted together at k = 2.
The whole k table for this layer
| k | active = k × 99 + 24 | operations = 8 × active | stored |
|---|---|---|---|
| 1 | 123 | 984 | 816 |
| 2 | 222 | 1,776 | 816 |
| 3 | 321 | 2,568 | 816 |
| 4 | 420 | 3,360 | 816 |
| 5 | 519 | 4,152 | 816 |
| 6 | 618 | 4,944 | 816 |
| 7 | 717 | 5,736 | 816 |
| 8 | 816 | 6,528 | 816 |
Each step of k adds exactly 99 active parameters and 792 operations, which is why no row is twice the row above it.
7 · the same counters at 64 experts
The same counters at any size
computed in this browserfour integersActive for one input: 2 × 2,048 + 512 = 4,608. Dense, if every expert ran: 64 × 2,048 + 512 = 131,584. The ratio is 0.035, so the dense layer runs 28.6 times as many parameters for the same input.
Hold E where it is and raise n: the ratio falls towards 0.031, which is the number usually quoted. The router runs for every input, so it never disappears from the bill: the true ratio stays a little above the ideal.
8 · the hand computation
A layer has 64 experts and routes the top 2. Each expert has E = 2,048 parameters and the router has R = 512. How many parameters run for one input?
A whole number, exactly.
Attempts: 0
9 · predict what raising k does
Before you take k to 4
Three counters, three questions. Recording them settles all three at once, because nothing here has to train.
All three are needed before they are recorded.
Record the three predictions first.
Ask for a hint
Rungs open as you run the session, not as time passes. You have run it 0 times.
1. A question: available after 3 evaluated runs.
2. A narrowing: available after 5 evaluated runs.
3. A worked fragment, on a smaller layer: available after 7 evaluated runs.
4. The answer: available after 9 evaluated runs.
04 · Compare
05 · Check your understanding
Your check: an outcome you produce
not yetIn one submission: the hand computation returns exactly 4,608 without the third-attempt assist, and all three predictions about raising k from 2 to 4 are correct. A value accepted after the third attempt marks the session done with the answer shown and does not satisfy the gate on its own.
Not met yet. Keep working the session.
Hand computation: not yet returned unassisted. Predictions: 0 of 3 correct , 3 not answered yet.
Save your takeaway
Your takeaway
Edit it into your own words, or keep the line as it is.
The artifact carries the recording's own hash, so the shelf entry names the run it came from. 10.2 studies the same layer at the same shape, from the same seeded initialisation, with one term added to the loss.