Unit 1 showed that the best learning rate moves when the model grows, so a ladder of sizes trained at one rate measures mis-tuning as much as size. This unit asks why it moves, and whether it has to. The answer is a piece of arithmetic: with Adam, every weight moves by about the learning rate, and a matrix with fan-in n sums n of those moves; when the moves line up with the input, the output changes n times as much. Undo that factor, layer by layer, and you get µP, a way of writing the model so the best learning rate stays put as width grows. You will run a real coordinate check in the browser, predict four results CS 312 measured on its own model (and see how far the theory carries), extend the idea to depth, and meet Muon, an optimizer that takes the same arithmetic one step further by controlling the size of every update direction at once. Three fictional audit drills close the unit: each is a confident conclusion built on a parametrization that was not what the log said it was.
Start with a prediction, before any argument. Lock it, then read on.
A ReLU network in the standard parametrization (Kaiming initialization, one learning rate for every matrix) takes one Adam step at a fixed learning rate. Look at the second hidden layer, an n×n matrix. As width n goes from 64 to 2048, how does the RMS change of that layer’s output features grow?
Roughly like n. The argument is below. The widget in section 4 runs this exact experiment on a toy network: the fitted log-log slope over widths 64–2048 is 0.88, and the last doubling (1024 → 2048) multiplies the change by 1.87. The slope is a little under 1 because a second, width-independent term (the change of the layer’s input) still matters at small widths.
derived from the alignment argument · measured in the toy MLP of section 4, computed by this page
Why n? Think about what Adam does on its first step. With bias correction, step one has m̂ = g and v̂ = g², so the update is g/|g|: every weight moves by exactly ±η. Later steps are not exactly signs, but stay of order η per entry. So under Adam the size of each weight change does not depend on width. What depends on width is how those changes add up.
Take one row of a hidden matrix and one example. The gradient of a linear layer is an outer product, δ x⊤: row i is δi times the input x. Its sign pattern is therefore sign(δi)·sign(xj), and the change that row makes to its output on that same input is
Every one of the n terms pushes the same way, because the update was built from the input it is now applied to. If the signs were random instead, the sum would grow like √n, the way a random walk does. These are the two reference cases CS 312 calls full alignment (n) and no alignment (√n). Batching many examples into one gradient blurs the outer product a little, but on the inputs the network actually trains on, the update stays strongly aligned.
So in the standard parametrization a hidden matrix’s effect on its output grows like η·n. To keep that effect the same size as you widen, η has to shrink like 1/n. CS 312 measured exactly this on its own model: in the course baseline, the best learning rate falls roughly as 1/n with width measured. This is the mechanism behind unit 1’s warning about ladders at one learning rate (unit 1): the larger rungs are not just bigger, they are hotter.
No. Each output coordinate of the input layer sums over the fan-in, which is the vocabulary or input dimension, and that does not grow with width. With Adam each weight still moves by about η, and each output sums a fixed number of those moves, so the change per coordinate is independent of n. Only matrices whose fan-in grows with width need the 1/n correction. Section 4’s widget shows it: layer 1 (fan-in 32) moves the same amount at every width under both parametrizations.
“Shrink the LR like 1/n” is a patch. To decide what the right patch is for every layer, lecture 4 states what a good infinite-width limit should look like. Take width to infinity at a fixed training step, with order-one inputs. Then we want:
There is a famous limit that is stable and nontrivial but fails the third condition: the neural tangent kernel (Jacot et al., 2018). Scale the output by 1/√n, keep weights and learning rate of order one, and each weight moves only O(n−1/2). The network stays a linear function of its parameters around initialization: training becomes kernel regression on fixed features, and the first layer barely moves. It is mathematically tidy and empirically much worse than real networks. Stability alone is not the goal; the point of µP is to find the limit where every layer keeps learning.
One warning that lecture 4 makes explicitly: all of this is leading-order arithmetic, “physics math”. It assumes the terms being added are independent enough not to cancel. Two quantities of size n can subtract to something of size log n. The rules below are a strong starting point, and section 5 shows how far they hold up on a real model.
Following Everett et al. (2024), write each layer with fan-in n using three exponents: a forward multiplier n−a, an initialization scale with variance σ²n−2b, and a learning rate ηn−c. The hyperparameters σ and η are the ones we hope to tune once.
Three facts fix the exponents, for Adam with full alignment:
Inputs with a fixed fan-in need a + b = a + c = 0. Choosing a = 0 where possible and writing m = n/n0 for the width multiplier over a reference width n0 (where everything matches the tuned baseline), assignment 2 summarizes one valid choice:
| input | hidden | readout | |
|---|---|---|---|
| (a, b, c) | (0, 0, 0) | (0, ½, 1) | (1, 0, 0) |
| forward multiplier | 1 | 1 | 1/m |
| initialization variance | 1/d | 1/fan-in | 1/n0 |
| Adam learning rate | η | η/m | η |
Read the bold cells as the three changes from the standard parametrization: hidden matrices get a learning rate that falls with width, and the readout gets a fixed initialization scale plus a 1/m multiplier in the forward pass. Everything else is Kaiming. At m = 1 nothing changes, which is the point: you tune at n0, then widen.
The direct update would be ηn1/2−a−c, so with a = 0 we need c = 1/2: a hidden LR of η/√m. That is assignment 2’s problem 4.0, and it is the reason the alignment exponent matters. Lecture 4 shows measured alignment during training between about 0.5 and 1.0, so the truth sits between the two prescriptions. Keep that in mind for section 5.
A coordinate check is the standard test of a parametrization (Yang et al., Tensor Programs V): train a few steps at several widths with the same base learning rate, and plot how much the features move against width. If the parametrization is right, the lines are flat. It costs seconds, and it catches most implementation mistakes before any real run.
Real computation, in this page: input 32 → ReLU hidden n → ReLU hidden n → 10 classes, one fixed batch of 16 random examples, Adam (β = 0.9, 0.999, ε = 10−8) at base LR 0.004, reference width 64. Both parametrizations start from the same weight draws; only the scales of the table above differ. This is a toy network, not CS 312’s transformer.
What it produced when we checked it in node with the same code (one step): over widths 64–1024 the fitted slopes are 0.86 (SP, layer 2) and −0.02 (µP, layer 2); logits 1.30 (SP) and −0.01 (µP). Adding width 2048: 0.88 and −0.01; logits 1.33 and 0.01. Layer 1 is flat under both (fan-in fixed). At 3 and 5 steps the SP slope for layer 2 falls (0.64 at 3 steps over 64–1024) because at large width the hot updates start to saturate; µP stays flat at every step count. measured on the toy, computed in your browser
Three things to notice. At width 64 the two lines start at the same point: m = 1, so the two parametrizations are the same network. Under SP the slope for layer 2 is a little under 1 at small widths because the layer’s output also changes through its input (layer 1 moved), a width-independent term; once the n-proportional term dominates, each doubling roughly doubles the change. And the logits under SP grow faster than n because two effects stack: the readout’s own update, and the readout amplifying the already inflated layer-2 change.
A coordinate check is the cheapest evidence you can ask for that a parametrization was implemented as described. If a log says “we switched to µP” and there is no flat coordinate-check plot, you do not yet know that the code does what the log says. Ask for the plot before any transfer claim.
The payoff of µP is learning-rate transfer: sweep the base learning rate at a small width, then reuse the best value at a large width. Here is the cartoon of what that looks like.
illustrative A made-up model, not a measurement: loss = E + A·n−1/4 + 0.02·(log2(η/η*))², with η* = 0.003·512/n under SP and 0.003 under µP. It shows the shape of the claim, not its size.
Reality is messier, and CS 312’s quiz 2 asked about two of the messy cases. The setting: the course’s d8-family model at 153.6M tokens, reference width n0 = 512, m = n/512. The course baseline uses embedding init G/n, readout init G/√n and readout multiplier 1, where G is one fixed matrix of standard Gaussian draws; µP changes those to G/n0, G/√n0, 1/m, and divides the hidden-matrix LR by m. Predict before reading the answers.
Almost-µP takes µP’s embedding initialization and hidden-matrix LR (η/m), but keeps the baseline’s readout (initialization G/√n, multiplier 1). Trained at width 1024 (m = 2), with the width-512 optimum at η* = 0.002378. Which interval contains η*(1024)/η*(512)?
0.869. The baseline’s optimum falls roughly as 1/n, which would give about 0.5 for one doubling; full µP is roughly flat, about 1. Almost-µP lands between them and much closer to µP, because it keeps the rule that matters most: every hidden matrix gets η/m. The readout is one matrix among many, and its share of the effect is smaller. Using the width-512 learning rate (0.003) at width 1024 cost a measured 0.0175 in loss, in the interval [0.01, 0.06) the course expected; assignment 2 found the full-µP penalty near zero.
measured CS 312 quiz 2 Q5, answer key and answer slides (seed 42)
Now the baseline parametrization (not µP), trained with Hyperball instead of AdamW: every linear matrix, readout included, takes an Adam-direction step and is projected back to its initial Frobenius norm (unit 2a); embeddings and gains use Adam; no weight decay. Widths 128, 256, 512, 1024. Fit η*(n) = c·np. Which interval contains p?
−0.561. The course’s reasoning: the AdamW baseline falls roughly as 1/n, µP is roughly flat, and measured update alignment is around 0.75 rather than the fully aligned 1, so a method that fixes some but not all of the scaling should land in between. Hyperball fixes each matrix’s norm, but does not divide the hidden learning rate by width, so it does not undo the alignment factor. Transferring width 512’s best learning rate to width 1024 cost a measured 0.033, between AdamW’s penalty and µP’s.
We have not derived −0.56 from first principles, and neither does the course: it is an empirical number, and the useful skill is placing it in the right interval with the bracket “between −1 and 0, alignment below 1”.
measured CS 312 quiz 2 Q6, answer key (fitted exponent for seed 42) and answer slides
What these two numbers teach is not “µP works” or “µP fails”, but how to reason when only part of a recipe is applied: find which matrices carry the width dependence, and ask which of them the change actually touched. The same question returns in drill 1.
Assignment 2 asks this repeatedly: does the prescription with better learning-rate transfer also reach a lower loss after tuning? Not necessarily. Transfer saves you the sweep at the large width; it does not promise that the transferred model beats a separately tuned standard-parametrization model. A director should keep these as two claims and ask for evidence of each.
Width is not the only axis. A transformer adds every block’s output into a residual stream, xl+1 = xl + cL·fl(xl), and with L blocks the same counting game applies. Independent random blocks at initialization add like a random walk, so the residual stream grows like √L. Updates during training are more coherent (all blocks are pushed by the same loss), so their effects add more like L. Assignment 2 compares three prescriptions, with reference depth 2 and r = L/2, applied on top of µP for width:
| prescription | residual multiplier cL | block LR | block Adam ε |
|---|---|---|---|
| standard µP | 1 | η | 10−8 |
| depth-µP | r−1/2 | η r−1/2 | 10−8 r−1/2 |
| CompleteP | r−1 | η | 10−8 r−1 |
Depth-µP (Yang et al., Tensor Programs VI; Bordelon et al., 2023) uses 1/√L, matched to the random-walk sum: the residual stream stays order one at initialization. CompleteP (Dey et al., 2025) uses 1/L, matched to the coherent sum, and keeps each block’s own learning rate at η so that every block still learns features of its own (the paper’s title is “Don’t be lazy”: under depth-µP’s shrinking block LR, deep blocks drift toward lazy, linearized behaviour). In assignment 2’s five-step stress test with independent blocks, the course’s own playtest found depth-µP did best and CompleteP over-corrected measured. Quiz 2 then changed one thing.
Depth L = 1000 (r = 500), width 64. The 1,000 blocks are now one shared block, initialized like one block of the stress test and applied 1,000 times. Let ρ = (RMS of the final residual stream at initialization, shared block) / (the same with 1,000 independent blocks), same prescription and seed. For which prescriptions is ρ below 10?
CompleteP only: measured ρ ≈ 32.2 (standard µP), 28.6 (depth-µP), 1.56 (CompleteP). A shared block adds the same contribution 1,000 times, coherently, instead of 1,000 independent ones that partly cancel, so the branch sum grows by a similar factor in every prescription, of the order of √1000 ≈ 32: the course’s table shows 33 → 1,062 (32×), 0.89 → 38.4 (43×), and 0.036 → 1.20 (33×). The multiplier scales the shared and the independent sums alike, so it largely cancels from the branch-sum ratio.
Then why is CompleteP’s ρ only 1.56? Because the residual stream is not only the branch sum: it also carries the embedding, with RMS about 1. Under CompleteP the branches add only 0.036 with independent blocks, so the stream’s RMS is essentially the embedding’s (1.00), and sharing raises it to 1.56. The trap is reading an RMS ratio as if it measured only the part you were thinking about.
measured CS 312 quiz 2 Q1(a), answer key and answer slides
Same shared-block model. For each prescription take the lowest validation loss after five updates over base learning rates {0.003, 0.01, 0.03}. Which prescription attains the lowest loss?
CompleteP. A shared block is the limiting case of coherent updates: every application receives the same update, so the 1,000 contributions are perfectly aligned. CompleteP’s 1/L multiplier is designed for exactly that, while depth-µP’s 1/√L is designed for incoherent sums. The course’s rationale is one line: CompleteP targets full alignment, and block copying is full alignment. The lesson for prediction: which prescription wins depends on which alignment regime the experiment is in, and changing the experiment can flip the winner.
measured CS 312 quiz 2 Q1(b), answer key
Why does the table scale Adam’s ε? Adam’s update is m̂/(√v̂ + ε). Everything above assumes that ratio is about ±1, which holds only while gradient entries are much larger than ε. But the prescriptions shrink gradients: a residual multiplier cL multiplies every gradient flowing into a block, and the readout’s 1/m multiplier scales the signal flowing back into the whole network. Once gradient entries approach 10−8, ε takes over, the update shrinks below η, and the model is silently running a different parametrization from the one the log names. Scaling ε with the multiplier, as assignment 2 does, keeps the ratio intact. Everett et al. (2024) report that ε matters at scale for exactly this reason and propose Adam-atan2, which removes it.
Where √v̂ ≈ 3×10−9, the update is about 3×10−9/(3×10−9 + 10−8) ≈ 0.23 of the intended ±1, so those matrices move at roughly a quarter of their nominal learning rate, and coordinates with smaller gradients even less. The learning rate in the config is no longer the learning rate the model sees. Either scale ε down, or treat any width or depth comparison from that run as confounded.
µP chooses initialization and learning rates so that an Adam update moves features by the right amount. Lecture 4 then asks the question the other way round: what if the update itself were built to move features by the right amount, in every direction at once?
Gradient descent picks the step that most decreases a linear model of the loss inside a small ball: minimize ⟨G, ΔW⟩ subject to ΔW being small. With the usual Euclidean ball you get −ηG. But the quantity µP cares about is how much ΔW can change the output, ‖ΔW x‖ ≤ ‖ΔW‖op‖x‖, which is controlled by the operator (spectral) norm, the largest singular value. Use that ball instead, and the answer is clean (Bernstein and Newhouse, 2024): if G = UΣV⊤, then
UV⊤ is the polar factor of G. Every direction the gradient points along gets the same step size, however small its singular value was. That is “idealized Muon”. Practical Muon (Keller Jordan, 2024) makes three changes: it orthogonalizes the momentum (Nesterov style) rather than the raw gradient; it replaces the SVD with five steps of a Newton–Schulz iteration, which uses only matrix multiplies; and it applies only to the 2D hidden matrices, leaving embeddings, the output head, and every scalar or vector parameter to AdamW.
Divide the matrix by its Frobenius norm, so all singular values are at most 1. Then repeat, with A = XX⊤:
Each step keeps U and V and applies the scalar polynomial p(σ) = aσ + bσ³ + cσ⁵ to every singular value. The coefficients are tuned for speed: a slope of 3.4445 at zero means small singular values grow fast. The price is that 1 is not a stable fixed point, so the values do not converge; they settle into a band.
A 64×128 matrix is divided by its Frobenius norm; its singular values then range from 0.013 to 0.85 (a ratio of 67). After Muon’s five Newton–Schulz steps, what do its singular values look like?
A band. Computed by the widget below: after five steps every singular value lies between 0.683 and 1.134 (ratio 1.66, down from 67). Applying p five times to any starting value between 0.002 and 1 lands in [0.68, 1.20]; iterate forever and the values keep cycling within [0.68, 1.13] instead of converging. Values below about 0.001 do not make it in five steps (0.001 only reaches 0.47). Jordan’s post says the coefficients were chosen to land within roughly 1 ± 0.3, and that this tolerance did not hurt the loss curve in his runs.
measured by this page’s computation (the “gradient-like” matrix) · tolerance claim from K. Jordan’s Muon post
Real computation, in this page: a seeded 64×128 matrix, divided by its Frobenius norm, then k iterations of Muon’s quintic. Singular values come from the eigenvalues of XX⊤ (Jacobi rotations).
Look at the gradient-like matrix at k = 0: four large singular values (the signal) and sixty small ones (noise). Muon raises the small directions to the same step size as the large ones. That is the bet Muon makes: that rare directions in the update are under-served by gradient descent and Adam. Whether that is exactly the right bet is still open; lecture 4 closes with evidence that the precise spectrum is not what matters, and that perfectly uniform singular values are not optimal either. Treat “Muon works because it orthogonalizes” as a hypothesis with good results behind it, not a settled mechanism.
Each drill is a fictional research log written by an agent, ending in a confident conclusion. Tick every problem that undermines the conclusion, and leave unticked anything that is true in general but not a problem here. The logs are invented for this course; the mistakes are the ones this unit is about.
Which of these undermine the conclusion in W-17?
Debrief. The log names a method (“µP”) and then runs something else: a half-converted readout, and an ε that quietly rescales the hidden learning rate at the large width. Add a data change between the two ends of the transfer, and no conclusion about µP is available. What to ask for: a coordinate check at 128 and 2048 (flat lines, seconds of compute), ε scaled or checked against gradient magnitudes, and a small learning-rate sweep at the target width on the same data, which is the only thing that shows where the optimum actually is.
Which of these undermine the conclusion in B-34?
Debrief. A decisive-looking gap, well outside the noise, can still answer the wrong question. Three configuration choices each push against Muon, and the comparison point is mid-schedule. What to ask for: Muon on hidden matrices only, RMS-matched update with weight decay, a small learning-rate sweep for Muon (it is a different optimizer and gets its own), and a comparison after the decay. A director who checked only “how many seeds?” would have passed this log.
Which of these undermine the conclusion in C-10?
Debrief. “X stops helping” is a claim about the task, but an unscaled ladder at one learning rate tests the parametrization. The log even contains the evidence (the growing residual RMS) and does not use it. What to ask for: a depth prescription (residual multiplier and its matching LR and ε rules), a short learning-rate sweep at the deepest rung, and only then a statement about depth.
This unit prepares CS 312 quiz 2’s Q1 (depth with a shared block), Q5 (almost-µP) and Q6 (Hyperball across widths); unit 2a covers Q2–Q4 and Q7. If you predicted above, you have seen the answers to Q1 and the first parts of Q5 and Q6. The course’s own interactive version grades locally in your browser: quiz-hyperparameter-scaling. Their answer slides make a point worth keeping: the quiz is designed so that remembering what your own experiments showed beats reasoning from scratch.
Nothing in this lab was run by us. It points at CS 312’s assignment 2, which is theirs: use their repository (the a2_release branch of deep-learning-alchemy/assignments, which has no licence, so clone it rather than copying code around) and their handout.
Adam moves each weight by about η regardless of width, and a gradient built from an input is aligned with that input, so a row’s n moves add up coherently: the output changes by about η·n. Keeping that change constant needs η ∝ 1/n. With unaligned updates it would be 1/√n.
Hidden matrices: learning rate η/m. Readout: initialization variance fixed at 1/n0 and a forward multiplier 1/m. Inputs and everything else: unchanged. At m = 1 the model is the tuned baseline.
The RMS change of each layer’s features (and of the logits) after a few steps at a fixed base learning rate, against width on log-log axes. Under a correct parametrization the lines are flat; under SP, layers with width-growing fan-in rise with slope near 1.
Each weight moves only O(n−1/2), so the network stays linear around its initialization and its features never change: kernel regression on fixed features. It performs much worse than real networks. µP is the limit where features keep changing by a constant amount.
Depth-µP: block contributions add incoherently (random walk, √L), as at initialization. CompleteP: they add coherently (L), as aligned updates do, while keeping each block’s LR at η so every block learns. A shared block is the fully coherent case, which is why CompleteP wins quiz 2 Q1.
The multipliers shrink gradients. Once gradient entries approach ε, Adam’s update falls below its intended size and the effective learning rate is no longer what the parametrization prescribes. Scaling ε with the multiplier keeps the update sign-like.
The polar factor UV⊤ of the momentum: the steepest-descent step under the operator norm, every singular value set to 1. Five quintic Newton–Schulz steps approximate it with matrix multiplies only; singular values land in a band (0.68–1.13 on this page’s matrices), not exactly at 1, by design.
Nesterov momentum; AdamW for embeddings, the output head and vector parameters; update RMS matching by 0.2·√max(A,B) plus weight decay (Moonlight, Kimi K2); QK-Clip on attention heads whose logits exceed a cap (Kimi K2).