Learn machine learning · part 10 · session 2
Load balancing
A router with nothing pushing back on it sends almost every input to the few experts that happened to start well. One extra term in the loss changes that, and it is not free.
After this session you will be able to say why a mixture-of-experts layer needs a second loss term, and what happens on both sides of it.
The gate: select a recorded cell whose eight experts all sit between 6 and 20 percent of the training inputs, and whose held-out accuracy is within 0.02 of the best cell in the grid.
01 · The question
A router with nothing pushing back on it has an easy way to lower the loss early: train whichever happened to start well, and keep sending inputs to those. An expert that receives no inputs receives no gradient either, so it stays close to where it was initialised. You are then paying to store eight experts and getting the work of fewer, which is the gap between what a layer stores and its .
The usual fix is a second term in the loss. It is smallest when the eight experts get roughly equal shares, and largest when one takes everything. Alpha sets how hard it pushes. The course script ran this layer once at each weight on the picker, from the same seeded initialisation as 10.1, and this page plays those runs back. Say what you expect at weight 0 first, then play it. Then find a weight that spreads the load without giving up held-out accuracy.
What was recorded
recorded runrecord-runs.mjs · seed 7Six runs of 1,500 steps, one per balancing weight in 0, 0.01, 0.05, 0.1, 1, 10, each carrying the task loss and the balancing loss at every step, fifteen checkpoints of the eight routing shares, and its own trained parameters as int8. The script is in the repository at scripts/ml/grids/load-balancing.mjs.
The shares are read on the 600 training points, because 200 points cannot resolve an eight-way split to a percentage point. The accuracy and the loss are read on the 200 points no run trains on.
The balancing loss, written out
For a batch of T inputs and n experts, let f_i be the fraction of the batch whose top-1 expert is i, and let p_i be the mean over the batch of the router's softmax probability for expert i. Then L_balance = n × Σ_i ( f_i × p_i ).
The vector f records where the inputs went and p records how strongly the router scored them there, so the product charges an expert that is both crowded and confidently chosen. Both vectors sum to 1, so the value is 1 when both are uniform and 8 when all the mass sits on one expert. The trained objective is L_total = L_task + α × L_balance, and α is the control. This is one common form, the one used here, not the only one.
02 · Predict
Before anything runs
At balancing weight 0, how many of the eight experts will still take more than 5 percent of the 600 training inputs after 1,500 steps?
The weight-0 cell does not play until this is recorded. Once it is, play it and watch the inputs pile onto a few experts, one at a time. Calling the collapse and then watching it is the point.
03 · Try it
1 to 6 · choose a weight, play it, and read both sides
Loading the recorded grid.
Play is off until the prediction above is recorded.
What the six recorded cells did
Shares on the 600 training points, accuracy and loss on the 200 held out. The band runs 0.06 to 0.20, and the accuracy floor is 0.02 below the best cell in the grid.
| weight | lowest share | highest share | spread, of a possible 3.000 | held-out accuracy | held-out loss | against the gate |
|---|
That is a statement about this layer at this size, and it is why the gate reads the band first. The bill for over-balancing arrives at scale, where a fixed per-expert buffer means a router pushed towards uniform sends inputs to experts that have no room for them and they are dropped. Deep dive D12 is that experiment, on the cell you pass with here.
7 · run my own balancing weight
A weight between the recorded six
trained in this browseroptional, off the gateThe recorded picker offers six cells. This path trains the same layer here, at any weight from 0 to 20, in a worker so the page stays usable while it runs. 1,500 steps is about 20 seconds here on a mid-range phone, and the panel states the wait again before anything starts. It is not the default and the gate reads the recorded grid only.
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: reach the stated target
not yetSelect a recorded cell whose run holds every one of the eight experts between 0.06 and 0.20 of the training inputs, with held-out accuracy no more than 0.02 below the best cell in the grid, read off the grid itself. The prediction in beat two has to be recorded before the weight-0 cell plays.
Not met yet. Keep working the session.
Save your takeaway
Your takeaway
Edit it into your own words, or keep the line as it is.
The artifact records the weight, the eight shares, the held-out numbers, the trained layer as int8, and that all of it came from the recorded grid rather than from a run in this browser.