Chapter 2 — Optimization: How Models Actually Learn
Chapter 1 solved least squares exactly (Ch 1, Section 3) and ran one gradient-descent loop on it (Ch 1, Section 4). Almost nothing after it has a closed form. Logistic regression, the lasso, every tree ensemble's leaf fit and every neural network must be searched for. This chapter is about that search. By the end you can compute a gradient and a step by hand, prove when a step size is safe and how many steps you'll need, explain what SGD, momentum, Adam, AdamW and Newton each fix, read a loss curve like a doctor reads an X-ray, and answer every optimization question an interviewer asks. One idea carries all of it: learning is walking downhill on the loss, and the curvature of the hill decides how big a step you can take and how fast you arrive.
How to read it. The chapter is a chain of 15 questions in three acts. Act I, Downhill (Sections 1–5), works on one bowl: which way is down, when the bottom is unique, how big a step, how many steps. Act II, Faster and noisier (Sections 6–9), buys cheaper gradients (SGD), memory (momentum), per-weight step sizes (Adam) and real curvature (Newton). Act III, In practice (Sections 10–15), covers what breaks, how it looks on a loss curve, landscapes without a single bowl, constraints, and how to build and ship it. Each section opens with its question and a 30-second version, and ends with the question it leaves open.
Two passes. On the first pass, read only the question, the 30-second version, and the figures of each section. You'll follow the whole arc in under an hour. On the second pass, read everything, doing the Remember box at the end of each section from memory before you look. Practice lives in its own tabs: Derivations (what to prove on a whiteboard) and Interview Bank (every question by level, with the intuition behind each answer).
1. Learning Is Optimization
The question: our six orders took 12 to 28 minutes. What does it mean to "train" an ETA model on them, and what are we actually searching for?
The 30-second version: training means choosing the weights that make the average loss on past orders as small as possible: empirical risk minimization. The loss over all possible weights is a surface; training is finding its lowest point. "Optimal" means the global minimum, but most methods can only promise a point where the slope is zero. And we never need the exact bottom: optimizing far past the statistical noise in the data buys nothing. First read: the problem, the surface picture, the vocabulary table.
The problem
You work on a food-delivery app. When a customer orders, the app promises an ETA, and the ops team wants it predicted from distance and number of items. Six recent orders:
| Order | Distance (m, as logged) | Distance (km) | Items | ETA (min) |
|---|---|---|---|---|
| A | 1,000 | 1 | 2 | 15 |
| B | 2,000 | 2 | 1 | 12 |
| C | 3,000 | 3 | 5 | 24 |
| D | 3,000 | 3 | 3 | 21 |
| E | 4,000 | 4 | 4 | 26 |
| F | 5,000 | 5 | 3 | 28 |
Two details matter all chapter. First, the dispatch system logs distance in meters; that column will wreck gradient descent in Section 4. Second, distance and items are correlated (corr = 0.5): bigger orders tend to travel further. That will make the bowl lopsided.
One convention, used everywhere. We center both features (distance − 3 km, items − 3). The intercept then separates out and equals the mean ETA, 21 minutes (the line passes through the means, Ch 1, Section 3), and the optimization problem becomes a 2-D bowl in (w_{\text{km}}, w_{\text{items}}) you can draw and check by hand. The loss is the mean squared error with the 1/n, as in Ch 1, Section 4:
- L
- the loss: mean squared error of the ETA predictions, in minutes², a function of the weights
- \mathbf w
- the weights (w_{\text{km}}, w_{\text{items}}): minutes per km and per item (bold = a vector)
- n
- the number of orders (6 in the toy)
- \mathbf y
- the six centered ETAs (minutes minus the 21-minute mean), stacked as a vector
- X
- the data matrix: one row per order, columns = centered km and centered items
- \lVert \mathbf y - X\mathbf w\rVert^2
- the sum of the six squared misses (squared length of the miss vector)
- \nabla L
- the gradient: the vector of slopes of the loss, one per weight; points uphill
- X^\top
- X transposed: rows become columns, so each feature column meets the misses
- H
- the Hessian: the matrix of second derivatives, the bowl's curvature in every direction
- X^\top X
- the features' sums of squares and cross-products; it alone sets the bowl's shape
The answer we're searching for is known exactly, which is what makes the toy useful: \text{ETA} = 6 + 3\cdot\text{km} + 2\cdot\text{items}, residuals (2, −2, −1, 0, 0, 1), and the best possible loss L^* = 10/6 \approx 1.667. Every method in this chapter is checked against (3, 2).
The 10,000-order simulation. Noise, batches, schedules and data ordering can't be shown on six rows, so we also simulate 10,000 orders (seed 2002) from the same truth plus noise of 3 minutes. Distance averages 3.18 km (sd 2.24), items 3.00 (sd 1.51), their correlation is 0.36, and ETA averages 21.6 minutes. OLS recovers (6.05, 2.99, 1.99). One more realistic detail: the warehouse table is stored sorted by distance, which matters in Section 10.
Every model is three choices: F-L-O
| Choice | Question it answers | ETA model | Where it's owned |
|---|---|---|---|
| Family | What shapes can the prediction take? | A weighted sum of distance and items | Ch 1, and later model chapters |
| Loss | What does "wrong" cost? | Squared error in minutes | Ch 1 (MSE), Ch 3 (cross-entropy) |
| Optimizer | How do we find the best member of the family? | Gradient descent and its relatives | This chapter |
Memory hook: every model is F-L-O; this chapter is the O.
Empirical risk minimization in one formula
What we care about is the expected loss on future orders, the risk. We can't compute it, so we minimize the average over the orders we have, the empirical risk:
- \hat{\mathbf w}
- the trained weights (hat = estimated from data)
- \arg\min_{\mathbf w}
- the weights \mathbf w that make what follows as small as possible
- \sum_{i=1}^{n}
- add up over every past order, i = 1 to n
- \ell
- the per-order loss: what one prediction's miss costs (here squared minutes)
- y_i
- order i's actual delivery time in minutes
- f(\mathbf x_i; \mathbf w)
- the model's predicted ETA for order i, from its features \mathbf x_i (distance, items)
Picture the loss as a landscape: the two weights are east–west and north–south, and the loss is altitude. Training is finding the lowest point of that landscape.
Vocabulary you'll be asked about
| Term | Meaning | On the ETA bowl |
|---|---|---|
| Global minimum | The lowest point anywhere | (3, 2), loss 1.667 |
| Local minimum | Lower than everything nearby | Same point: a bowl has only one |
| Stationary (critical) point | Where the gradient is zero | Same point |
| Saddle point | Stationary, but downhill in some direction | None here; common in neural nets (Section 11) |
| Plateau | A region where the gradient is nearly zero but the loss is high | Raw meters, Section 4 |
The trap in that table: "the optimizer converged" usually means it reached a stationary point, not the global minimum. Only convexity (Section 3) closes that gap.
Three errors, and why we stop early
Least squares has a closed form, X^\top Xw = X^\top y, costing about np^{2} to form X^\top X plus p^{3} to solve (Ch 1, Section 3); iteration is needed when p is huge, the data stream in, or no formula exists. Either way, Bottou and Bousquet (2008) split the gap between our trained model and the best possible predictor into three parts:
| Error | Caused by | Shrinks with |
|---|---|---|
| Approximation | The family can't represent the truth | A richer family (F) |
| Estimation | We fit on a finite sample, not the population | More data |
| Optimization | We stopped before the exact minimum | More steps, a better optimizer (O) |
On the 10k orders, the distance weight's standard error is 0.0144 min/km: another sample of orders would move it about that much. Pushing the optimization error to 10^{-10} when the estimation error is 10^{-2} buys nothing. Don't optimize past the noise. That's why large-scale training uses cheap, noisy steps (Section 6) and stops on validation loss rather than at machine precision.
Remember: Every model is F-L-O; this chapter is the O. Training = empirical risk minimization = finding the lowest point of the loss surface. "Converged" means "slope is zero", not "global best". Stop when optimization error is well below the statistical noise.
Next question → We know we're looking for the bottom of a surface we can't see. Standing at one point, how do we know which way is down? (Section 2)
2. Which Way Is Down? Gradients and the Taylor Picture
The question: standing at w = (0, 0) on the six-order bowl, which way should we step, and why does a small step against the gradient always lower the loss?
The 30-second version: the gradient is the vector of slopes, one per weight, and points in the steepest uphill direction. A first-order Taylor expansion says a small step of -\eta\nabla L lowers the loss by about \eta\lVert \nabla L\rVert ^{2}. The second-order term, set by the curvature (the Hessian), is what eventually stops a step from being "small enough". At (0, 0) the gradient is (−13.3, −11.7), so we raise both weights. First read: the hillside picture, the by-hand gradient, the Taylor box.
You need this: derivatives, partial derivatives, gradients
A derivative is a slope: how much f changes per tiny change in its input. With several inputs, a partial derivative \partial L/\partial w_{j} is the slope along one axis with the others held fixed. Stack all of them and you get the gradient, \nabla L, an arrow in weight space. Two facts about it carry the chapter: it points in the direction of steepest ascent, and its length is how steep that ascent is.
The gradient at w = 0, by hand
With centered data the intercept is already 21, so at w = (0, 0) every prediction misses by the centered ETA \tilde y = (-6, -9, 3, 0, 5, 7). The gradient is "each feature column, weighted by the misses" (Ch 1, Section 4):
- \nabla L
- the gradient: the vector of the loss's slopes, one per weight; points uphill
- \mathbf 0
- the starting point: both weights zero
- \sum_i
- add up over the six orders i
- d_{1i}
- order i's centered distance (km minus 3)
- d_{2i}
- order i's centered item count (items minus 3)
- \tilde y_i
- order i's centered ETA: minutes minus the 21-minute mean (tilde = centered)
- \lVert\nabla L\rVert
- the gradient's length: how steep the loss is at this point
Both slopes are negative: raising either weight lowers the loss. One step with \eta = 0.1 moves to w_{1} = (0, 0) - 0.1\cdot (-13.33, -11.67) = (1.333, 1.167), and the loss falls from 33.33 to 9.77.
Why -\nabla L is the steepest way down. Move a tiny distance in a unit direction u. The loss changes by the directional derivative \nabla L^\top u. By Cauchy–Schwarz, \nabla L^\top u \ge -\lVert \nabla L\rVert, with equality only when u points exactly against the gradient (Card 1). So among all directions, -\nabla L/\lVert \nabla L\rVert lowers the loss fastest. That's the whole justification for w -= lr * grad. Ch 1's hook still holds: \text{Step} = -\eta \times \text{slope}.
The Taylor picture
Near any point, a smooth loss looks like its Taylor expansion. Keep one term and you get a tilted plane; keep two and you get a bowl:
- L
- the loss (mean squared ETA error) as a function of the weights
- \mathbf w
- the current weights (w_{\text{km}}, w_{\text{items}}) (bold = a vector)
- \Delta
- a small step in weight space
- \nabla L^\top \Delta
- the slope term: gradient dotted with the step (^\top = transpose), the first-order change in loss
- \Delta^\top H\,\Delta
- the curvature term: how much the bowl bends along the step
- H
- the Hessian: the matrix of second derivatives, the bowl's curvature in every direction
- H_{jk}
- the Hessian's entry in row j, column k
- \frac{\partial^2 L}{\partial w_j\,\partial w_k}
- how fast the slope along w_j changes as you nudge w_k: a rate of a rate, not a fraction to divide
Put in the gradient step \Delta = -\eta\nabla L:
- L(\mathbf w - \eta\nabla L)
- the loss after one step of size \eta against the gradient
- \eta
- the learning rate: how far each step moves (0.1 in the example)
- \nabla L^\top H\,\nabla L
- the bowl's curvature along the gradient, times the gradient's squared length
The first-order gain grows like \eta; the curvature penalty grows like \eta^{2}. So small steps always win, and long steps eventually lose. On the toy at \eta = 0.1, the linear model predicts the loss will fall to 1.944. It actually falls to 9.769, which is exactly what the quadratic model predicts (our loss is a quadratic, so Taylor-2 is exact). The linear model over-promised because it ignores the curvature term, here \tfrac12 \cdot 0.01 \cdot 1{,}564.8.
The Hessian and the smoothness constant L. The Hessian H is the matrix of second derivatives: the curvature of the bowl in every direction. For MSE it's H = (2/n)X^\top X, which doesn't depend on w. On the toy:
The smoothness constant L (overloaded notation, unfortunately standard) is the largest curvature (in size) anywhere: the gradient can't change faster than L per unit of distance. For a quadratic it's the largest eigenvalue of H. On the toy, L = 5.
The descent lemma: a guaranteed drop
If the curvature never exceeds L, the Taylor bound becomes a guarantee (Card 2):
- L(\mathbf w)
- the loss at the current weights
- L
- here the smoothness constant: the bowl's largest curvature (5 on the toy), not the loss
Any \eta < 2/L lowers the loss; \eta = 1/L gives the largest guaranteed drop, \lVert \nabla L\rVert ^{2}/(2L). On the toy at \eta = 1/5 = 0.2 the lemma guarantees a drop of 31.39; the actual drop is 31.48. The bound is nearly tight because the gradient at w = 0 points almost along the steepest direction.
Zero slope is necessary, not sufficient. At a minimum \nabla L = 0, but also at maxima and saddles; the curvature decides which (Sections 3 and 11).
Memory hook: "Slope says which way, curvature says how far."
Remember: The gradient points uphill; step against it. First-order Taylor promises a drop of \eta\lVert \nabla L\rVert ^{2}; the curvature term \tfrac12\eta^{2}\nabla L^\top H\nabla L charges for long steps. With smoothness L, any \eta < 2/L decreases the loss and \eta = 1/L guarantees \lVert \nabla L\rVert ^{2}/(2L). Slope says which way, curvature says how far.
Next question → Every step goes downhill, but downhill to where? When is the bottom we reach the only bottom? (Section 3)
3. Convexity: When Is a Local Minimum the Global One?
The question: gradient descent stops where the slope is zero. When can we promise that point is the best possible weights, and is our ETA loss such a case?
The 30-second version: a convex function is a bowl with no false bottoms: every local minimum is global. Three equivalent tests: the chord between two points lies above the curve, the tangent lies below it, the curvature (Hessian) is never negative (PSD). Least squares has Hessian (2/n)X^\top X, always PSD, so it's convex; with full rank it's strongly convex and the minimum is unique. Logistic loss is convex too; neural-network losses are not. First read: the three tests, the table of which losses are convex.
A convex set contains the straight segment between any two of its points: a disk, a box, a half-plane. A convex function is one whose region on or above the graph is a convex set: it curves up everywhere, with no dents and no second valley. Pour water in and it all runs to one place.
The three tests
- L
- the loss as a function of the weights, e.g. the ETA model's MSE
- \mathbf a
- one point in weight space (a pair of weights)
- \mathbf b
- another point in weight space
- \theta
- the mixing fraction: 1 gives \mathbf a, 0 gives \mathbf b, between gives a point on the segment
- \nabla L(\mathbf a)
- the gradient at \mathbf a: the vector of slopes, one per weight; points uphill
- \nabla L(\mathbf a)^\top(\mathbf b - \mathbf a)
- the rise the tangent plane at \mathbf a predicts for the move to \mathbf b (^\top = transpose)
- H(\mathbf w)
- the Hessian at \mathbf w: the matrix of second derivatives, the curvature in every direction
- \succeq
- is positive semi-definite: curves up, or stays flat, in every direction
- \mathbf w
- any point in weight space (bold = a vector)
For a twice-differentiable function all three are equivalent (Card 3). Memory hook: "Chord above, tangent below, curvature up." The tangent test is the useful one in proofs: a convex function lies above every tangent plane, so a point with zero gradient has a flat tangent plane that the whole function sits above. That point is the global minimum.
You need this: positive semi-definite (PSD). A symmetric matrix H is PSD (H \succeq 0) if v^\top Hv \ge 0 for every vector v: curving up, or flat, in every direction. Equivalently, all its eigenvalues are ≥ 0. Positive definite (H \succ 0) means strictly positive: curving up in every direction.
Why local means global. Suppose w_{0} is a local minimum and some w_{1} were lower. Walk along the chord from w_{0} toward w_{1}. By the chord test, every point on it is no higher than the chord, which slopes down toward w_{1}. So there are lower points arbitrarily close to w_{0}, contradicting "local minimum" (Card 4). Convexity forbids false bottoms.
MSE is convex, and when it's strictly so. For least squares, v^\top Hv = (2/n)v^\top X^\top Xv = (2/n)\lVert Xv\rVert ^{2} \ge 0 for every v: PSD, so convex, always. If X has full column rank, Xv \ne 0 for v \ne 0, H is positive definite, and the minimum is unique. If two columns are collinear (distance in km and in miles), there's a flat valley of equally good weights: still convex, no longer unique.
Second-order conditions at one point. For any smooth loss, a local minimum needs \nabla L = 0 and H \succeq 0 there; \nabla L = 0 with H \succ 0 guarantees one. The gap between the two is the degenerate case (w^{4} and w^{3} both have H = 0 at 0; only one is a minimum), settled by higher-order terms (Section 11). Convexity is the global version of the same test: H \succeq 0 at every w, not just at one point.
Strong convexity μ and smoothness L
Two numbers summarize a bowl: its flattest curvature \mu and its steepest curvature L.
- \mu
- the flattest curvature the bowl ever has (5/3 on the toy); \mu > 0 means strongly convex
- I
- the identity matrix: curvature exactly 1 in every direction
- \preceq
- is no more curved than, in every direction (the difference is PSD)
- L
- here the smoothness constant: the steepest curvature anywhere (5 on the toy), not the loss
\mu-strongly convex means curving up by at least \mu in every direction. On the toy, \mu = 5/3 \approx 1.667 and L = 5. Their ratio, \kappa = L/\mu = 3, will decide the step size and the speed in Sections 4–5.
Which losses are convex (in the weights)
| Loss with a linear model | Convex? | Smooth? | Unique minimum? | Notes |
|---|---|---|---|---|
| MSE | Yes | Yes | If X has full rank | Hessian (2/n)X^\top X |
| MAE | Yes | No (kink at 0) | Not always | Needs subgradients (Section 12) |
| Huber | Yes | Once-differentiable | Usually | Smooth core, linear tails |
| Logistic (cross-entropy) | Yes | Yes | Unless the data are separable | Hessian (1/n)X^\top SX, S = \operatorname{diag}(p(1 - p)) (Card 3, Ch 3) |
| Hinge (SVM) | Yes | No | With the L2 term, yes | Ch 11 |
| MSE on a sigmoid output | No | Yes | No | Why Ch 3 uses cross-entropy |
| Any neural network | No | Mostly | No | Symmetries alone create many equal minima |
Operations that keep convexity: SAM
You rarely test a new loss from scratch. You build it from convex pieces with operations that preserve convexity. Memory hook: SAM keeps you convex:
- Sums with non-negative weights: \text{loss} + \lambda\cdot \text{penalty} is convex if both are.
- Affine maps inside: if \ell is convex, then \ell(Xw + b) is convex in w. This is why "convex loss of a linear model" is always convex.
- Max of convex functions: \text{hinge} = \max(0, 1 - y\cdot\text{score}) is convex.
Convex isn't the same as easy. MAE, hinge and the lasso penalty are convex but have kinks, where the gradient doesn't exist. Plain gradient descent there jitters instead of landing (Section 12 fixes it with proximal steps). And convexity alone gives slow worst-case rates (Section 5); it's strong convexity that gives fast ones.
The trap interviewers set: "convex" means convex in the weights, not in the inputs. A model \hat y = w_{1}x + w_{2}x^{2} is curved in x but its MSE is convex in (w_{1}, w_{2}). A network's output may be linear in x along some slice and still have a wildly non-convex loss in its weights.
Remember: Convex = no false bottoms, so local = global. Chord above, tangent below, curvature up. MSE's Hessian (2/n)X^\top X is PSD, so it's convex; full rank makes it strongly convex and the answer unique. SAM (non-negative Sums, Affine inside, Max) keeps you convex. Convex is about the weights.
Next question → On a convex bowl gradient descent can't get lost, but it can still crawl or explode. What decides how big a step we can take? (Section 4)
4. The Learning Rate and the Hessian
The question: on the six-order bowl, \eta = 0.3 converges in 10 steps, \eta = 0.39 needs 135, and \eta = 0.41 explodes. Where do those cliffs come from, and why did the run on distance in meters never move?
The 30-second version: near the bottom every loss looks like a quadratic, and gradient descent on a quadratic splits into independent 1-D problems, one per Hessian eigenvector. Each shrinks its error by (1 - \eta\lambda_{i}) per step. So \eta must be below 2/\lambda_{\max} or the steepest direction explodes, and progress along the flattest direction is limited by \lambda_{\min}. Their ratio, the condition number \kappa, decides everything: big \kappa means a narrow ravine and zigzags. Centering and scaling features shrink \kappa; correlation between features doesn't go away with scaling. First read: Figure 2.2, the regime table, the meters story.
You need this: eigenvalues and eigenvectors
A matrix usually turns a vector as well as stretching it. An eigenvector v is a direction it only stretches: Hv = \lambda v, and \lambda is the eigenvalue, the stretch factor. A symmetric matrix like H has a full set of perpendicular eigenvectors. Think of them as the bowl's own axes: along each one, the bowl is a simple 1-D parabola with curvature \lambda.
Gradient descent in the eigenbasis
Write the error as e = w - w^*. For a quadratic, \nabla L = H(w - w^*), so one step is:
- \mathbf e
- the error: current weights minus the optimum (3, 2) (bold = a vector)
- k
- the step number
- \eta
- the learning rate (step size)
- H
- the Hessian: the bowl's curvature matrix, (2/n)X^\top X
- I
- the identity matrix
Rotate into the eigenbasis (c_{i} = v_{i}^\top e) and the matrix becomes a list of numbers. Each coordinate evolves on its own (Card 5):
- c_{i,k}
- the error's coordinate along eigenvector i after k steps
- c_{i,0}
- that coordinate at the start
- \lambda_i
- the curvature (eigenvalue) along eigen-direction i: 5 or 5/3 on the toy
That's the whole theory of the learning rate. Each mode's factor 1 - \eta\lambda_{i} must have absolute value below 1, so:
- \lambda_{\max}
- the largest curvature (biggest Hessian eigenvalue): 5 on the toy
This is the bound Ch 1, Section 4 stated without proof.
The toy's two axes
The toy Hessian has eigenvalue \lambda = 5 along (1, 1)/\sqrt{2} and \lambda = 5/3 along (1, -1)/\sqrt{2}. Read them as stories:
- The steep direction (1, 1) is "more distance and more items together". The data pin that down well: it changes predictions a lot.
- The flat direction (1, −1) is "trade distance for items". Because distance and items are correlated, the data can barely tell those apart, so the loss hardly changes. This is where gradient descent dawdles and zigzags.
So \eta_{\max} = 2/5 = 0.4. Now the cliffs in the question make sense:
| \eta | Steep-mode factor 1 - 5\eta | Flat-mode factor 1 - (5/3)\eta | Steps to 10^{-6} | What you see |
|---|---|---|---|---|
| 0.05 | 0.75 | 0.917 | 55 | Smooth, slow |
| 0.1 | 0.5 | 0.833 | 27 | Smooth |
| 0.2 = 1/\lambda_{\max} | 0 | 0.667 | 12 | Steep mode gone in one step |
| 0.3 = 2/(\lambda_{\max} + \lambda_{\min}) | −0.5 | 0.5 | 10 | Best fixed step: slight zigzag |
| 0.35 | −0.75 | 0.417 | 24 | Zigzag |
| 0.39 | −0.95 | 0.35 | 135 | Violent zigzag, barely converging |
| 0.41 | −1.05 | 0.317 | Diverges | Each bounce 5% bigger |
A negative factor means the error flips sign every step: the iterate bounces across the valley. Its size is what matters. Memory hook: "1/\lambda smooth, 2/\lambda the cliff." Below 1/\lambda_{\max} every mode shrinks monotonically; between 1/\lambda_{\max} and 2/\lambda_{\max} the steep mode oscillates but shrinks; above 2/\lambda_{\max} it grows without bound.
Watch \eta = 0.2 kill the steep mode. From (0, 0) it jumps to (2.667, 2.333), and from there walks to (2.778, 2.222), (2.852, 2.148)… exactly along the flat direction (1, −1), the error shrinking by 2/3 each step. At \eta = 0.41 the iterates go (5.47, 4.78), (0.19, −0.71), (5.88, 4.91): bouncing ever wider along (1, 1).
The best fixed step and the condition number
With one \eta for all modes, the slowest mode sets the pace: we want to minimize the largest \lvert 1 - \eta\lambda_{i}\rvert. The worst modes are always the two extremes, and the best compromise makes them equal and opposite (Card 6):
- \eta^*
- the best fixed learning rate (star = the optimum): 0.3 on the toy
- \lambda_{\min}
- the smallest curvature (Hessian eigenvalue): 5/3 on the toy
- \rho
- the worst per-step shrink factor over all directions: 0.5 on the toy
- \max_i
- the largest over all eigen-directions i
- \kappa
- the condition number: steepest over flattest curvature, how elongated the bowl is (3 on the toy)
On the toy: \eta^* = 2/(5 + 5/3) = 0.3, \kappa = 3, \rho = 0.5. The condition number \kappa is how elongated the bowl is. \kappa = 1 is a round bowl: the gradient points straight at the bottom and one step of \eta = 1/\lambda lands there. Large \kappa is a narrow ravine: the gradient points mostly across the ravine, so steps zigzag between the walls while creeping along the floor.
The meters story
The dispatch logs distance in meters. Use that column, centered, and the steep eigenvalue jumps to 3.33 million (a weight per meter multiplies distances in the thousands, so the curvature along it is a million times the km version) while the flat one is 2.5. The same data, four ways:
| Features | \kappa, six orders | \eta_{\max}, six orders | \kappa, 10k orders |
|---|---|---|---|
| km and items, centered (our convention) | 3 | 0.4 | 3.02 (2.15 standardized) |
| km and items, uncentered, with an intercept | 183 | 0.0468 | 143 |
| Meters and items, centered | 1.33 \times 10^{6} | 6.0 \times 10^{-7} | — |
| Meters and items, uncentered, with an intercept | 8.9 \times 10^{7} | 9.4 \times 10^{-8} | 8.7 \times 10^{7} |
Now run it. With \eta = 0.1 the loss passes 10^{6} on the first step (it lands near 3 \times 10^{12}). With \eta = 5 \times 10^{-7}, safely under the cliff, 1,000 steps later the loss is still 6.65 instead of 1.667. The items weight has reached 0.0025, 0.125% of its way to 2; the distance weight has soaked up items' share and sits at the equivalent of 4.0 min/km instead of 3. That's the plateau in the question: one learning rate can't serve a direction with curvature 3 million and another with curvature 2.5. Any \eta small enough for the first freezes the second.
Center, then scale, then worry about correlation
- Center. An uncentered feature with mean 3 is correlated with the intercept column of 1s, which tilts and stretches the bowl. Centering removes it: the toy's \kappa falls from 183 to 3.
- Scale. Dividing by the standard deviation equalizes the curvatures along the feature axes: meters stop dwarfing items.
- Correlation stays. Two standardized features with correlation r give eigenvalues proportional to 1 + r and 1 - r:
- r
- the correlation between the two standardized features, distance and items (0.5 on the toy)
No rescaling of individual columns removes that; only rotating the axes does. Memory hook: "Center, then scale, then worry about correlation."
Preconditioning is the general fix: step along P\nabla L for a matrix P that makes the bowl look round. Standardizing is a diagonal P; whitening (rotating to decorrelate, then scaling) or P = H^{-1} makes \kappa = 1 exactly. That last choice is Newton's method (Section 9). Adam (Section 8) is a diagonal preconditioner learned on the fly.
Ridge conditions the bowl too. A penalty (\lambda/2)\lVert w\rVert ^{2} adds \lambda to every eigenvalue, so \kappa becomes (\lambda_{\max} + \lambda)/(\lambda_{\min} + \lambda): 3 → 2.25 on the toy at \lambda = 1. On the raw 10k model (penalizing all three weights), \lambda = 1 takes \kappa from 143 to 37.8 and GD's step count from 991 to 261. But the answer moves to (2.17, 3.19, 2.72) from (6.05, 2.99, 1.99): faster, because it is a different problem (Drill 20). The same \lambda is what makes a logistic loss strongly convex (\mu \ge \lambda).
Choosing η in practice
You rarely know \lambda_{\max}. Three practical answers:
- LR range test. Increase \eta geometrically over a few hundred steps and plot loss vs \eta; pick a value somewhat below where the loss starts to blow up. This finds the cliff empirically.
- Backtracking (Armijo) line search. Start big; halve until the step achieves a fraction c of the first-order predicted drop: L(w - \eta\nabla L) \le L(w) - c\cdot \eta\lVert \nabla L\rVert ^{2}. On the toy at w = 0 with c = 10^{-4}, \eta = 1 and 0.5 are rejected and \eta = 0.25 is accepted, dropping the loss to 3.76. Standard for full-batch convex problems; too noisy for minibatches.
- Estimate \lambda_{\max} by power iteration with Hessian-vector products (Section 9).
Remember: GD on a quadratic = independent 1-D problems along the Hessian's eigenvectors, each shrinking by (1 - \eta\lambda_{i}). 1/\lambda smooth, 2/\lambda the cliff. Best fixed step 2/(\lambda_{\max} + \lambda_{\min}), rate (\kappa - 1)/(\kappa + 1). Center, then scale, then worry about correlation: \kappa = (1 + r)/(1 - r) survives standardizing.
Next question → We know a step size that works. How many steps will it take to get within \varepsilon of the bottom, and what makes one method faster than another? (Section 5)
5. How Fast? Convergence Rates
The question: how many steps until the ETA model is within 10^{-6} of the best loss, and how does that count change with \kappa and with the shape of the loss?
The 30-second version: on a strongly convex, smooth loss, gradient descent converges linearly: the error shrinks by a constant factor \rho = (\kappa - 1)/(\kappa + 1) every step, so the step count grows like \kappa\cdot \log(1/\varepsilon). Without strong convexity it's sublinear: error \propto 1/k. Noisy gradients (SGD) are slower still, 1/\sqrt{k} or 1/k. Without convexity, only the gradient is guaranteed to shrink. Newton is quadratic: correct digits double per step. Momentum improves \kappa to \sqrt{\kappa}, the best any first-order method can do in the worst case. First read: the rates table and Figure 2.3.
What "rate" means
Plot log(error) against the step number k. A straight line down is a linear rate: error \propto \rho^{k}, a fixed number of correct digits per step ("linear" only on the log plot; it's exponentially fast). A curve that flattens is sublinear (1/k, 1/\sqrt{k}, 1/k^{2}): each step adds fewer digits than the last. A curve that plunges is quadratic: e_{k+1} \propto e_{k}^{2}, doubling the correct digits. Memory hook: "Sublinear crawls, linear adds digits, quadratic doubles them."
Linear convergence and the step count
On a quadratic with the best fixed step, each mode's error shrinks by at most \rho per step (Section 4). To shrink the distance to w^* by a factor \varepsilon (Card 7):
- \lVert \mathbf e_k\rVert
- the distance from the weights after k steps to the optimum \mathbf w^*
- \lVert \mathbf e_0\rVert
- the starting distance to the optimum
- \rho
- the per-step shrink factor, (\kappa - 1)/(\kappa + 1) with the best fixed step
- k
- the number of steps
- \varepsilon
- the target: shrink the distance to this fraction of where it started (e.g. 10^{-6})
- \kappa
- the condition number: steepest over flattest curvature (3 on the toy, 143 on the raw 10k model)
The bound uses \ln(1/\rho) = \ln((\kappa + 1)/(\kappa - 1)) = 2(1/\kappa + 1/(3\kappa^{3}) + \dots) \ge 2/\kappa, which is nearly tight for large \kappa. The step count is proportional to \kappa: a bowl ten times more elongated takes ten times as many steps.
The loss shrinks twice as fast as the distance. Loss suboptimality is quadratic in the error (L - L^* = \tfrac12e^\top He), so it shrinks by \rho^{2} per step. On the toy at \eta = 0.3, \rho = 0.5, so the loss gap shrinks by 0.25 per step: −0.602 on a \log_{10} scale, straight down. Reaching 10^{-6} takes \ln(10^{6})/\ln(4) \approx 10 steps, as Section 4 measured.
On the raw 10k model (unscaled km and items with an intercept), \kappa = 143.4 and \rho = 0.98615. The number of steps depends on what you count, so always name the criterion:
| Criterion (shrink by 10^{6}) | Theory | Simulated |
|---|---|---|
| Distance \lVert w - w^*\rVert | (\kappa/2)\cdot \ln 10^{6} \approx 990.6 | 991 |
| Loss gap L - L^* | half that, ≈ 495.3 | 496 |
Standardize the same features and \kappa drops to 2.15: 13 steps by distance and 6 by loss instead of hundreds (WP1 in the Interview Bank works this through).
Sublinear: a loss with a flat bottom
Strong convexity fails when the bowl's bottom is flat. Take f(w) = w^{4}: convex, minimum at 0, but f''(0) = 0. GD from w_{0} = 1 with \eta = 0.1 reaches w = 0.303, 0.109 and 0.0352 after 10, 100 and 1,000 steps, close to the approximation 1/\sqrt{8\eta k}. So the loss w^{4} falls like 1/k^{2} (fitted log-log slope −2.0), and getting to 10^{-4} takes 120 steps.
That 1/k^{2} is faster than the worst-case bound for convex smooth functions, which is 1/k (w^{4} isn't smooth everywhere, but on \lvert w\rvert \le 1, where GD stays, its curvature is at most 12). But it's still sublinear: there's no constant factor per step, because the gradient 4w^{3} vanishes faster than the distance to the minimum. The flatter the bottom, the weaker the pull.
The rates table
| Method | Setting | Rate | Steps for accuracy \varepsilon |
|---|---|---|---|
| GD | Strongly convex, smooth | Linear, \rho = (\kappa - 1)/(\kappa + 1) | \approx \kappa\cdot \log(1/\varepsilon) |
| GD | Convex, smooth (\eta = 1/L) | L(w_{k}) - L^* \le L\lVert w_{0} - w^*\rVert ^{2}/(2k) | \approx 1/\varepsilon |
| Subgradient method | Convex, non-smooth (MAE, hinge) | 1/\sqrt{k} with decaying \eta | \approx 1/\varepsilon^{2} |
| SGD, decaying \eta | Strongly convex | 1/k | \approx 1/\varepsilon |
| SGD, decaying \eta | Convex | 1/\sqrt{k} | \approx 1/\varepsilon^{2} |
| Heavy ball, tuned | Strongly convex quadratic | Linear, \rho = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) | \approx \sqrt{\kappa}\cdot \log(1/\varepsilon) |
| Nesterov | Strongly convex, smooth | Linear, \rho \approx 1 - 1/\sqrt{\kappa} | \approx \sqrt{\kappa}\cdot \log(1/\varepsilon) |
| Nesterov | Convex, smooth | 1/k^{2} | \approx 1/\sqrt{\varepsilon} |
| GD (\eta = 1/L) | Non-convex, smooth | \min_{j} \lVert \nabla L(w_{j})\rVert ^{2} \le 2L(L(w_{0}) - L^*)/k | \approx 1/\varepsilon^{2} to reach \lVert \nabla L\rVert \le \varepsilon |
| Newton | Near a non-degenerate minimum | Quadratic | \approx \log \log(1/\varepsilon) |
On the 10k data, SGD with \eta_{t} = c/t gives a fitted log-log slope of −1.02: the 1/k row, measured.
The speed limit. A first-order method uses only gradients (each step stays in the span of the gradients seen so far). Nesterov proved none can beat roughly (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) on all strongly convex quadratics: about \sqrt{\kappa}\cdot \log(1/\varepsilon) gradient evaluations. The fine print: it is a worst case over problems with more dimensions than steps. On a d-dimensional quadratic, conjugate gradient finishes in at most d steps (Section 9). Plain GD's \kappa is not optimal; momentum reaches the \sqrt{\kappa} limit (Section 7); going below it needs curvature (Section 9).
Without convexity, the guarantee is about the slope. Sum the descent lemma over k steps at \eta = 1/L and the losses telescope (Card 7, step 10): the smallest squared gradient seen so far is at most 2L(L(w_{0}) - L^*)/k, so about 1/\varepsilon^{2} steps make \lVert \nabla L\rVert \le \varepsilon. That promises a nearly flat point, not a minimum, and not the global one. On the double well of Section 11, started at w = 2 with smoothness 42 on [−2, 2], the bound after 50 steps is 16.0; the actual smallest squared slope is 1.2 \times 10^{-10}, at the local minimum 1.13. The promise is kept, at the wrong valley.
Linear rates without convexity: the PL condition
Neural nets aren't convex, yet overparameterized ones often train at a linear rate. The Polyak–Łojasiewicz (PL) condition explains how:
- \lVert\nabla L(\mathbf w)\rVert^2
- the squared length of the gradient (the vector of slopes) at \mathbf w: how steep it is there
- L(\mathbf w)
- the loss at weights \mathbf w
- L^*
- the lowest loss possible (star = the optimum)
- \mu
- the PL constant: how much slope is guaranteed per unit of excess loss
- \mathbf w
- any weight vector (bold = a vector)
In words: wherever the loss is far above its minimum, the slope is proportionally steep. That alone, with smoothness, gives L(w_{k}) - L^* \le (1 - \mu/L)^{k}(L(w_{0}) - L^*), with no convexity needed. Strongly convex functions satisfy PL, and so (approximately, near their solutions) do wide networks that can interpolate the data.
When to stop is unchanged from Ch 1, Section 4: gradient norm below a tolerance, relative loss change below a tolerance, or (the ML answer) validation loss stops improving.
Remember: Sublinear crawls, linear adds digits, quadratic doubles them. GD on a strongly convex bowl: \rho = (\kappa - 1)/(\kappa + 1), about (\kappa/2)\cdot \ln(1/\varepsilon) steps by distance, half that by loss. Flat bottoms (w^{4}) are sublinear. First-order methods can't beat \sqrt{\kappa} (in high dimension); momentum reaches it. PL gives linear rates without convexity; plain non-convex smoothness only promises a small gradient, in about 1/\varepsilon^{2} steps.
Next question → Every rate above assumed the exact gradient over all 10,000 orders. What if computing it once is too expensive? (Section 6)
6. SGD and Minibatches
The question: with 10,000 orders (and in production, 50 million), can we step using only 32 orders at a time, and what does that noise cost us?
The 30-second version: a minibatch gradient is an unbiased estimate of the full gradient, so it points downhill on average; its variance falls like 1/B. With a constant learning rate SGD never settles: it rattles around the minimum at a noise floor proportional to \eta/B. Decaying the learning rate removes the floor (Robbins–Monro), which is why schedules exist: warmup, step, cosine, warmup-stable-decay. The noise is not only a cost: it lets SGD escape saddles and acts as a regularizer. First read: the estimator box, Figure 2.4, the schedule table.
Recap from Ch 1, Section 4: batch, mini-batch and stochastic GD use all / some / one row per step; an iteration is one step, an epoch is one pass. This section supplies the why.
You need this: the mean and variance of an average
If X_{1}, …, X_B are independent draws with mean m and variance \sigma^{2}, their average has mean m and variance \sigma^{2}/B. Averaging keeps the target and divides the wobble by B, so the standard deviation shrinks only by \sqrt{B}. These are the estimator words from Ch 1, Section 5: unbiased (aimed at the truth) and variance (how much it wobbles).
The minibatch gradient is an estimator
The full gradient is an average of per-order gradients, G = (1/N)\sum\nabla\ell_{i}. A minibatch of B random orders gives (Card 8):
- \mathbf g_B
- the minibatch gradient: the average slope over a random batch of B orders
- \frac1B
- one over the batch size: turns the batch's sum into an average
- \sum_{i\in\mathcal B}
- add up over the orders i in the batch \mathcal B
- \nabla\ell_i(\mathbf w)
- order i's own gradient: the slopes of its squared miss, one per weight
- \mathbb E
- the average over every possible random batch
- \mathbf G
- the full gradient: the average over all N orders
- \lVert \mathbf g_B - \mathbf G\rVert^2
- the squared gap between batch and full gradient: the noise
- \operatorname{tr}\Sigma
- the total variance of per-order gradients (trace = sum of the diagonal of their covariance \Sigma)
- B
- the batch size: orders per step (e.g. 32)
- N
- the number of orders in the training set (10,000)
\Sigma is the covariance of per-example gradients. The last factor is the finite-population correction for sampling without replacement (how shuffled epochs work): it's ≈ 1 for B \ll N and exactly 0 at B = N, when the "minibatch" is the whole dataset.
On the 10k orders (standardized features, at w = 0, 2,000 batches per size):
| Batch size B | Measured E\lVert g_B - G\rVert ^{2} | Theory \operatorname{tr} \Sigma/B \cdot (N - B)/(N - 1) |
|---|---|---|
| 1 | 8,584 | 8,330 |
| 8 | 967 | 1,041 |
| 32 | 275 | 260 |
| 128 | 67.5 | 64.3 |
| 1,024 | 7.07 | 7.30 |
The log-log slope is −1.02: variance \propto 1/B, as promised. Memory hook for what a minibatch gradient is: "Unbiased, cheap, noisy."
Cost vs noise. B up by a factor k → the noise's standard deviation down by \sqrt{k}, the cost per step up by k. Per example processed, small batches are more efficient. On a budget of 50,000 examples with \eta tuned per batch size, B = 1 and B = 32 reach excess losses of 1.67 \times 10^{-4} and 1.32 \times 10^{-4}, while B = 1{,}024, which gets only about 50 steps, stops at 2.59 \times 10^{-3}. Large batches win on wall-clock time only because a GPU processes a batch of 1,024 almost as fast as a batch of 32.
The noise floor
Here's the catch. Near the minimum the true gradient goes to zero, but the minibatch noise doesn't. With a constant \eta, SGD stops descending and rattles around w^*. Model one direction of the bowl as (\lambda/2)e^{2} with gradient noise of variance \sigma^{2}/B. Each step pulls the error in by a fraction \eta\lambda and kicks it by \eta\xi (Card 9, WP2):
- e_k
- the error along one direction of the bowl after k steps (weight minus its optimum)
- e_{k+1}
- the same error one step later
- \eta
- the learning rate, held constant
- \lambda
- the curvature of the bowl along this direction
- \xi_k
- step k's gradient noise: random, mean 0, variance \sigma^2/B
- \mathbb E[L - L^*]
- the average excess loss where SGD settles: the noise floor
- \sigma^2
- the variance of one order's gradient along this direction
Summed over all directions, the floor is about (\eta/4B)\cdot \operatorname{tr} \Sigma, with \Sigma measured at w^* (\operatorname{tr} \Sigma = 109 on the 10k data). It depends on \eta/B: halve the step or double the batch, and the floor halves. Simulated against theory:
| \eta | B | Simulated floor | Theory (\eta/4B)\cdot \operatorname{tr} \Sigma |
|---|---|---|---|
| 0.01 | 1 | 0.287 | 0.273 |
| 0.01 | 32 | 0.0092 | 0.0085 |
| 0.1 | 8 | 0.407 | 0.341 |
| 0.1 | 1 | 8.24 | 2.73 |
The theory is for small \eta. In the last row it under-predicts badly, because per-example noise here grows as w wanders from w^* (the noise is multiplicative, not additive), which the formula ignores. The trap: a flat SGD loss curve is not convergence. It may be a floor that only a smaller step will break.
Breaking the floor: Robbins–Monro
To converge to w^* itself, the step must shrink, but not too fast:
- \sum_{t=1}^{\infty}
- add up over every step t of training, forever
- \eta_t
- the learning rate at step t, e.g. c/t
\eta_{t} = c/t satisfies both. One subtlety (Drill 17): the 1/t rate needs c large enough relative to the flattest curvature (2c\mu > 1); with c too small, convergence slows to t^{-2c\mu}.
Schedules: warm up, hold, cool down
Every modern schedule is a variation on three phases. On the 10k orders (20 epochs, B = 32, peak \eta = 0.1, 5 seeds), the final excess loss:
| Schedule | Shape of \eta(t) | Final excess loss |
|---|---|---|
| Constant | Flat at the peak | 0.0373 |
| Step | ×0.1 at 50% and 75% of training | 1.00 \times 10^{-4} |
| Cosine | Half a cosine from peak to 0 | 1.58 \times 10^{-4} |
| Linear warmup + cosine | Ramp up over 5%, then cosine | 1.65 \times 10^{-4} |
| Warmup-stable-decay (WSD) | Ramp, hold flat to 80%, linear decay to 0 | 2.57 \times 10^{-3} |
| 1/t | peak/(1 + t/steps per epoch) | 1.34 \times 10^{-3} |
Every decaying schedule beats constant by more than ten times; the floor is what the cool-down removes. Warmup barely matters on this convex problem; it earns its place in deep nets and with Adam (Section 8, Q31), where early gradients are large and curvature estimates unreliable. (1cycle, not run here, ramps up to a high peak and anneals below the start.) WSD (MiniCPM, 2024) is popular for LLMs because the long stable phase can be extended or branched without knowing the total budget in advance. Memory hook: "Warm up, hold, cool down."
Averaging the iterates. Another way through the floor: keep \eta constant and average the iterates. The noise in successive iterates partly cancels. On the 10k orders with constant \eta, the last iterate sits at excess 0.0373; the average over the second half of training sits at 5.76 \times 10^{-6}. This is Polyak–Ruppert averaging. Its cousin, an exponential moving average of the weights, is standard in diffusion models and many LLM recipes.
Variance reduction. A finite training set allows a third way. SVRG (Johnson and Zhang, 2013) computes the full gradient G(\tilde w) at a snapshot \tilde w once per round and steps along \nabla\ell_{i}(w) - \nabla\ell_{i}(\tilde w) + G(\tilde w). That is still unbiased, but its variance shrinks to zero as w and \tilde w approach w^*, so a constant \eta converges linearly. On the 10k orders (B = 1, \eta = 0.005), SVRG's excess loss is 6.2 \times 10^{-7} after five rounds and at machine precision after ten; plain SGD with the same \eta and the same number of gradient evaluations hovers around 0.155 (floor theory 0.136). SAGA stores the last gradient of every example instead of a snapshot. Both shine on convex finite sums and rarely pay in deep learning, where data augmentation makes the sum infinite.
Batch size and learning rate: the linear scaling rule
The floor depends on \eta/B. So if you multiply B by k, multiply \eta by k and the noise level stays the same while each epoch needs k times fewer steps: the linear scaling rule (Goyal et al., 2017, who added warmup to make it work on ImageNet). It breaks in two ways (Card 19):
- Stability. k\eta must stay below 2/\lambda_{\max}. On the 10k data (2 epochs, starting from B = 32 at \eta = 0.02, final excess 0.0087): B = 256 at 8\eta reaches 0.0207, while B = 256 at the old \eta stalls at 0.899. B = 2{,}048 at 64\eta = 1.28 exceeds the full-batch cliff of 0.733 and diverges; at 32\eta = 0.64 it survives but, with only 9 steps, ends at 0.332.
- The critical batch size. Once the batch is so large that the noise is negligible, doubling it no longer halves the number of steps. The gradient noise scale B_{\text{noise}} = \operatorname{tr} \Sigma/\lVert G\rVert ^{2} (McCandlish et al., 2018) marks the turn. On the 10k data at w = 0 it's only 3.75 (the gradient is huge compared with the noise); as training approaches w^*, \lVert G\rVert \to 0 and B_{\text{noise}} grows, which is why large batches pay off later in training.
Noise as a regularizer. SGD's noise isn't only a cost. Backward-error analysis shows that GD with step \eta follows, to first order, the gradient flow of a modified loss L + (\eta/4)\lVert \nabla L\rVert ^{2} (Barrett and Dherin, 2021), and SGD adds a similar penalty on how much minibatch gradients disagree (Smith et al., 2021). Both push toward flat regions. Noise also kicks the iterate off saddle points. Section 11 picks this up.
The assumption behind all of it: "unbiased" requires each minibatch to be a random sample, i.e. data i.i.d. and shuffled. Section 10 shows what happens when it isn't.
Remember: A minibatch gradient is unbiased, cheap, noisy: variance \operatorname{tr} \Sigma/B, times (N - B)/(N - 1) without replacement. Constant \eta leaves a floor \approx (\eta/4B)\cdot \operatorname{tr} \Sigma (small \eta); decay it (\sum\eta_{t} = \infty, \sum\eta_{t}^{2} < \infty), average the iterates, or (on a finite convex sum) cancel the noise with SVRG. Warm up, hold, cool down. Scale \eta with B until the cliff or the critical batch.
Next question → SGD makes each step cheap, but on a ravine it still zigzags. Can the optimizer remember where it was going? (Section 7)
7. Momentum and Nesterov
The question: on the unscaled 10k model (\kappa \approx 143), plain GD at its best fixed step needs about a thousand steps to shrink the distance to w^* by 10^{6} (991; 496 under the loss criterion). Can a ball that keeps rolling do better?
The 30-second version: momentum keeps an exponential moving average of past gradients and steps along it. Zigzags across the ravine cancel; the consistent push along the valley adds up. With the best \beta, the rate improves from (\kappa - 1)/(\kappa + 1) to (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1): asymptotically \sqrt{\kappa} \approx 12 times fewer steps (realized on the 10k run: 7.0× by loss, 8.8× by distance, because the optimally tuned heavy ball has a k\cdot \rho^{k} transient). Nesterov evaluates the gradient at the look-ahead point, which damps overshoot. The price: \beta adds inertia, so the effective step is \eta/(1 - \beta), and too much momentum oscillates. First read: the heavy-ball picture, the EMA box, Figure 2.5.
The heavy ball
Gradient descent is a hiker with no memory: every step points wherever the ground slopes here. In a ravine that means bouncing from wall to wall. A heavy ball rolling down the same ravine keeps its velocity: the side-to-side pushes alternate in sign and cancel, while the gentle push along the floor always points the same way and accumulates. That's momentum (Polyak, 1964).
The update, two equivalent forms:
- \mathbf v
- the velocity: a decaying running sum of past gradients (bold = one per weight)
- k
- the step number
- \beta
- the momentum coefficient: the fraction of the old velocity kept each step (e.g. 0.9)
- \nabla L(\mathbf w_k)
- the gradient at the current weights: one slope per weight; points uphill
- \mathbf w
- the ETA model's weights
- \eta
- the learning rate
- (\mathbf w_k - \mathbf w_{k-1})
- the previous step; momentum repeats a fraction \beta of it
Substitute one into the other and they're identical. Some libraries scale the gradient by (1 - \beta) inside the average; then \eta means something different, a source of silent bugs when porting hyperparameters.
Momentum is an exponential moving average
Unroll the velocity (Card 10):
- \nabla L
- the gradient (vector of slopes, one per weight); the subscript is the step it was taken at
- \mathbf g
- a gradient that stays the same every step (a steady slope)
- \frac{\mathbf g}{1 - \beta}
- where the velocity settles: 1/(1 - \beta) copies of the gradient (10 at \beta = 0.9)
Two consequences:
- Window. Weights \beta^{j} decay with a characteristic memory of about 1/(1 - \beta) steps: \beta = 0.9 remembers about 10, \beta = 0.99 about 100.
- Effective step. On a steady slope, the step grows to \eta/(1 - \beta): with \beta = 0.9, 10 times \eta (Drill 5). Adding momentum without lowering \eta is a bigger step, not just a smoother one.
Memory hook: "\beta is memory: window 1/(1 - \beta), step \eta/(1 - \beta)."
In practice: \beta = 0.9 (window 10, step 10\eta) is the default for SGD + momentum; \beta = 0.99 (window 100, step 100\eta) overshoots and spirals unless \eta is cut.
The optimal β on a quadratic
In the eigenbasis each mode now follows a two-term recursion (position and velocity), a 2 × 2 linear system per mode. Choose \eta and \beta so every mode's 2 × 2 matrix contracts at the same rate (Card 11):
- \eta^*
- the best learning rate for heavy ball (star = the optimum)
- \lambda_{\max}
- the largest curvature (Hessian eigenvalue) of the bowl
- \lambda_{\min}
- the smallest curvature
- \beta^*
- the best momentum coefficient
- \sqrt\kappa
- the square root of the condition number \kappa = \lambda_{\max}/\lambda_{\min} (about 12 on the raw 10k model)
- \rho
- the per-step shrink factor of the error
Compare with GD's (\kappa - 1)/(\kappa + 1): \kappa has become \sqrt{\kappa}. That's the first-order speed limit from Section 5.
- Toy (\kappa = 3): \beta^* = 0.0718 (tiny, because the bowl is nearly round), \eta^* = 0.3215, \rho = 0.268 vs 0.5 for GD. The loss gap falls 1.14 decades per step instead of 0.60; 7 steps to 10^{-6} instead of 10.
- Raw 10k (\kappa = 143.4, \sqrt{\kappa} = 11.98): \beta^* = 0.715, \eta^* = 0.0682, \rho = 0.846 vs 0.986. The ratio of log-rates is 12.0 = \sqrt{\kappa}.
But the realized speedup is smaller: 496 → 71 steps by loss (7.0×) and 991 → 112 by distance (8.8×). The reason is a classic follow-up: at the optimal (\eta, \beta), every mode's 2 × 2 matrix has a repeated eigenvalue, so the error behaves like k\cdot \rho^{k} rather than \rho^{k}. The factor k costs extra steps before the asymptotic rate takes over (the same reason Drill 4's 69 becomes 93 in simulation).
Nesterov: look before you leap
Nesterov's accelerated gradient computes the gradient at the point momentum is about to carry you to, not where you stand:
- \mathbf u_k
- the look-ahead point: where momentum alone would carry the weights
- \nabla L(\mathbf u_k)
- the gradient measured at the look-ahead point, not at \mathbf w_k
If the ball is about to overshoot, the look-ahead gradient already points back, so it brakes early. On a pure quadratic, a perfectly tuned heavy ball is as fast or faster: Nesterov with its standard settings (\eta = 1/L, \beta = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) = 0.846) needs 76 steps by loss and 188 by distance on the raw 10k model. Nesterov's real advantage is a guarantee: it achieves the accelerated rate on every smooth convex function (1 - 1/\sqrt{\kappa} per step with the constant \beta above when strongly convex; 1/k^{2} without strong convexity, using the growing momentum \beta_{k} = (k - 1)/(k + 2)), where heavy ball can fail to converge on some strongly convex non-quadratic functions (Lessard, Recht and Packard, 2016). In the curved detour valley below, tuned Nesterov needs 52 steps against heavy ball's 85.
The optimizer race
The detour valley f(a, b) = (1 - a)^{2} + 10(b - a^{2})^{2} is a Rosenbrock-style curved ravine: start at (−1.5, 2), minimum at (1, 1), Hessian eigenvalues at the minimum 101.6 and 0.394 (\kappa = 258). Steps to get within 10^{-3} of (1, 1), each method with \eta (and \beta) tuned by grid search: GD 522, heavy ball 85, Nesterov 52, Adam 84, L-BFGS 23 (27 function evaluations), Newton 10.
Momentum and SGD noise. Momentum averages minibatch noise too, but doesn't remove the floor: its effective step \eta/(1 - \beta) sets the floor's level just as \eta did. The EMA idea is everywhere: Adam's two moments (Section 8, \beta_{1} and \beta_{2} play exactly this role), EMA of weights (Section 6), BatchNorm's running statistics (Ch 17). Whenever you see a \beta next to a running average, think window 1/(1 - \beta).
Remember: \beta is memory: window 1/(1 - \beta), step \eta/(1 - \beta). Optimal heavy ball turns \kappa into \sqrt{\kappa}: rate (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1). Realized speedups are smaller than \sqrt{\kappa} because of the k\cdot \rho^{k} transient (7.0× and 8.8× on the 10k run vs \sqrt{\kappa} \approx 12). Nesterov looks ahead and is guaranteed on all smooth convex functions. Momentum doesn't remove SGD's floor.
Next question → Momentum fixes direction, but every weight still shares one learning rate. What if distance and items need very different step sizes? (Section 8)
8. Adaptive Methods: AdaGrad → RMSProp → Adam → AdamW
The question: distance in meters and items in units need step sizes a million times apart. Can each weight get its own learning rate, set automatically from its own gradient history?
The 30-second version: AdaGrad divides each weight's step by the root of its summed squared gradients, which helps rare features but shrinks steps forever. RMSProp uses a moving average instead. Adam adds momentum and corrects both moving averages for starting at zero: an unbiased estimator again. The result is a per-weight step of about \eta regardless of gradient scale. AdamW then decouples weight decay from the gradient, because inside Adam an L2 penalty gets divided by \sqrt{\hat v} and stops acting like weight decay. First read: the family table, the four D's, the bias-correction box.
The family tree
| Method | Per-weight update (g = gradient) | What it fixes | What it breaks |
|---|---|---|---|
| SGD + momentum | w -= \eta\cdot v, v = \beta v + g | Zigzag (Section 7) | One \eta for all weights |
| AdaGrad (2011) | w -= \eta\cdot g/\sqrt{\Sigma g^{2}} | Rare features get big steps | Steps shrink forever: \eta/\sqrt{t} |
| RMSProp (2012) | w -= \eta\cdot g/\sqrt{\text{EMA of } g^{2}} | AdaGrad's dying step | No momentum, biased early |
| Adam (2015) | w -= \eta\cdot \hat m/(\sqrt{\hat v} + \varepsilon) | Both, plus bias correction | L2 penalty gets rescaled |
| AdamW (2019) | Adam step, plus w -= \eta\lambda w | Decay independent of gradient scale | Still diagonal: ignores correlation |
AdaGrad's flaw in numbers: on a constant gradient its effective step is \eta/\sqrt{t}, so 0.316\eta, 0.1\eta and 0.0316\eta at t = 10, 100 and 1,000. Fine for a convex problem with sparse features; fatal for a long non-convex run.
Adam: the four D's
- \mathbf m
- the first moment: a moving average of recent gradients, one per weight (the direction)
- \mathbf v
- the second moment: a moving average of squared gradients (each weight's typical size²)
- t
- the step number, counting from 1
- \beta_1
- how much of the old \mathbf m to keep each step (0.9: about a 10-step window)
- \beta_2
- how much of the old \mathbf v to keep (0.999: about a 1,000-step window)
- \mathbf g
- the minibatch gradient; \mathbf g_t^2 squares it weight by weight
- \hat{\mathbf m}
- \mathbf m corrected for starting at zero (hat = the corrected estimate)
- \hat{\mathbf v}
- \mathbf v corrected for starting at zero
- \mathbf w
- the weights
- \eta
- the learning rate: roughly each weight's step size
- \epsilon
- a tiny constant (10^{-8}) so the division never blows up
- \lambda
- the decoupled weight-decay rate (e.g. 0.1)
Memory hook: the four D's of AdamW: Direction, Denominator, Debias, Decay. Defaults: \beta_{1} = 0.9, \beta_{2} = 0.999, \varepsilon = 10^{-8}. Each weight's step is (its average gradient) / (its typical gradient size): roughly \pm \eta whatever the gradient's units.
Bias correction: an unbiased estimator again
Both averages start at zero, so early on they're dragged toward zero. For a stationary gradient, unrolling the EMA gives (Card 12):
- \mathbb E
- the average over the random minibatches
- \sum_{j=0}^{t-1}
- add up over the past t gradients, j steps back
- j
- how many steps back a gradient is; \beta_1^{\,j} is its weight
- \mathbb E[\mathbf g]
- the true average gradient, which \mathbf m is meant to estimate
After 10 steps with \beta_{1} = 0.9, m carries only 1 - 0.9^{10} = 0.651 of the true mean (Monte Carlo: 0.650); dividing by that factor makes it unbiased. The same holds for v with \beta_{2}, and there the effect is huge because \beta_{2} is close to 1.
What goes wrong without it: on step 1, m_{1} = 0.1g and v_{1} = 0.001g^{2}, so the uncorrected step is \eta\cdot 0.1/\sqrt{0.001} = 3.16\eta instead of \eta (Drill 7). Correct m but forget v and the step is inflated by \sqrt{1/(1 - \beta_{2}^{t})}: 10.0×, 3.24× and 1.26× at t = 10, 100 and 1,000 (Drill 8). Early oversized steps are one reason fine-tuning with Adam needs warmup (Q31); RAdam (2020) rectifies the variance of \hat v directly instead.
What Adam is invariant to, and what it isn't
- Gradient scale: yes. Multiply the whole loss by 1,000 and Adam's trajectory is unchanged (maximum difference 1.5 \times 10^{-9}, i.e. \varepsilon effects). GD's would differ by 2.39 \times 10^{12} after 5 steps.
- A feature's units: no, but much better than GD. Under x \to c\cdot x, the gradient scales by c and the curvature by c^{2}. GD needs step sizes c^{2} apart; Adam's steps are \approx \eta in each weight's own units, so it needs them c apart. For meters (c = 1{,}000), GD's needed disparity of 10^{6} becomes 10^{3}. With \eta = 0.05, Adam's first step moves the distance weight by 50 min/km (vs 0.05 in km) yet, 300 steps later, reaches loss 1.676 vs 1.667 in km. GD at any single \eta diverges or freezes. Shrink Adam's \eta a thousandfold to tame meters and the items weight stalls at 0.0117 after 300 steps.
- Correlation: no. Adam is a diagonal preconditioner: it rescales each axis, but a ravine that runs diagonally (two correlated features) stays a ravine. Recall "Center, then scale, then worry about correlation": Adam does the scaling for you, not the worrying.
\varepsilon keeps the division safe when \hat v \approx 0. It also caps the step for weights with tiny gradients, so in low-precision training it's often raised (10^{-6} or larger).
L2 penalty vs decoupled weight decay
Under plain SGD, adding (\lambda/2)\lVert w\rVert ^{2} to the loss and shrinking the weights directly are identical (Card 13):
- \nabla L
- the gradient of the data loss alone: one slope per weight; points uphill
- \lambda
- the L2 penalty strength in (\lambda/2)\lVert\mathbf w\rVert^2; its gradient is \lambda\mathbf w
- (1 - \eta\lambda)
- the factor every weight shrinks by each step
Inside Adam they aren't. The penalty's gradient \lambda w goes into m and v and gets divided by \sqrt{\hat v} along with everything else. AdamW (Loshchilov and Hutter, 2019) applies the shrink outside the adaptive step instead. Two consequences:
- A weight with no data gradient. Start w = 1, \eta = 10^{-3}, \lambda = 0.1 (Drill 9). AdamW decays it smoothly: (1 - 10^{-4})^{1000} = 0.905 after 1,000 steps. Adam + L2 normalizes the penalty gradient away, so it marches at about \eta per step (0.902 after 100), then slows as \hat v remembers the larger early gradients: 0.258 after 1,000, 0.021 after 2,000, never crossing 0. And the answer is the same for \lambda = 10^{-3}, 0.1 or 1: under Adam, the penalty's strength stops mattering.
- Where they converge. On the raw 10k model with \lambda = 0.1, Adam + L2 lands on the ordinary ridge solution (6.21, 2.97, 1.96), matching it to 3–4 digits, because its fixed point is wherever the penalized gradient is zero. AdamW's fixed point is \nabla L + \lambda\sqrt{\hat v_{j}}\cdot w_{j} = 0: ridge with a per-weight strength \lambda\sqrt{\hat v_{j}} (0.45 for distance, 0.38 for items here). It lands on (6.69, 2.89, 1.89), shrinking the higher-gradient weight most.
Neither is "wrong", but they are different models, and in deep nets (which never reach a fixed point) AdamW's decay is a clean (1 - \eta\lambda) shrink you can tune separately from \eta. **The trap: torch.optim.Adam(weight_decay=…) is L2, not AdamW.**
When SGD + momentum still wins. On some vision benchmarks well-tuned SGD + momentum has generalized better than Adam (Wilson et al., 2017); much of the gap closes with AdamW and equal tuning effort. Treat it as an empirical question per domain (Q29).
Newer optimizers, one line each. AMSGrad keeps the running max of \hat v so the step never grows (Reddi et al., 2018, built a simple convex problem on which Adam provably fails to converge). RAdam switches the adaptive step on only once \hat v is reliable (built-in warmup). Adafactor stores v as row and column factors to save memory. LAMB/LARS scale each layer's step by \lVert w\rVert /\lVert \text{update}\rVert for huge batches. Lion steps by the sign of a momentum mix, with one state per weight. Shampoo preconditions each layer with Kronecker factors. Muon orthogonalizes each weight matrix's momentum update.
Remember: Direction, Denominator, Debias, Decay. Bias correction makes \hat m and \hat v unbiased (E[m_{t}] = (1 - \beta_{1}^{t})E[g]); without it the first step is 3.16\eta. Adam is invariant to gradient scale, cuts a feature-unit disparity from c^{2} to c, and ignores correlation. Inside Adam, L2 is divided by \sqrt{\hat v}; AdamW decouples the decay.
Next question → Adam guesses curvature from gradient sizes, one weight at a time. What if we used the actual curvature, the full Hessian? (Section 9)
9. Second-Order Methods
The question: Newton's method lands on the six-order optimum (3, 2) in one step from anywhere. Why doesn't everyone use it?
The 30-second version: Newton fits the local quadratic (Taylor to second order) and jumps to its bottom: w \leftarrow w - H^{-1}\nabla L. On a quadratic that's exact in one step; near any well-behaved minimum, correct digits double each step. For logistic regression, Newton is IRLS, the ~10-iteration solver in statsmodels. Gauss–Newton drops the hard part of the Hessian for least-squares problems; quasi-Newton (BFGS, L-BFGS) builds a Hessian estimate from gradient differences. Deep learning avoids them because the Hessian is p \times p, minibatch curvature is noisy, and Newton is attracted to saddle points. First read: the one-step picture, the IRLS comparison, the three toos.
Newton's step: jump to the bottom of the local bowl
Take the second-order Taylor model from Section 2 and minimize it over \Delta: set its gradient \nabla L + H\Delta to zero (Card 14):
- \Delta
- the Newton step: the jump to the bottom of the local quadratic bowl
- H^{-1}
- the inverse Hessian: undoes the curvature, dividing each direction by its eigenvalue
- \nabla L
- the gradient: the vector of slopes, one per weight; points uphill
- \mathbf w
- the weights (w_{\text{km}}, w_{\text{items}})
- k
- the iteration number
It's GD with the perfect preconditioner: H^{-1} rescales every eigen-direction by 1/\lambda_{i}, so \kappa becomes 1. On the toy, from (0, 0), (−5, 10) or (100, −3): one step to (3, 2), gradient exactly zero. That's the normal equation again (Ch 1, Section 3): for least squares, one Newton step is the closed form. Newton is also affine invariant: rescale, rotate or re-center the features and its iterates map exactly onto the old ones, so feature scaling doesn't matter to it (Q22).
Quadratic convergence, and when it fails
Near a minimum where H is positive definite (and changes smoothly: a Lipschitz Hessian), from a close enough start, the error squares each step: correct digits double. The classic quant-screen example (Drill 11): \sqrt{2} by Newton on x^{2} - 2 from x_{0} = 1 gives the iterates 1, 1.5, 1.416667, 1.414216, 1.41421356237, with 0, 1, 2, 5 and 11 correct digits.
Three ways it fails, and the fixes:
- Degenerate minimum. On f(w) = w^{4}, the Newton step is w - 4w^{3}/(12w^{2}) = (2/3)w: linear, factor 2/3 (Drill 10). Quadratic convergence needs H \succ 0 at the minimum.
- Far from the minimum, the quadratic model is a bad map and the step can overshoot. Fix: damping, w - \alpha\cdot H^{-1}\nabla L with a line search, or a trust region (only trust the model within a radius).
- Non-convex or singular H. Newton goes to whatever stationary point the local quadratic has, saddles and maxima included. Fix: Levenberg–Marquardt, (H + \lambda I)^{-1}\nabla L, which blends Newton (\lambda \to 0) and GD (\lambda large).
Gauss–Newton: Newton for least-squares problems
For a nonlinear model fitted by squared error, L = (1/n)\sum r_{i}(w)^{2}, the Hessian is (2/n)(J^\top J + \sum r_{i}\nabla^{2}r_{i}), where J is the Jacobian of the residuals. Gauss–Newton drops the second term, which is small when residuals are small or the model is nearly linear:
- H
- the true Hessian (curvature matrix) of the least-squares loss
- n
- the number of orders
- J
- the Jacobian: how each order's residual changes with each parameter (one row per order)
- J^\top J
- the Gauss–Newton curvature: the Jacobian transposed times itself, always PSD
- (J^\top J)^{-1}
- its matrix inverse: undoes that curvature
- J^\top
- the Jacobian transposed: rows become columns
- \mathbf r
- the residuals: each order's actual ETA minus the model's prediction
On the 10k orders, fit \text{ETA} = a + c\cdot\text{km}^q + d\cdot\text{items} (true q = 1). Pure Gauss–Newton from (0, 1, 0.5, 0) overshoots and diverges; with step halving on the loss it converges in 11 iterations to \hat q = 0.991, (a, c, d) = (5.97, 3.06, 1.99). Damping isn't optional far from the solution.
IRLS: Newton for logistic regression (preview of Ch 3)
Next chapter's logistic loss has gradient X^\top(p - y)/n and Hessian X^\top SX/n with S = \operatorname{diag}(p(1 - p)) (Card 3). The Newton step is then a weighted least-squares solve, re-weighted every iteration: iteratively reweighted least squares (IRLS). On the 10k orders with the label "late" (ETA > 30 min, 15.9% of orders):
| Method | Features | Iterations | Why |
|---|---|---|---|
| Our IRLS | Standardized | 10 to \lVert \Delta w\rVert < 10^{-8} (7 to relative suboptimality 10^{-6}) | Quadratic convergence |
| statsmodels Logit (Newton) / GLM (IRLS) | Standardized | 10 / 9 (Newton: 10 on raw features too) | Same algorithm; affine invariant |
| GD at \eta = 1/L | Standardized | 553 | Hessian \kappa = 27.5 at the optimum |
| GD at the best fixed \eta (14) | Standardized | 70 | Tuned step |
| GD at \eta = 1/L | Raw km and items | 49,989 | \kappa = 2{,}581 |
All converge to the same coefficients (−5.05, 4.11, 1.78), and sklearn's unpenalized lbfgs matches to 6 \times 10^{-9}.
Quasi-Newton: BFGS and L-BFGS
BFGS never forms H. It builds an approximation of H^{-1} from what it observes: after a step s = w_{k+1} - w_{k} the gradient changed by y = \nabla L_{k+1} - \nabla L_{k}, and any good curvature model must satisfy the secant condition Bs = y. Each iteration makes the smallest rank-2 correction that does. Its line search must satisfy the Wolfe conditions (enough decrease, plus a curvature check), which guarantee s^\top y > 0 and so keep the estimate positive definite. L-BFGS keeps only the last m pairs (s, y), typically 5–20, so memory is O(mp) instead of O(p^{2}). In the detour valley: 23 iterations, 27 function evaluations, between GD's 522 and Newton's 10. (On the 2-D toy, scipy's BFGS takes 8; with exact line searches BFGS finishes a 2-D quadratic in 2.)
Conjugate gradient: momentum that knows the bowl. For a quadratic, CG makes each new direction conjugate to all earlier ones (p_{i}^\top Hp_{j} = 0) and line-searches exactly along it, so no step undoes an earlier one. It solves a d-dimensional quadratic in at most d steps: 2 on the toy (through (2.67, 2.34) to (3, 2)) and 3 on the raw 10k model (relative error 3 \times 10^{-13}, against GD's 991). With many more dimensions than steps its worst-case rate is the \sqrt{\kappa} rate of tuned heavy ball, without needing \mu or L. Run on Hessian-vector products, CG is the inner solver of Hessian-free and truncated-Newton methods.
Hessian-vector products without the Hessian. Hv is the directional derivative of the gradient, computable by autodiff for about the cost of two gradients (Pearlmutter's trick). Run power iteration with it to get \lambda_{\max}: on the raw 10k model, 30 iterations give 49.93, matching the exact eigenvalue. That's how you find the cliff 2/\lambda_{\max} in a model too big to form H (WP1).
Why deep learning doesn't: the three toos
- Too big. A 1-million-parameter model's Hessian has 10^{12} entries: 4 TB in float32, and a solve costs on the order of 10^{18} operations. LLMs have billions of parameters.
- Too noisy. A minibatch Hessian is a far noisier estimate than a minibatch gradient, and inverting noise amplifies it.
- Too saddle-happy. Newton jumps to the nearest stationary point of the local quadratic, and in high dimensions most are saddles (Section 11). GD at least moves downhill.
Memory hook: the three toos: too big, too noisy, too saddle-happy.
Second-order ideas live on in diagonal or structured form: Adam (diagonal, Section 8), K-FAC and Shampoo (Kronecker-factored blocks per layer), the natural gradient (Fisher information instead of the Hessian), and Gauss–Newton approximations in some LLM optimizer research.
Remember: Newton: w \leftarrow w - H^{-1}\nabla L, exact in one step on a quadratic, quadratic convergence near a non-degenerate minimum, only linear (2/3) on w^{4}. Gauss–Newton uses J^\top J; IRLS is Newton for logistic regression; L-BFGS learns curvature from gradient differences in O(mp) memory; CG solves a d-dimensional quadratic in d steps. Too big, too noisy, too saddle-happy.
Next question → We now have a toolbox of optimizers. Their guarantees rest on assumptions. When one breaks in a real training run, how do we see it, and what do we do? (Section 10)
10. Assumptions, Biased Gradients, and Reading the Loss Curve
The question: the nightly ETA retrain reads orders in the order they're stored (sorted by distance), and its loss curve has a sawtooth. Which assumption broke, and how do we read any loss curve?
The 30-second version: every guarantee in Sections 3–9 rests on five assumptions, SCUBA: Smooth, Convex, Unbiased gradients, Bounded noise, Aligned scales. Most real failures are U (data not shuffled, so the gradient is biased toward whatever comes next) or A (unscaled features). See it on the loss curve: the four N's are NaN, Noise, Nothing (plateau) and Notches (sawtooth or spikes). Test with a gradient check, a shuffle check, a one-batch overfit test, and an LR sweep. Fix with the four S's: Shuffle, Stratify, Sample by importance, Scale. First read: the SCUBA table, Figure 2.6, the symptom → cause → fix grid.
Remember: SCUBA
| Letter | Assumption | In plain words | If it breaks… |
|---|---|---|---|
| S | Smooth (L-Lipschitz gradient) | Curvature is bounded, so a small enough step is safe | No safe \eta: kinks (Section 12), exploding gradients |
| C | Convex | One bottom | Many stationary points; the start matters (Section 11) |
| U | Unbiased gradient | Each minibatch is a fair random sample | SGD converges to the wrong point, or cycles |
| B | Bounded noise variance | Gradient noise is finite | Heavy tails (storms, bad batches) → spikes |
| A | Aligned scales (small \kappa) | Features centered and scaled | Crawl or divergence; one \eta can't fit all |
The ladder, like Ch 1's: S for a safe step, C for one bottom, U for the right direction, B for a floor, A for speed. Break S and no \eta is safe; break C and you lose "global"; break U and you converge, confidently, to the wrong answer; break B and the floor (Section 6) is infinite; break A and every rate in Section 5 gets multiplied by \kappa.
Data bias: when the gradient is biased
U is the assumption data problems break. SGD's guarantee needs E[g_B] = \nabla L, which holds only if each batch is a random sample of the training set.
- Sorted data. The warehouse table is sorted by distance. One epoch of SGD in stored order (B = 32, \eta = 0.05, standardized features) ends at w far from w^*: distance 18.8 in weight space vs 0.22 shuffled, excess loss 350.5 vs 0.035. The intercept ends near 40 minutes instead of 21.6, because the last batches are the longest trips and a constant-step SGD remembers recent batches best. Over five epochs the loss sawtooths between 14.7 and 201 as the order cycles from short trips to long ones; shuffled, it stays between 9.008 and 9.237 (L^* = 9.007).
- Class-sorted batches. For next chapter's P(late) model, feeding all late orders first and then all on-time ones each epoch leaves the intercept 1.08 too low and the average predicted probability at 0.057 against a true late rate of 0.159. Stratified batches: intercept off by only 0.105, average prediction 0.156.
- Duplicated or oversampled rows silently reweight the loss. Fine if intended; then correct the predictions (Ch 3 shows the intercept shift).
- Curriculum ordering on purpose (easy examples first) is a deliberate bias; it can help early training but must end in random sampling.
Correct it: the four S's
- Shuffle every epoch (not once).
- Stratify batches so each carries the class or segment mix of the whole.
- Sample by importance, then reweight. Draw example i with probability p_{i} and scale its gradient so the estimate stays unbiased:
- \mathbf g
- the reweighted gradient estimate from one sampled order i
- N
- the number of training orders (10,000)
- p_i
- the probability of drawing order i
- \nabla\ell_i
- order i's own gradient: the slopes of its squared miss, one per weight
- \mathbb E
- the average over repeated random draws
- \sum_i
- add up over all N orders
- \frac1N
- one over the number of orders: turns the sum into an average
- \nabla L
- the full gradient over all orders; points uphill
- p_i^{\text{opt}}
- the sampling probability that gives the lowest variance
- \lVert\nabla\ell_i\rVert
- the length of order i's gradient
On the 10k orders at w = 0, sampling ∝ gradient norm cuts the variance from 8,330 to 3,558 (2.34× lower), and the estimate stays unbiased (Monte Carlo mean within 0.2% of the full gradient).
- Scale the features (A in SCUBA, Section 4).
Memory hook: the four S's of a fair gradient: Shuffle, Stratify, Sample by importance, Scale.
See it: the four N's
| Shape | Name | Real 10k run that made it | Usual cause |
|---|---|---|---|
| Shoots up, then NaN | NaN | Raw km, full batch, \eta = 1.0 (cliff 0.040): error ×48.9 per step, so loss ×2,390; non-finite at step 90 in float64, step 10 in float32 | \eta > 2/\lambda_{\max}; unscaled features; fp16 overflow |
| Fat, jagged band far above L^* | Noise | Standardized, B = 1, \eta = 0.1: late average 18.6, spikes to 129 (L^* = 9.0) | \eta/B too large; tiny batches |
| Drops once, then flat and high | Nothing | Raw meters, \eta = 1/\lambda_{\max} \approx 3.3 \times 10^{-8}: 542.7 → 53.1, then stuck (L^* = 9.0) | Ill-conditioning (A); vanishing gradients; dead units |
| Repeating sawtooth or sudden spikes | Notches | Stored (sorted) order, no shuffle: 14.7 ↔ 201 each epoch | Ordered data (U); bad batches; loss spikes |
Memory hook: NaN, Noise, Nothing, Notches.
Test it
| Test | What it checks | Pass mark |
|---|---|---|
| Gradient check (central differences) | The gradient code (S, and your bugs) | Relative error \lesssim 10^{-6} in float64 |
| Overfit one batch | Model, loss and optimizer can fit anything at all | Training loss → ≈ 0 on 1–10 examples |
| LR range test | Where the cliff is (A, S) | A clear blow-up point; pick well below it |
| Shuffle on/off A/B | Ordering bias (U) | Same final loss either way |
| Gradient norms per layer | Explosions, dead layers (B, S) | Stable, same order of magnitude |
The gradient check, done right: central differences [L(w + h\cdot e_{j}) - L(w - h\cdot e_{j})]/(2h), relative error \lVert g_{\text{analytic}} - g_{\text{numeric}}\rVert /\max(\lVert g_{\text{analytic}}\rVert , \lVert g_{\text{numeric}}\rVert ). On the toy: correct gradient 9.0 \times 10^{-12}; forget the factor 2 and it's 0.50. Run the same check in float32 with h = 10^{-5} and a correct gradient scores 3.1 \times 10^{-3}: rounding error swamps the difference, a false alarm. Gradient-check in float64.
import numpy as np
# Shuffle check: one epoch of SGD (B = 32, eta = 0.05) in stored order vs shuffled
rng = np.random.default_rng(2002) # the 10,000-order simulation
km = np.clip(rng.gamma(2.0, 1.6, 10_000), 0.2, 15)
items = np.minimum(1 + rng.poisson(1.2 + 0.25 * km), 12).astype(float)
eta_min = 6 + 3 * km + 2 * items + rng.normal(0, 3.0, 10_000)
X = np.column_stack([np.ones(10_000), (km - km.mean()) / km.std(), (items - items.mean()) / items.std()])
w_star = np.linalg.lstsq(X, eta_min, rcond=None)[0]
def one_epoch(order, lr=0.05, B=32):
w = np.zeros(3)
for s in range(0, len(order), B):
i = order[s:s + B]
w -= lr * (-2 / len(i)) * X[i].T @ (eta_min[i] - X[i] @ w)
return w
stored = np.argsort(km, kind="stable") # the warehouse table is sorted by distance
shuffled = np.random.default_rng(10).permutation(10_000)
for name, order in [("stored order", stored), ("shuffled", shuffled)]:
print(f"{name:12s} ||w - w*|| = {np.linalg.norm(one_epoch(order) - w_star):.2f}")
Fix it: symptom → cause → fix
| Symptom | Likely cause | Fix |
|---|---|---|
| NaN or loss explodes in the first steps | \eta above the cliff; unscaled inputs | Lower \eta; standardize; warmup |
| NaN after many steps | fp16 overflow; a bad batch; log(0) | bf16 or loss scaling; clip; check data; stable log-sum-exp |
| Plateau far above a baseline | Ill-conditioning; dead units; \eta far too small | Scale features; Adam; LR range test |
| Plateau just above the best loss | SGD noise floor | Decay \eta; bigger batch late; average iterates |
| Sawtooth with the epoch period | Ordered data | Shuffle every epoch |
| Rare large spikes | Heavy-tailed examples; Adam with tiny \hat v | Clip gradients; skip the batch; lower \eta or \beta_{2} |
| Train loss falls, validation rises | Overfitting, not optimization | Regularize, stop early (Ch 4) |
Gradient clipping on the storm orders. Ch 1, Q6 had storm deliveries that take far longer than normal. In the 10k orders, 101 (1%) have storms that add 40 minutes. Their gradients are huge and every batch that contains one jolts the weights. Clipping rescales any minibatch gradient whose norm exceeds a threshold c down to norm c (direction kept). With c at the 95th percentile of batch-gradient norms (5.84), the largest single step halves, 0.633 → 0.292, but the average weight error barely moves, 0.311 → 0.294. Sell clipping as what it is: insurance against blow-ups, not an accuracy gain. The fix for the storms' bias is in the loss (Huber, or a storm feature), not the optimizer.
A debugging order that saves days: data (shuffled? scaled? labels right?) → loss (correct formula, numerically stable) → gradient (check it) → step size (range test) → optimizer (only then tune \beta, \varepsilon, schedules).
Remember: SCUBA: Smooth, Convex, Unbiased, Bounded noise, Aligned scales. Unshuffled data bias the gradient: SGD converges confidently to the wrong place. Read curves by the four N's: NaN, Noise, Nothing, Notches. Fix gradients with the four S's: Shuffle, Stratify, Sample by importance (reweight 1/(Np_{i})), Scale. Gradient-check in float64. Clipping is insurance, not accuracy.
Next question → C in SCUBA fails for every neural network. Without convexity, what does the landscape look like, and does it matter where we start? (Section 11)
11. Non-Convex Landscapes
The question: if the ETA model becomes a neural network, the bowl is gone. Will gradient descent get stuck, and does it matter which minimum it finds?
The 30-second version: a non-convex loss has many stationary points; the Hessian's eigenvalue signs classify them: all positive is a minimum, mixed is a saddle. In high dimensions saddles vastly outnumber bad local minima, and plain GD can sit near one for a long time while noise escapes it fast. Where you start matters (initialization: break symmetry, keep scale right). Among minima, GD with step \eta can only settle where the curvature is below 2/\eta, so large steps and small batches prefer flatter minima, which often generalize better (an argument, not a theorem). First read: the 1-D double well, the saddle picture, flat vs sharp.
Two neighborhoods, the start decides
The double well f(w) = w^{4} - 3w^{2} + w has a global minimum at −1.30 (f = -3.51), a local minimum at 1.13 (f = -1.07), and a maximum at 0.17 between them. GD with \eta = 0.01 from w_{0} = -2, −0.5 and 0.1 ends at the global minimum; from 0.3 and 2 it ends at the local one. The maximum is the basin boundary: which side you start on decides where you end. Convexity's guarantee is gone.
Signs of the Hessian. At a stationary point the gradient term of the Taylor expansion vanishes and the curvature decides everything (Card 17). Memory hook: "Signs of the Hessian: all up = minimum, all down = maximum, mixed = saddle." (Some zero eigenvalues: degenerate, look at higher-order terms.) The double well's curvatures are 14.3 and 9.35 at its minima and −5.65 at its maximum.
Saddles: the mountain pass
The saddle toy f(a, b) = \tfrac12a^{2} - \tfrac12b^{2} + \tfrac14b^{4} has a saddle at (0, 0) (curving up in a, down in b) and minima at (0, ±1). Along a negative-curvature direction \lambda < 0, GD multiplies the error by 1 + \eta\lvert \lambda\rvert each step, so it does escape, but starting from a tiny offset it takes a while:
- b_k
- the saddle toy's b after k steps: the offset along the downhill direction
- b_0
- the starting offset (10^{-6} in the example)
- b_{\text{target}}
- the offset that counts as escaped (0.5 in the example)
- k
- the number of steps
- \eta
- the learning rate (0.1)
- \lambda
- the negative curvature along b at the saddle; |\lambda| is its size
From b_{0} = 10^{-6} with \eta = 0.1, GD takes 139 steps to reach \lvert b\rvert \ge 0.5 (on the pure quadratic \tfrac12(a^{2} - b^{2}), 145 steps to \lvert b\rvert = 1, matching \ln 10^{6}/\ln 1.1). Add Gaussian noise of 0.01 per step and the average escape takes 47 steps: the noise replaces the tiny start with an offset of order 0.01. In high dimensions, a stationary point with high loss almost always has some downhill direction, so saddles vastly outnumber bad local minima (Dauphin et al., 2014, an empirical and random-matrix argument rather than a theorem). That's the real reason "stuck in a local minimum" is usually the wrong diagnosis (Q35), and why Newton, which is attracted to saddles, is a poor fit (Section 9).
Two theorems bound the danger. From a random start, GD with a fixed step below 1/L almost surely never converges to a strict saddle (one with a negative eigenvalue; Lee et al., 2016), though in contrived cases it can take exponentially long to pass a chain of them (Du et al., 2017). Adding a small random kick whenever the gradient is tiny (perturbed GD) reaches an approximate local minimum in a step count that depends on the dimension only through log factors (Jin et al., 2017). SGD's noise plays that kick for free.
Plateaus and vanishing gradients look like saddles on the loss curve (the "Nothing" N): long flat stretches with tiny gradients, typical of saturated activations or poor initialization. Ch 17 owns the deep-net fixes.
Initialization. Two rules. Break symmetry: if two hidden units start with identical weights they get identical gradients forever, so start random. Keep the scale right: too large and activations explode, too small and gradients vanish; Xavier and He initialization (Ch 17) set the variance by layer width.
Big steps can't sit in sharp holes
Near any minimum the loss is a quadratic with Hessian H, and Section 4's cliff applies locally: GD with step \eta can only settle in a minimum whose sharpest curvature satisfies \lambda_{\max} < 2/\eta (Card 20). Take a landscape with a deep, narrow hole (curvature 44.4, at −2) and a wide, shallow basin (curvature 0.9, at +2). The threshold is \eta = 2/44.4 = 0.045. From 1,000 random starts with a little noise:
| \eta | Settles in the sharp hole | Ends in the flat basin | Trapped bouncing inside the sharp well |
|---|---|---|---|
| 0.01 (2/\eta = 200 > 44.4) | 19% | 64.1% | 0% |
| 0.1 (2/\eta = 20 < 44.4) | 0% | 67.3% | 19% (spread ≈ 0.16) |
The rest stall on a far plateau where the gradient is nearly zero. At \eta = 0.1 the sharp mode's factor is 1 - 0.1\cdot 44.4 = -3.44, so no start can settle there; but the hole's walls flatten out, so the bounce never throws those 19% clear. "Can't sit" means can't settle, not "always leaves". Memory hook: "Big steps can't sit in sharp holes", the local version of "2/\lambda the cliff". SGD noise does the same job: a larger \eta/B makes sharp minima harder to stay in.
Edge of stability. In real networks trained with full-batch GD, sharpness tends to rise during training until \lambda_{\max} hovers near 2/\eta, and the loss keeps falling non-monotonically there (Cohen et al., 2021). The cliff isn't avoided; training lives on it.
Do flat minima generalize better? Large-batch training finds sharper minima and sometimes generalizes worse (Keskar et al., 2017), which motivates the idea. But sharpness isn't invariant: rescaling one layer up and the next down leaves the function unchanged and makes any minimum as sharp as you like (Dinh et al., 2017). Treat "flat generalizes" as a useful heuristic, not a law. SAM (sharpness-aware minimization, Foret et al., 2021) acts on it directly: take the gradient at the worst point within a small radius, then step, so training prefers neighborhoods that are low everywhere.
Implicit regularization: which minimum does GD pick?
When there are more weights than equations, infinitely many weight vectors fit the data perfectly. Add 8 binary "menu flag" features to the six orders: 10 weights, 6 rows, every one fit exactly. GD from w = 0 converges to the minimum-norm interpolating solution X^{+}y (X^{+} is the pseudo-inverse, X^\top(XX^\top)^{-1} here), to within 1.2 \times 10^{-14} (norm 5.44). Why (Card 18): every gradient -(2/n)X^\top r lies in the row space of X, so iterates starting at 0 never leave it, and the only interpolating solution in the row space is X^{+}y. Start from a random w_{0} instead and GD reaches a different perfect fit, 2.06 away: X^{+}y plus the part of w_{0} the data can't see (the null-space component). The optimizer chose the answer the loss couldn't (Drill 16: minimizing (w_{1} + w_{2} - 2)^{2} from (0, 0) gives (1, 1); from (3, 0) gives (2.5, −0.5)).
This is implicit regularization; with GD's implicit (\eta/4)\lVert \nabla L\rVert ^{2} penalty and SGD's noise (Section 6) it's a leading explanation for why huge networks generalize (Ch 4). Early stopping is a third form: k steps of GD from 0 shrink each eigen-direction roughly like ridge with \lambda \approx 1/(\eta k) (Ch 4 owns it).
Most local minima are fine. In wide networks, the minima SGD finds tend to have similar loss, and they're often joined by low-loss paths (mode connectivity, Garipov et al. and Draxler et al., 2018). Two seeds ending at slightly different losses (Q12) is normal variation, not a stuck run.
Remember: Signs of the Hessian: all up = minimum, all down = maximum, mixed = saddle. Saddles outnumber bad minima in high dimensions; GD escapes at rate (1 + \eta\lvert \lambda\rvert ) per step, and noise escapes faster. Start random (symmetry) at the right scale. Big steps can't sit in sharp holes: settling needs \lambda_{\max} < 2/\eta. GD from 0 finds the minimum-norm solution.
Next question → So far the weights could go anywhere. What if they must obey rules (non-negative, a budget), or the loss has a kink like lasso's? (Section 12)
12. Constraints and Regularization as Optimization
The question: ops says the per-item time can't exceed 1.5 minutes (their SLA model), and the lasso version of the ETA model needs exact zeros. How does gradient descent respect a rule or handle a kink?
The 30-second version: a constraint turns into a price via the Lagrangian; at the optimum the KKT conditions hold: the gradient is balanced by the constraint's push, and that push (the multiplier) is non-zero only when the constraint is tight. On the toy, capping w_{\text{items}} at 1.5 gives w = (3.25, 1.5) with multiplier 1.25. Projected GD steps, then snaps back into the allowed set. Proximal GD steps on the smooth part, then applies the penalty's own "shrink" operator; for L1 that's soft-thresholding, which produces exact zeros where plain (sub)gradient descent only jitters around zero. First read: the toy KKT example, Figure 2.8, step-then-snap vs step-then-shrink.
You need this: the Lagrangian
Two ways to say "keep the weights in check" are the same problem: a constraint (minimize L subject to g(w) \le 0) or a penalty (minimize L + \lambda\cdot \text{penalty}). Ch 1, Section 10 drew ridge and lasso both ways. The bridge between them is the Lagrangian, which charges a price \mu \ge 0 per unit of constraint violation:
- \mathcal L(\mathbf w, \mu)
- the Lagrangian: the loss plus a price for breaking the constraint
- L(\mathbf w)
- the ETA model's training loss at weights \mathbf w
- \mu
- the multiplier: the price per unit of violation (1.25 for the 1.5-minute cap)
- g(\mathbf w)
- the constraint, written so it holds when \le 0 (e.g. w_{\text{items}} - 1.5)
KKT on the toy, by hand
Minimize the toy loss subject to w_{\text{items}} \le 1.5. The unconstrained answer has w_{\text{items}} = 2, so the cap binds. The KKT conditions (Card 15):
| Condition | Meaning | On the toy |
|---|---|---|
| Stationarity | \nabla L + \mu\nabla g = 0: the downhill pull is balanced by the wall's push | \partial L/\partial w_{\text{km}} = 0 and \partial L/\partial w_{\text{items}} + \mu = 0 |
| Primal feasibility | The constraint holds | w_{\text{items}} \le 1.5 |
| Dual feasibility | The wall can only push, not pull | \mu \ge 0 |
| Complementary slackness | \mu\cdot g(w) = 0: a price only if the wall is touched | w_{\text{items}} = 1.5, so \mu may be > 0 |
Memory hook: KKT = "Some People Dislike Constraints": Stationarity, Primal, Dual, Complementary slackness. Solve: with w_{\text{items}} = 1.5, \partial L/\partial w_{\text{km}} = (2/6)(10w_{\text{km}} + 5\cdot 1.5 - 40) = 0 gives w_{\text{km}} = 3.25. Then \partial L/\partial w_{\text{items}} = (2/6)(5\cdot 3.25 + 10\cdot 1.5 - 35) = -1.25, so \mu = 1.25 \ge 0. The distance weight rose from 3 to 3.25 to absorb part of what items can no longer explain, and the loss rose from 1.667 to 1.979.
The multiplier is a price: dL^*/d(\text{cap}) = -\mu = -1.25. Relax the cap by one minute per item and the best loss falls by about 1.25 per unit (to first order). That's the number to bring to the ops meeting. (Drill 15: a budget w_{\text{km}} + w_{\text{items}} \le 4 gives w = (2.5, 1.5), \mu = 2.5.)
Projected GD: step, then snap
Take a normal gradient step, then project: move to the nearest allowed point. For a box that's just clipping each coordinate. On the toy with \eta = 0.3: from (0, 0) the step goes to (4.0, 3.5), the snap to (4.0, 1.5); the next step and snap land on (3.25, 1.5), the KKT answer, in 2 iterations. Projection works whenever it's cheap: boxes, balls, the simplex. scipy's L-BFGS-B does this for box bounds (Section 14).
Kinks: why GD never lands on zero. The lasso, (1/n)\lVert y - Xw\rVert ^{2} + \lambda\lVert w\rVert _{1}, has a kink at every w_{j} = 0. There the gradient doesn't exist; a subgradient is any slope that stays below the function (for \lvert w\rvert at 0, anything in [−1, 1]). Subgradient descent picks one, usually sign(w), and overshoots zero every time. On the toy at \lambda = 11, where the exact answer has w_{\text{items}} = 0, subgradient descent (\eta = 0.01) keeps w_{\text{items}} jittering between about −0.001 and 0.21, and in its last 1,000 steps it is never exactly zero (Q25).
Proximal GD: step, then shrink
Split the loss into a smooth part and the penalty. Step on the smooth part only, then apply the penalty's proximal operator: the point that best trades staying close against paying the penalty. For L1 it's soft-thresholding (Card 16):
- \operatorname{prox}_{t|\cdot|}(z)
- the proximal operator of t|w|: best trade between staying near z and paying the penalty
- z
- the weight after the plain gradient step, before the shrink
- \arg\min_w
- the w that makes what follows smallest
- w
- a candidate value for one weight
- t
- the shrink threshold: \eta\lambda in ISTA
- \operatorname{sign}
- the sign: +1 for positive, −1 for negative, 0 at 0
Shrink toward zero by t, and anything within t of zero becomes exactly zero: prox(3) = 2, prox(0.5) = 0, prox(−2) = −1 with t = 1. ISTA (iterative shrinkage-thresholding) alternates the two with threshold t = \eta\lambda. On the toy with \eta = 0.1 and \lambda = 4, the first step goes to (1.333, 1.167) and the shrink by 0.4 gives (0.933, 0.767); it converges to (2.2, 1.2). At \lambda = 11 it lands on (0.7, 0), an exact zero where subgradient descent jittered. Memory hook: "Projected: step, then snap. Proximal: step, then shrink."
The toy lasso path. Both weights start at (3, 2) and shrink by \lambda/5 together (because (1, 1) is the Hessian's steep eigenvector) until w_{\text{items}} hits zero at \lambda = 10; distance alone then shrinks to zero at \lambda_{\max} = 40/3 \approx 13.33 = \lVert (2/n)X^\top y\rVert _\infty, the smallest \lambda that zeros everything. L1 kills, L2 shrinks (Ch 1, Section 10), and this is the mechanism.
Faster cousins. FISTA adds Nesterov momentum to ISTA, improving its convex rate from 1/k to 1/k^{2}. Coordinate descent (what sklearn's Lasso uses) updates one weight at a time with an exact 1-D soft-threshold, which is very fast when most weights end at zero.
Weight decay is a prox too. The prox of (\lambda/2)\lVert w\rVert ^{2} with step \eta is w/(1 + \eta\lambda) \approx (1 - \eta\lambda)w: a multiplicative shrink applied after the gradient step. That's exactly AdamW's decoupled decay (Section 8): step on the loss, then shrink.
Duality, in one paragraph. Minimizing the Lagrangian over w for a fixed \mu gives a function of \mu alone, the dual, which is always concave and whose maximum lower-bounds the original minimum. For a convex problem with a strictly feasible point (Slater's condition) the two are equal, and the KKT conditions are then both necessary and sufficient. Without convexity KKT is only necessary, and only under a "constraint qualification" that rules out degenerate walls. SVMs (Ch 11) are solved in their dual, where the kernel trick appears.
Remember: A constraint becomes a price \mu. KKT = Some People Dislike Constraints: Stationarity, Primal, Dual, Complementary slackness; \mu > 0 only when the constraint binds, and \mu is the loss saved per unit of slack. Projected: step, then snap. Proximal: step, then shrink. Soft-thresholding gives exact zeros; subgradients only jitter.
Next question → We have every piece in theory. Can we build GD, SGD, momentum and Adam ourselves and watch them match the closed form? (Section 13)
13. Build It From Scratch
The question: can we write GD, minibatch SGD, momentum and Adam in about 50 lines of NumPy and land on (6, 3, 2) for the toy and the OLS answer for the 10k orders?
The 30-second version: one function computes the gradient; one loop applies any of four update rules; a gradient check verifies the gradient; all four optimizers converge to the lstsq answer on scaled data. Delete the scaling line and watch GD on raw features diverge or crawl. First read: run the code once.
Every line maps to a section: the gradient (2), standardizing (4), minibatches reshuffled each epoch (6, 10), the heavy ball (7), the first three D's of Adam (8), and the gradient check (10).
import numpy as np
def grad(w, X, y): # gradient of (1/n)||y - Xw||^2
return -2 / len(y) * X.T @ (y - X @ w)
def grad_check(w, X, y, h=1e-5): # central differences, float64
L = lambda v: np.mean((y - X @ v) ** 2)
num = np.array([(L(w + h * e) - L(w - h * e)) / (2 * h) for e in np.eye(len(w))])
ana = grad(w, X, y)
return np.linalg.norm(ana - num) / max(np.linalg.norm(ana), np.linalg.norm(num))
def fit(X, y, method="gd", eta=0.1, steps=2000, batch=None, beta=0.9,
b1=0.9, b2=0.999, eps=1e-8, debias=True, seed=0, tol=1e-10):
rng = np.random.default_rng(seed)
w = np.zeros(X.shape[1]); v = np.zeros_like(w); m = np.zeros_like(w); s = np.zeros_like(w)
n = len(y); perm, pos = rng.permutation(n), 0
for t in range(1, steps + 1):
if batch: # minibatch, reshuffled every epoch
if pos + batch > n: perm, pos = rng.permutation(n), 0
idx = perm[pos:pos + batch]; pos += batch
g = grad(w, X[idx], y[idx])
else:
g = grad(w, X, y)
if method in ("gd", "sgd"):
step = eta * g
elif method == "momentum": # heavy ball, PyTorch form
v = beta * v + g; step = eta * v
elif method == "adam": # the four D's (no decay here)
m = b1 * m + (1 - b1) * g # Direction
s = b2 * s + (1 - b2) * g * g # Denominator
mh, sh = (m / (1 - b1 ** t), s / (1 - b2 ** t)) if debias else (m, s) # Debias
step = eta * mh / (np.sqrt(sh) + eps)
w = w - step
if not batch and np.linalg.norm(step) < tol:
return w, t
return w, steps
def standardize(F):
mu, sd = F.mean(0), F.std(0)
return np.column_stack([np.ones(len(F)), (F - mu) / sd]), mu, sd
def to_raw(w, mu, sd): # back to minutes per km / per item
slopes = w[1:] / sd
return np.concatenate([[w[0] - slopes @ mu], slopes])
# ---- run it: the six orders, then the 10,000 ----
F = np.array([[1, 2], [2, 1], [3, 5], [3, 3], [4, 4], [5, 3.]]); y = np.array([15, 12, 24, 21, 26, 28.])
X, mu, sd = standardize(F)
print("lstsq ", to_raw(np.linalg.lstsq(X, y, rcond=None)[0], mu, sd).round(4))
for name, kw in [("gd", dict(eta=0.3)), ("momentum", dict(eta=0.1, beta=0.5)),
("adam", dict(eta=0.05, steps=5000)), ("sgd", dict(eta=0.05, steps=3000, batch=2))]:
w, k = fit(X, y, name, **kw)
print(f"{name:9s}", to_raw(w, mu, sd).round(4), "" if name == "sgd" else f"({k} steps)")
print("grad check", f"{grad_check(np.array([1.0, -0.5, 0.3]), X, y):.1e}")
rng = np.random.default_rng(2002) # the 10k orders (Section 1)
km = np.clip(rng.gamma(2.0, 1.6, 10_000), 0.2, 15)
items = np.minimum(1 + rng.poisson(1.2 + 0.25 * km), 12).astype(float)
eta_min = 6 + 3 * km + 2 * items + rng.normal(0, 3.0, 10_000)
X, mu, sd = standardize(np.column_stack([km, items]))
print("10k lstsq", to_raw(np.linalg.lstsq(X, eta_min, rcond=None)[0], mu, sd).round(4))
w, k = fit(X, eta_min, "gd", eta=0.5); print("10k gd ", to_raw(w, mu, sd).round(4), f"({k} steps)")
w, _ = fit(X, eta_min, "sgd", eta=0.01, steps=20000, batch=32); print("10k sgd ", to_raw(w, mu, sd).round(4))
What it prints, and what to notice:
- On the toy,
lstsq, GD (62 steps to a step norm below 10^{-10}), momentum (75) and Adam (2,338) all give (6, 3, 2). Adam is the slowest here: near the minimum its steps stay about \eta in size until \hat v shrinks, so it creeps in at the end. Adam's strength is robustness to scale, not speed on a well-conditioned bowl. - Minibatch SGD with B = 2 ends near (5.92, 3.00, 2.02): within its noise floor, not on the answer (Section 6).
- The gradient check scores 5.5 \times 10^{-11}.
- On the 10k orders, GD at \eta = 0.5 (exactly 2/(\lambda_{\max} + \lambda_{\min}) for the standardized bowl) reaches the OLS answer (6.05, 2.99, 1.99) in 27 steps; SGD with B = 32 lands near (6.10, 2.99, 1.98).
Common mistake: no scaling. On raw km and items the cliff is 2/49.93 = 0.040 (Section 9), so \eta = 0.5 has blown up within 50 steps, and \eta = 0.01 is still at (6.046, 2.9945, 1.9946) after 2,000. A second one: forgetting to reshuffle. One epoch in the warehouse's stored order ends at (27.7, 1.40, 2.62) instead of (6.05, 2.99, 1.99): the four S's (Section 10) in one line.
Try it: set debias=False and run one Adam step on the 10k data. Each weight moves 0.158 instead of \eta = 0.05: the 3.16× overshoot of Section 8.
Remember: If you can write this from memory, you understand this chapter: one gradient, one loop, four update rules, one gradient check. \text{Step} = -\eta \times \text{slope}, plus memory (momentum) and a per-weight denominator (Adam).
Next question → How does this look in PyTorch and scikit-learn, and which defaults bite? (Section 14)
14. Use It Like a Pro
The question: which optimizer, schedule and settings do practitioners actually reach for, and which library defaults cause silent bugs?
The 30-second version: the PyTorch loop is zero-grad, forward, backward, clip, step, schedule. Use AdamW (not Adam with weight_decay) for deep nets, SGD + momentum for some vision models, L-BFGS for small full-batch convex problems. Five gotchas: Adam(weight_decay=…) is L2, not decoupled; schedulers stepped per epoch vs per iteration; clipping must come after backward and before step (and after unscaling under AMP); sklearn's SGDRegressor is regularized and decays its learning rate by default; sklearn's LogisticRegression stops at max_iter = 100 with only a warning. First read: the training-loop hook, the gotchas table.
The six beats of a training step
Memory hook: "Zero, Forward, Backward, Clip, Step, Schedule." Here they are on the 10k orders with a small network, AdamW, warmup + cosine counted in steps, and clipping (the block skips itself if PyTorch isn't installed):
import importlib.util
import math
def train_eta_model():
import numpy as np
import torch
torch.manual_seed(0)
rng = np.random.default_rng(2002) # the 10k orders (Section 1)
km = np.clip(rng.gamma(2.0, 1.6, 10_000), 0.2, 15)
items = np.minimum(1 + rng.poisson(1.2 + 0.25 * km), 12)
y = 6 + 3 * km + 2 * items + rng.normal(0, 3.0, 10_000)
F = np.column_stack([km, items])
X = torch.tensor((F - F.mean(0)) / F.std(0), dtype=torch.float32) # scale first (A in SCUBA)
Y = torch.tensor(y, dtype=torch.float32).unsqueeze(1)
model = torch.nn.Sequential(torch.nn.Linear(2, 16), torch.nn.ReLU(), torch.nn.Linear(16, 1))
decay = [p for p in model.parameters() if p.ndim > 1] # decay weights, not biases
no_decay = [p for p in model.parameters() if p.ndim <= 1]
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
{"params": no_decay, "weight_decay": 0.0}], lr=3e-3, betas=(0.9, 0.95))
total, warmup = 3000, 150 # schedule counted in STEPS, not epochs
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda t: (t + 1) / warmup if t < warmup else
0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * (t - warmup) / (total - warmup))))
loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(X, Y), batch_size=64, shuffle=True)
step = 0
while step < total:
for xb, yb in loader: # shuffle=True: a fresh order every epoch
opt.zero_grad(set_to_none=True) # Zero
loss = torch.nn.functional.mse_loss(model(xb), yb) # Forward
loss.backward() # Backward
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # Clip
opt.step() # Step
sched.step() # Schedule (per iteration)
step += 1
if step == total:
break
with torch.no_grad():
print("train MSE:", round(torch.nn.functional.mse_loss(model(X), Y).item(), 2))
if importlib.util.find_spec("torch") is None:
print("PyTorch is not installed; skipping the training loop.")
else:
train_eta_model()
Mixed precision. With fp16, wrap the forward pass in torch.autocast and use a GradScaler: it multiplies the loss so small gradients don't underflow, and the order becomes scaler.scale(loss).backward() → scaler.unscale_(opt) → clip → scaler.step(opt) → scaler.update(). Clip after unscaling or the threshold is meaningless. bf16 has fp32's exponent range, so it needs no scaler; it's the usual choice on hardware that supports it (TPUs, recent GPUs). **torch.optim.LBFGS** needs a closure that recomputes the loss, because it evaluates the function several times per step.
scikit-learn and SciPy. scipy.optimize.minimize(..., method="L-BFGS-B", bounds=...) handles box constraints by projection (Section 12): with w_{\text{items}} capped at 1.5 on the toy it returns (3.25, 1.5), the KKT answer. sklearn's estimators carry defaults that are not plain least squares (table below).
Defaults that bite
| Setting | Default | What bites | Fix |
|---|---|---|---|
torch.optim.Adam(weight_decay=λ) | L2 added to the gradient | Penalty divided by \sqrt{\hat v} (Section 8) | AdamW; exclude biases and norm weights from decay |
scheduler.step() placement | Whatever your loop does | A per-iteration schedule stepped per epoch barely moves; warmup never ends | Count schedules in steps; step once per optimizer step |
clip_grad_norm_ placement | — | Before backward: clips nothing; under fp16 before unscale_: wrong threshold | After backward, after unscale, before step |
SGDRegressor() | penalty='l2', alpha=1e-4, learning_rate='invscaling', eta0=0.01, power_t=0.25 | It's ridge with a decaying step, and unscaled features wreck it | StandardScaler in a pipeline; set penalty and schedule on purpose |
LogisticRegression() | L2 with C = 1.0, lbfgs, max_iter=100 | Penalized coefficients ((−4.97, 4.03, 1.75) vs MLE (−5.05, 4.11, 1.78) here); a ConvergenceWarning you might not see | Scale; raise max_iter; C=np.inf for no penalty |
DataLoader(shuffle=False) | No shuffling | Sorted data → biased gradients (Section 10) | shuffle=True for training |
| Adam \varepsilon = 1e-8 | — | Huge steps for weights with tiny \hat v in low precision | Raise \varepsilon (e.g. 1e-6) |
Before you launch a training run: features centered and scaled (A); data shuffled every epoch, batches stratified if classes are rare (U); gradient checked or autodiff trusted; one batch overfit; LR range test done; schedule counted in steps with warmup; clipping on and gradient norms logged; loss, learning rate and validation curve all plotted; a checkpoint to rewind to.
Remember: Zero, Forward, Backward, Clip, Step, Schedule. AdamW, not Adam's weight_decay. Schedules in steps. Clip after backward and after unscaling. sklearn's SGDRegressor and LogisticRegression are regularized by default; read the defaults.
Next question → Where do these choices matter most in the real world, and what is the recipe that trains an LLM? (Section 15)
15. In the Real World, and How It Connects to LLMs
The question: which optimizers run real systems, and what is the optimization recipe behind a large language model?
The 30-second version: convex models in production (ads CTR, credit, pricing) use L-BFGS, coordinate descent or SGD at scale; deep nets use SGD + momentum or AdamW. The LLM recipe is AdamW with \beta_{2} \approx 0.95, linear warmup, cosine or warmup-stable-decay to a small floor, weight decay ≈ 0.1, gradient clipping at 1.0, bf16 mixed precision. Loss spikes are the signature LLM failure, and the fixes are this chapter's tools: lower peak \eta, longer warmup, clip, skip bad batches, rewind. First read: the domain table and the LLM-recipe table.
Where optimization choices earn their keep
| Domain | Optimizer in practice | Why |
|---|---|---|
| Credit scoring, insurance pricing | IRLS / Newton, L-BFGS | Small convex models; exact, reproducible fits for regulators |
| Ads and recommendation CTR | SGD / AdaGrad / FTRL on sparse features | Billions of rows, rare features need larger steps |
| Delivery ETAs, demand forecasting | Gradient-boosted trees (Ch 10), small nets with Adam | Nonlinear tabular data |
| Operations constraints (SLAs, budgets, fairness caps) | Projected or constrained solvers, Lagrangian penalties | KKT and shadow prices are the business answer |
| LLM pretraining and fine-tuning | AdamW + schedule + clipping | The recipe below |
The LLM recipe
GPT-3 (Brown et al., 2020) and LLaMA (Touvron et al., 2023) published essentially the same optimization recipe. Memory hook: "Always Warm up, Cool down, Clip, in bf16, with AdamW."
| Ingredient | Typical value | Why (section) |
|---|---|---|
| AdamW, \beta_{1} | 0.9 | Momentum window ≈ 10 steps (7, 8) |
| \beta_{2} | 0.95 (not 0.999) | Shorter memory for \hat v: adapts faster when gradient scale jumps, fewer spikes (8) |
| Weight decay (decoupled) | 0.1 | Clean shrink, independent of \hat v (8, 12) |
| Linear warmup | e.g. 2,000 steps (LLaMA) | Early \hat v unreliable; early curvature high (6, 8) |
| Decay | Cosine to 10% of the peak, or WSD | Removes the noise floor (6) |
| Gradient clipping | Global norm 1.0 | Insurance against rare huge gradients (10) |
| Precision | bf16 mixed precision, fp32 master weights | fp32's range at half the memory (14) |
Loss spikes. The signature failure of large runs is a spike: the loss jumps (say from 2.1 to 7) and partly recovers (Q30). The usual suspects are the chapter's: a peak \eta near the edge of stability (11); Adam's \hat v tiny for some weights after a long quiet stretch, so a sudden large gradient produces a huge step (8, the reason for \beta_{2} = 0.95 and larger \varepsilon); a bad batch (10); fp16 overflow (14). The fixes are the chapter's too: lower peak \eta or longer warmup, clip, skip the offending batches, and rewind to a checkpoint before the spike, which PaLM (Chowdhery et al., 2022) did by restarting a little before each spike and skipping a few hundred batches.
Scale changes the optimizer's problems
- Data parallelism is a bigger batch: k \text{ GPUs} \times B each. That means the linear scaling rule with warmup, up to the critical batch size (6, Q28).
- Gradient accumulation sums gradients over several small batches before one step: the same math as a big batch, with less memory.
- Asynchronous updates apply stale gradients, which acts like extra momentum and destabilizes training unless \eta falls (Q39). Most large runs use synchronous all-reduce.
- Optimizer memory. Adam stores two extra numbers per weight, 8 bytes per parameter for fp32 moments: twice the fp32 weights themselves. That's why Adafactor, 8-bit optimizers and sharded optimizer states (ZeRO, Ch 23) exist.
How this chapter becomes an LLM
| In this chapter | Becomes | Where in the book |
|---|---|---|
| Gradient of the MSE on six orders | Backpropagation through billions of weights | Ch 16 |
| \kappa and feature scaling | LayerNorm, careful initialization, Adam's per-weight scaling | Ch 17, 21 |
| Minibatch noise, \eta/B | Batch-size and learning-rate scaling, critical batch size | Ch 23 |
| Warm up, hold, cool down | Warmup + cosine or WSD schedules | Ch 22 |
| AdamW's four D's | The default LLM optimizer | Ch 22–24 |
| Flat vs sharp, implicit regularization | Why overparameterized models generalize | Ch 4, 23 |
Remember: Convex production models use L-BFGS, IRLS or coordinate descent; deep nets use SGD + momentum or AdamW. The LLM recipe: Always Warm up, Cool down, Clip, in bf16, with AdamW (\beta_{2} = 0.95, decay 0.1, clip 1.0, cosine to 10%). Spikes are fixed with this chapter's tools.
Next → The whole chapter on one page (Section 16). Then practice: the Derivations tab, then the Interview Bank. And next chapter: the ops team now wants P(late), a yes/no answer, not minutes. The loss becomes cross-entropy, the closed form is gone for good, and Newton becomes IRLS; everything in this chapter is how we'll fit it. (Ch 3)
16. The Chapter on One Page
One map of the whole chapter and every memory hook in one table. Use it before an interview, or to check a second read: if any row feels unfamiliar, go back to its section.
The memory map
Every hook in one place
| Hook | What it unlocks | Section |
|---|---|---|
| F-L-O (Family, Loss, Optimizer) | The three choices of every model; this chapter is the O | 1 |
| Slope says which way, curvature says how far | Gradient vs Hessian; the Taylor picture | 2 |
| \text{Step} = -\eta \times \text{slope} (Ch 1) | Gradient descent | 2 |
| Chord above, tangent below, curvature up | Three convexity tests | 3 |
| SAM (non-negative Sums, Affine inside, Max) | Building convex losses | 3 |
| 1/\lambda smooth, 2/\lambda the cliff | The four learning-rate regimes | 4 |
| Center, then scale, then worry about correlation | Conditioning; \kappa = (1 + r)/(1 - r) | 4 |
| Sublinear crawls, linear adds digits, quadratic doubles them | Convergence rates | 5 |
| Unbiased, cheap, noisy | The minibatch gradient | 6 |
| Warm up, hold, cool down | Every modern schedule | 6 |
| \beta is memory: window 1/(1 - \beta), step \eta/(1 - \beta) | Momentum and every EMA | 7 |
| Direction, Denominator, Debias, Decay | AdamW's four parts | 8 |
| Too big, too noisy, too saddle-happy | Why deep learning skips Newton | 9 |
| SCUBA (Smooth, Convex, Unbiased, Bounded noise, Aligned scales) | The five assumptions behind every guarantee | 10 |
| NaN, Noise, Nothing, Notches | Reading a bad loss curve | 10 |
| Shuffle, Stratify, Sample by importance, Scale | Fixing a biased gradient | 10 |
| Signs of the Hessian: all up, all down, mixed | Minimum, maximum, saddle | 11 |
| Big steps can't sit in sharp holes | Stability selects flat minima | 11 |
| Some People Dislike Constraints (KKT) | Stationarity, Primal, Dual, Complementary slackness | 12 |
| Projected: step, then snap. Proximal: step, then shrink. | Constrained and lasso updates | 12 |
| L1 kills, L2 shrinks (Ch 1) | Why the prox gives exact zeros | 12 |
| Zero, Forward, Backward, Clip, Step, Schedule | A PyTorch training step | 14 |
| Always Warm up, Cool down, Clip, in bf16, with AdamW | The LLM recipe | 15 |
Remember: if you can redraw the memory map on a blank page and say the hook for every row of the table, you have the chapter. Now practice.
17. Practice
The chapter ends here. Practice lives in two tabs, so this page stays a story you can read straight through.
- Ch 2 · Derivations — 20 derivation cards built on four moves, TRUE (Taylor, Rotate to the eigenbasis, Unroll the recursion, Expectation), each with its key move and common slip; a formula sheet; and a mastery tracker.
- Ch 2 · Interview Bank — 39 questions across Levels 1–3, each answered in four beats; a Level 4 with two worked research-level problems with figures (the run that diverged at \eta = 0.05, and SGD's noise floor) and 20 drills solved with one recipe, Rotate, Recurse, Read off; and rapid-fire flashcards for the day before.
Suggested order: finish the chapter's second read → Derivations Cards 1–8 → Interview Bank Levels 1–2 → Cards 9–20 → Levels 3–4.