An optimizer is four dials: a learning rate, a batch size, a preconditioner and momentum. Two simple models predict most of how they interact. The noisy quadratic treats the loss as a bowl with steep and flat directions and the gradient as a noisy measurement of its slope. It explains why the best learning rate grows with batch size and then stops growing, why larger batches eventually stop paying, and why momentum and preconditioning help mostly at large batch. The rotational picture of the weights starts from the fact that most weights in a normalized transformer only matter through their direction. Gradient steps then make their norms grow, weight decay pulls them back, and the two settle into an equilibrium where the learning rate and the weight decay act only through their product. That product explains measured results from CS 312 that look strange otherwise, such as why the right response to an eight times larger batch is eight times more weight decay. Hyperball fixes the norm by construction, so that the learning rate directly sets how far each weight turns. Twelve sealed predictions, most checked against CS 312’s measured d8 runs, and three audit drills.
Plain stochastic gradient descent takes one step downhill on a minibatch: θ ← θ − η ∇ℓ(θ). Every optimizer in use is that update with four things changed:
AdamW, the d8 default, uses all four. Elementwise, for a weight matrix W with gradient g at step t:
The “W” stands for decoupled weight decay (Loshchilov and Hutter). Adding λ |W|²/2 to the loss instead would send the decay through the preconditioner, dividing it by √v and shrinking each coordinate by a different amount. AdamW shrinks every weight by the same factor 1 − ηλ per step. Hold on to that factor: it is the subject of sections 6 to 8. The d8 recipe uses β = (0.9, 0.95), ε = 10−8, peak η = 0.003, λ = 0.1, and exempts embeddings, normalization gains and biases from decay.
Classical optimization theory asks how fast a method reaches a precise stationary point, on every function in some class. Training a language model asks for neither. Nobody runs a model to convergence; the goal is a good loose solution within a fixed budget, on one narrow family of functions. Adam is the standard example. Its original convergence proof contained an error, and Reddi and colleagues later built simple problems on which it does not converge. Practice did not change, because the guarantee was never what anyone relied on. CS 312’s stance follows from this. Convex intuition is a source of hypotheses, not a verdict. Build top-down models of the large-scale behaviour (how a recipe responds to learning rate, batch, scale), then reconcile them with bottom-up mechanisms. The rest of this unit does exactly that, twice.
SGD’s gradient, and therefore its step, gets 10 times larger: the same η is now a 10 times larger learning rate, enough to diverge. Adam divides by the running √v, which also grows 10 times, so its step barely changes. The only exception is ε, which becomes relatively smaller. Adam is nearly invariant to the scale of the gradient; its step size per coordinate is set by η. Keep this in mind for section 6: a step whose size ignores the gradient’s scale is what makes weight norms matter.
Zhang and colleagues (2019) asked how far one can get by pretending the loss is a quadratic bowl and the gradient is its slope plus Gaussian noise. Further than you would expect: their noisy quadratic model (NQM) predicted how the best learning rate, the effect of momentum and the benefit of preconditioning change with batch size in real networks. Its appeal is that everything about it can be solved exactly.
Rotate coordinates so that the curvature matrix H is diagonal, with eigenvalues hi. Then the loss is ½ Σ hi θi2 and the minibatch gradient is Hθ + ε, with noise ε ~ N(0, C/B): averaging B independent examples divides the variance by B. If C is diagonal in the same coordinates, every coordinate evolves on its own. With SGD at step size α:
The expected loss has two parts derived. The first is what is left of the start, shrinking by a factor (1 − αh)² per step. The second is a noise floor the run settles on, proportional to α/B. Three consequences follow, and between them they hold most of this unit:
Square the update and take expectations; the noise is independent of θ, so the cross term vanishes: v′ = (1 − αh)² v + α²c/B, where v = E[θ²]. At the fixed point v = α²c / (B(1 − (1 − αh)²)) = αc / (B h (2 − αh)). The loss is ½hv = αc / (2B(2 − αh)). It depends on α and B only through α/B as long as αh is small, which is the seed of the linear scaling rule in section 3.
CS 312’s assignment 2 uses the smallest interesting case: two coordinates with H = diag(1, 10), noise covariance I (the same noise in both directions), a fixed budget of N = 8,192 examples (so N/B steps), and a random start, w1 ~ N(0, 1) and w2 ~ N(0, 0.1). The widget below is that toy. For SGD, momentum and the preconditioned variant the expected loss is computed exactly, by carrying the second moments through every step. Adam has no closed form, so it is simulated. Every view also shows the average of 64 seeded simulated runs, so you can see the exact curve and the simulation agree.
H = diag(1, 10), gradient noise N(0, I/B), 8,192 examples, constant step size. Exact where a closed form exists; 64 seeded runs otherwise.
This widget needs JavaScript.
Things to try, in this order:
Hold the data budget fixed at 8,192 examples and ask how the best SGD step size changes with the batch. Predict before switching the widget to its sweep view.
In the toy above, with plain SGD and the step size tuned separately at each batch size, the batch goes from 16 to 32 (so from 512 steps to 256). By what factor does the best step size change?
About ×2 (exactly 0.00944 → 0.0186, a factor of 1.97) derived. While αh is small, (1 − αh)2T ≈ e−2αhT with T = N/B, so the leftover depends on α/B, and so does the noise floor αc/(4B). Double B and double α and the expected loss after the same number of examples is unchanged. This is the linear scaling rule, which Goyal and colleagues used to train ImageNet in an hour.
Computed exactly from the closed form in this page; you can check it in the sweep view.
The argument has an expiry date. It needs αhmax to be small, and the limit α < 2/hmax does not move with the batch. Once the best step size gets near that limit, a larger batch still lowers the noise but can no longer buy a proportionally larger step, and the flat direction gets fewer steps to finish. Past that point a larger batch just wastes data. McCandlish and colleagues called the turning point the critical batch size; Zhang and colleagues showed the NQM reproduces it, and that it depends on the optimizer.
Fit the linear rule η* = c·B to the toy’s tuned SGD step sizes at B = 1 to 64. They fit it almost perfectly. Now use the rule at B = 512 (16 steps). What happens at the extrapolated step size?
It diverges derived. The rule gives η ≈ 0.000589 × 512 ≈ 0.30, above the stability limit 2/hmax = 0.2. The tuned step at B = 512 is 0.174, and even then the best loss is 5.7 times the B = 1 value. A rule fitted where its assumption holds says nothing about where the assumption breaks, and nothing in the fit warns you. Problem 3.1 of the assignment makes you find this range yourself.
Computed exactly from the closed form in this page.
Switch the widget to tune η at every batch size. The left panel shows the best step size against B on log axes: slope 1 for SGD until it bends toward the stability limit. The right panel shows the best loss after tuning: flat while the linear rule holds, then rising. SGD’s best loss is up 9% at B = 64, 22% at B = 128 and 5.7 times at B = 512 derived. In this toy the critical batch is around 64 to 128 examples.
For adaptive methods the analogous small-step argument suggests η* ∝ √B instead of ∝ B (the square-root rule of Malladi and colleagues). Real language models follow neither rule cleanly, and published scaling rules for the optimal learning rate disagree with each other; the assignment opens with a figure collecting them, from the Step Law paper. Assignment 2 has you measure the exponent for Adam in the toy, then in d8 at batch sizes 8 to 256. Section 7 shows that for AdamW the cleanest thing to scale with batch is not the learning rate alone.
Momentum. With s ← μs + g and θ ← θ − ηs, a run that takes small steps behaves like SGD with step size η/(1 − μ): the sum of agreeing gradients is just a larger step. Once the step size is tuned, that buys nothing. The difference shows up when the stability limit binds. Heavy-ball momentum stays stable up to ηh < 2(1 + μ), so its effective step η/(1 − μ) can reach 2(1 + μ)/((1 − μ)hmax): 38/hmax at μ = 0.9, against SGD’s 2/hmax derived. The flat direction gets large steps while the steep one oscillates but does not blow up. Small batches need small steps anyway, because the noise floor is high, so the extra headroom is useless there. Large batches have a low floor and can use it.
Preconditioning. Zhang and colleagues model it as stepping with H−p: each direction gets step ηh−p. At p = ½ the toy’s curvature spread shrinks from 10 to √10, so the steep direction stops holding the flat one back. Adam’s division by √v is a rough, noisy version of this, and in their experiments Adam and K-FAC raised the critical batch well above SGD’s.
The sweep view puts numbers on both derived. Tuned momentum (μ and η both tuned) beats tuned SGD by 1.02× at B = 1, 1.7× at B = 16 and 7.3× at B = 512. Preconditioning beats it by 1.46× at B = 1 and 7.7× at B = 512. In the toy preconditioning helps even at batch 1, which differs from the paper. The reason is that the assignment’s noise is the same in both directions, so the preconditioner also cuts the noise injected along the steep one; Zhang and colleagues assumed the noise covariance equals the curvature, and under that assumption the small-batch gain disappears. One more trap: with μ fixed at 0.9 instead of tuned, momentum is worse than SGD at B ≥ 128 in the toy. Sixteen to sixty-four steps are too few for a memory that spans about ten of them. “Momentum helps at large batch” assumes enough steps for it to act.
CS 312 ran the d8 version (quiz 2, building on assignment problem 3.2). At 614.4M tokens and two batch sizes, each at its own tuned learning rate and weight decay, only β1 of AdamW changed: 0, 0.9 or 0.98, with β2 = 0.95. At B = 8 (η = 0.00075, λ = 0.1) the default β1 = 0.9 reached 2.9366; at B = 256 (η = 0.0015, λ = 1.944) it reached 2.9682 measured.
At B = 8, rank the three runs (β1 = 0, 0.9, 0.98) by final validation loss, best first.
0.98 < 0.9 < 0 measured. At B = 8 the momentum curve is nearly flat, so the order between 0.9 and 0.98 is the hard part; removing momentum entirely still costs a little. The NQM says the small-batch effect of momentum should be small, not that more momentum should be worse.
CS 312 quiz 2, question 4, answer slides (Fall 2026).
At B = 256, the same ranking.
0.9 < 0.98 < 0 measured. At large batch the curve over β1 is sharp: no momentum is much worse, and too much is also worse. Fewer steps leave less room for a long memory, as in the toy.
CS 312 quiz 2, question 4, answer slides (Fall 2026).
How much does removing momentum (β1 = 0 instead of 0.9) cost in final validation loss, at each batch size?
0.0276 at B = 8 and 0.2503 at B = 256 measured: about nine times larger at the larger batch, the direction the NQM predicts. Note the design: each batch kept its own tuned η and λ while β1 changed, so the β1 = 0 runs were not retuned. Removing momentum changes the effective step size, so part of the 0.25 may be a learning-rate mismatch rather than the absence of momentum as such. The assignment asks you to compare this with the toy, where η was retuned for every μ.
CS 312 quiz 2, question 4 (b); differences computed from the course’s measured losses.
The NQM is the workhorse of this unit because it is exactly solvable and often points the right way. It is also wrong about several things that matter, and knowing where is part of using it:
The working rule: trust the NQM for the direction of an effect and for which regime you are in, and measure the size.
CS 312’s first quiz contained a question that looks like bookkeeping. Four d8 runs cross the learning-rate schedule (linear decay or constant) with weight decay (0.1 or 0). Every step, each run logs the L2 norm of all its parameters and the L2 norm of its gradient before clipping. Predict before reading on.
At the end of training, which runs have a parameter norm noticeably (3–15%) below the highest value it reached during the run? The others end within 3% of their peak.
Only linear decay with weight decay 0.1 measured; the other three end within 3% of their peak. Without weight decay nothing pulls the norm down, so it ends at its maximum. With weight decay and a constant learning rate it reaches an equilibrium and stays there. Only when the learning rate falls does the equilibrium itself move down, and the norm follows it. The rest of this section and the next explain why.
CS 312 quiz 1, question 6 (a), answer slides (Fall 2026).
In which runs does the gradient norm before clipping at least double over training (median of the last 10% of steps against the first 10%)?
None of them measured. Defazio (2025) reported gradient norms that rise sharply near the end of training with weight decay and a decaying learning rate, and proposed a fix (AdamC). The d8 runs did not show it at this scale and length. The answer slides note that the gradient norm seemed driven more by weight decay than by the schedule.
CS 312 quiz 1, question 6 (b), answer slides (Fall 2026).
Most weight matrices in a transformer with RMSNorm feed, directly or a few operations later, into a normalization that divides out their scale. QK-norm normalizes queries and keys; the pre-norm before each block normalizes the residual stream. For such a weight, multiplying it by any c > 0 leaves the network’s output unchanged, or nearly: L(cw) = L(w). Two facts follow by differentiating that identity derived:
Now take a step u orthogonal to w. By Pythagoras, |w + u|² = |w|² + |u|²: every step makes the norm grow, whichever way it points. A larger norm means the same step turns the weight by a smaller angle, and the angle is all the loss cares about. For SGD the effect compounds, because the gradient itself shrinks like 1/|w|: the effective learning rate is η/|w|² (van Laarhoven, 2017). Without weight decay the effective learning rate decays on its own, on a schedule nobody chose and nobody logs. Weight decay multiplies the norm by 1 − ηλ every step, pulling it back. The question is where the two balance.
Four more predictions first, because the next derivation gives them away. All four are from CS 312’s second quiz, on d8.
The first two use a 3 × 3 grid of peak learning rates {0.0015, 0.003, 0.006} and weight decays {0.1, 0.2, 0.4}, at 1.2288B tokens with a cosine schedule. The best cell was η = 0.003, λ = 0.2, with loss 2.8354 measured.
Three other cells of that grid: A (η = 0.006, λ = 0.1), C (η = 0.006, λ = 0.4) and D (η = 0.0015, λ = 0.1). Rank them by loss, best first.
A, then D, then C measured (the full measured order of the six off-optimum cells was A < B < F < E < D < C). Look at the products ηλ: A has 6×10−4, the same as the best cell, and is the closest to it, although its learning rate and weight decay are each off by a factor of two. D’s product is four times too small and C’s four times too large. Both are off by the same factor, and C is worse; with the learning rate also doubled, C is the one pushing toward instability.
CS 312 quiz 2, question 2 (a), answer slides (Fall 2026); our three-cell version of the question.
How far above the best cell (2.8354) is C?
About 0.032 measured. CS 312’s rule of thumb from the assignment’s fits is a quadratic penalty in how far the product is off: about 0.007 × (log2 of the factor)². Off by 4 = 2², that is 0.007 × 4 ≈ 0.028 derived.
CS 312 quiz 2, question 2 (b), answer slides (Fall 2026).
The next two start from the default d8 at 614.4M tokens, where the best pair at B = 64 is η = 0.003, λ = 0.2 (loss 2.9178) measured. The batch becomes 512: the same tokens in eight times fewer steps. Five recipes: A keeps both; B multiplies η by √8; C multiplies η by 8; D multiplies λ by 8; E multiplies both by √8.
At B = 512, which recipe gives the lowest loss?
D, eight times the weight decay measured. C, D and E all multiply the product ηλ by 8; D does it without touching the learning rate, which the noisy quadratic says cannot grow without limit. The answer slides summarize the assignment’s batch experiments this way: scaling weight decay with batch is predictable, scaling the learning rate tapers off.
CS 312 quiz 2, question 7 (a), answer slides (Fall 2026).
How much worse is recipe A (keeping the batch-64 pair) than the best recipe?
About 0.106 measured. The answer slides’ estimate: each doubling the product is off costs about 0.01 × (doublings)², and 8 is three doublings, so 0.01 × 9 ≈ 0.09 derived.
CS 312 quiz 2, question 7 (b), answer slides (Fall 2026).
Here is why the product is the knob. Take a scale-invariant weight w with d entries, trained with AdamW. Assume:
Each step shrinks w by 1 − ηλ and adds an orthogonal |u|:
At equilibrium the angle each weight turns per step is about √(2ηλ) derived. It depends on η and λ only through their product, and not on ρ or d. Since the loss only sees the direction, the product is the effective learning rate; η alone sets only the equilibrium norm. Kosson and colleagues (2024) call this rotational equilibrium; the same algebra with SGD (gradient ∝ 1/|w|) gives the same √(2ηλ). The approach to equilibrium takes about 1/(2ηλ) steps. For d8’s defaults, 2ηλ = 6×10−4, so that is about 1,700 of its 9,375 steps, and the equilibrium angle is about 0.024 radians, 1.4° per step derived. Real updates are only approximately orthogonal and uncorrelated, so read these as orders of magnitude.
A weight w in two dimensions whose loss depends only on its angle. Each step has fixed size η, is at right angles to w, and points toward the best direction when a noisy sign says so (an Adam-like normalized step). Exact simulation, 3,000 steps, seeded.
This widget needs JavaScript.
In the toy the step size is exactly constant, so the equilibrium is exact: |w|² = η/(λ(2 − ηλ)), and the angle per step is a function of ηλ alone derived. Things to try:
The grid. If only ηλ mattered, the loss surface over (log η, log λ) would be a valley along lines of constant product. The measured grid is close to that: the cells that keep the best product (A and B) are the best of the rest, and penalties grow with how far the product is off. Not exactly, because η also matters on its own: it sets the step size of the parameters that are not scale-invariant (embeddings, gains, the readout), and the stability limit from section 2 still applies. That is why C, with η doubled as well, is worse than D.
The batch. At a fixed token budget, B = 512 means eight times fewer steps. If the right amount of rotation over the whole run is what matters, each step must rotate more, and the timescale 1/(ηλ) must shrink eight times when measured in steps. Wang and Aitchison read AdamW as an exponential moving average of updates with that timescale. Bergsma and colleagues (Power Lines) found that the best timescale, expressed as a fraction of the training data, B/(ηλD), stays roughly constant, which makes the right ηλ proportional to B. Recipes C, D and E all satisfy that. D wins because it gets the whole factor from λ, which has no stability limit, while C puts it all into η, which does. Quiz 2’s B = 256 setting in section 4 (η = 0.0015, λ = 1.944) is this rule applied.
Double λ. After the switch there are half as many steps per token, so each step must rotate √2 times more: √(2ηλ) × √2 means ηλ × 2. Doubling η also doubles the product, but it doubles the step size of the parameters that are not scale-invariant too, and moves you toward the stability limit. Doubling λ changes only the equilibrium. Assignment 2 has a worked example question on exactly this switch, with five policies to rank; try it before reading their answer.
If the learning rate and weight decay matter mostly through the norm equilibrium they produce, why not set the norm directly? Hyperball (Wen and colleagues, 2026), used in assignment 2, does this for every linear-layer weight matrix. Take Adam’s update direction U, scale it relative to the weight, step, then project back to the weight’s initial Frobenius norm R = |W0|F:
No weight decay is needed, because the norm cannot drift. The relative step is η on every step, so the angle turned per step is about arctan η, whatever happened earlier in the run. The learning rate becomes a number with a meaning: the fraction of the weight’s norm to move per step. Switch the widget above to Hyperball to see it: the radius never changes and the angle per step stays fixed at arctan η.
Two subtleties the lecture flags. Not every weight should be norm-constrained: the readout’s scale encodes the model’s confidence, and embeddings may use their norms, so the course applies Hyperball to the internal matrices and the readout, while embeddings and normalization gains use Adam at 0.1041η. And the idea has relatives: weight normalization, steps on the sphere or on the Stiefel manifold, column-wise norms, ball projection. Wen and colleagues report lower loss and norm control by construction; the open question in the assignment is whether it also makes the best learning rate scale more predictably.
CS 312 fitted Hyperball’s optimal peak learning rate on d8 (weight decay off) at four token budgets. Read off their plot, the fitted optima are roughly 0.014, 0.013, 0.012 and 0.010 at 0.15B, 0.31B, 0.61B and 1.23B tokens measured (approximate, read from a figure). Which interval contains the fitted optimum at 2.46B tokens?
About 0.008 measured. The trend continued. A straight line in log–log through the four points gives about 0.009 derived; the last two alone give about 0.008. The answer slides call it an easy question for that reason: follow the line. Whether AdamW’s ladder in assignment problem 1 is as regular is exactly what the assignment asks you to find out.
CS 312 quiz 2, question 3, answer slides (Fall 2026).
Three fictional research logs, each written like an AI research agent’s notes. Each one breaks a rule from this unit while sounding reasonable. Tick every objection you would raise, then check. The director’s checklist from unit 0 applies throughout.
Which objections would you raise?
Two of the checklist’s questions fail here: was each arm tuned (no), and were the runs compared after the decay (no, and nobody noticed because the config hid it). The repair: express the schedule in tokens, rerun B=256 at λ × 2 and at η × √2 (two runs), and compare after the decay. The self-review checked whether the gap was larger than the noise, which it was, and stopped there.
Which objections would you raise?
Checklist question 1 (was each arm tuned) fails in the most common way: the baseline got a 3 × 3 grid and the challenger got the baseline’s winner. Question 5 would have caught it too: one Lion run at lr 3e-4 would likely have changed the picture. The repair is a Lion grid around lr 1e-4 to 3e-4 with wd 0.3 to 1.0, the same size as AdamW’s.
Which objections would you raise?
The worst move in this log is not the bundling. It is declaring the result noise without looking at the noise, then making the regression the new reference point. The repair: rerun the baseline from the same commit; run the two changes separately at seeds 0 and 1; adopt only a version that is no worse than the old baseline within its seed spread.
Nothing on this page was run on a GPU for this companion. The measured numbers are CS 312’s; the toy numbers are computed by this page’s widgets. The commands below are the ones the assignment handout gives, on the a2_release branch, with the Modal setup from unit 0.
--links:
uv run python -m experiments.a2.provided_sweeps --part P1a --linksP1a with P1b, P1d, P1e or P2a. Fit the quadratic in log learning rate yourself, and make the extrapolation in 1 (b) before you look at the larger budgets.assets/u2a.js and loads in node: require('./u2a.js').sgdFinal(eta, B, 0, 8192 / B). Adam and RMSProp are where your simulation does the work.launch_training_jobs. At roughly 15 minutes per run on an H100 for B = 64, that is at least 12 GPU-hours, more because small batches run slower. Budget it before launching, and run the unit 0 harness check first.The steepest direction limits the step size: α < 2/hmax, or it diverges. The flattest direction limits the duration: it needs about 1/(αhmin) steps. The gap between the two is why preconditioning, which narrows the curvature spread, helps.
For small steps, both the leftover from the start and the noise floor depend on α/B at a fixed number of examples, so doubling both leaves the outcome unchanged. It stops when α approaches 2/hmax: the steep direction caps the step, and beyond that larger batches only reduce the number of steps. That turning point is the critical batch.
At small steps momentum is SGD with step η/(1 − μ), which tuning matches. Small batches need small steps because of the noise floor. Large batches can use large steps, and momentum stays stable up to ηh < 2(1 + μ), so its effective step can exceed SGD’s 2/hmax. Measured on d8: removing momentum cost about 0.03 at B = 8 and 0.25 at B = 256.
Scale invariance makes the gradient orthogonal to the weight, so every step adds its squared length to the squared norm (Pythagoras). Nothing removes it unless weight decay does.
The angle is about √(2ηλ): the product. η alone sets the equilibrium norm, |w|² ≈ ηρ²d/(2λ), and the step size of the parameters that are not scale-invariant.
Multiplying weight decay by 8, which multiplies ηλ by 8 (the Power Lines rule ηλ ∝ B) without pushing the learning rate toward instability. Keeping the batch-64 pair cost about 0.106.
It scales each matrix’s Adam step relative to the weight’s norm and projects back to the initial norm, with no weight decay. The norm cannot drift, so η is the relative step per update, about the angle turned per step.