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.

Figure 2.0 · Chapter map · 15 questions, 3 acts
Figure 2.0 · Chapter map · 15 questions, 3 acts

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:

OrderDistance (m, as logged)Distance (km)ItemsETA (min)
A1,0001215
B2,0002112
C3,0003524
D3,0003321
E4,0004426
F5,0005328

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(\mathbf w) = \frac{1}{n}\lVert \mathbf y - X\mathbf w\rVert^2, \qquad \nabla L = -\frac{2}{n}X^\top(\mathbf y - X\mathbf w), \qquad H = \frac{2}{n}X^\top X
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

ChoiceQuestion it answersETA modelWhere it's owned
FamilyWhat shapes can the prediction take?A weighted sum of distance and itemsCh 1, and later model chapters
LossWhat does "wrong" cost?Squared error in minutesCh 1 (MSE), Ch 3 (cross-entropy)
OptimizerHow do we find the best member of the family?Gradient descent and its relativesThis 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} = \arg\min_{\mathbf w}\; \frac{1}{n}\sum_{i=1}^{n} \ell\big(y_i, f(\mathbf x_i; \mathbf w)\big)
\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.

Figure 2.1a · The loss is a bowl: contours of the six-order loss over (w_km, w_items) with −∇L arrows (computed from the six-order toy)
Figure 2.1a · The loss is a bowl: contours of the six-order loss over (w_km, w_items) with −∇L arrows (computed from the six-order toy)

Vocabulary you'll be asked about

TermMeaningOn the ETA bowl
Global minimumThe lowest point anywhere(3, 2), loss 1.667
Local minimumLower than everything nearbySame point: a bowl has only one
Stationary (critical) pointWhere the gradient is zeroSame point
Saddle pointStationary, but downhill in some directionNone here; common in neural nets (Section 11)
PlateauA region where the gradient is nearly zero but the loss is highRaw 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:

ErrorCaused byShrinks with
ApproximationThe family can't represent the truthA richer family (F)
EstimationWe fit on a finite sample, not the populationMore data
OptimizationWe stopped before the exact minimumMore 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(\mathbf 0) = -\frac{2}{6}\begin{pmatrix}\sum_i d_{1i}\,\tilde y_i \\ \sum_i d_{2i}\,\tilde y_i\end{pmatrix} = -\frac{2}{6}\begin{pmatrix}40 \\ 35\end{pmatrix} = \begin{pmatrix}-13.33 \\ -11.67\end{pmatrix}, \qquad \lVert\nabla L\rVert = 17.72
\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(\mathbf w + \Delta) \approx L(\mathbf w) + \nabla L^\top \Delta + \tfrac12\,\Delta^\top H\,\Delta, \qquad H_{jk} = \frac{\partial^2 L}{\partial w_j\,\partial w_k}
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) \approx L(\mathbf w) \;-\; \underbrace{\eta\,\lVert\nabla L\rVert^2}_{\text{slope: always a gain}} \;+\; \underbrace{\tfrac12\,\eta^2\,\nabla L^\top H\,\nabla L}_{\text{curvature: the price of a long step}}
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.

Figure 2.1b · The Taylor picture: the loss along −∇L from w = 0, with the linear and quadratic models and the step η = 0.1 (computed from the six-order toy)
Figure 2.1b · The Taylor picture: the loss along −∇L from w = 0, with the linear and quadratic models and the step η = 0.1 (computed from the six-order toy)

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:

H = \frac{2}{6}\begin{pmatrix}10 & 5\\ 5 & 10\end{pmatrix} = \begin{pmatrix}3.333 & 1.667\\ 1.667 & 3.333\end{pmatrix}

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 - \eta\nabla L) \;\le\; L(\mathbf w) - \eta\Big(1 - \frac{L\eta}{2}\Big)\lVert\nabla L\rVert^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

\text{Chord: } L(\theta\mathbf a + (1-\theta)\mathbf b) \le \theta L(\mathbf a) + (1-\theta)L(\mathbf b),\;\; \theta\in[0,1]
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
\text{Tangent: } L(\mathbf b) \ge L(\mathbf a) + \nabla L(\mathbf a)^\top(\mathbf b - \mathbf a) \qquad \text{Curvature: } H(\mathbf w) \succeq 0 \;\;\text{for all } \mathbf w
\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 I \preceq H(\mathbf w) \preceq L\, I
\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 modelConvex?Smooth?Unique minimum?Notes
MSEYesYesIf X has full rankHessian (2/n)X^\top X
MAEYesNo (kink at 0)Not alwaysNeeds subgradients (Section 12)
HuberYesOnce-differentiableUsuallySmooth core, linear tails
Logistic (cross-entropy)YesYesUnless the data are separableHessian (1/n)X^\top SX, S = \operatorname{diag}(p(1 - p)) (Card 3, Ch 3)
Hinge (SVM)YesNoWith the L2 term, yesCh 11
MSE on a sigmoid outputNoYesNoWhy Ch 3 uses cross-entropy
Any neural networkNoMostlyNoSymmetries 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_{k+1} = \mathbf e_k - \eta H \mathbf e_k = (I - \eta H)\,\mathbf e_k \quad\Rightarrow\quad \mathbf e_k = (I - \eta H)^k\,\mathbf e_0
\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} = (1 - \eta\lambda_i)^k\, c_{i,0}
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:

0 < \eta < \frac{2}{\lambda_{\max}}
\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:

\etaSteep-mode factor 1 - 5\etaFlat-mode factor 1 - (5/3)\etaSteps to 10^{-6}What you see
0.050.750.91755Smooth, slow
0.10.50.83327Smooth
0.2 = 1/\lambda_{\max}00.66712Steep mode gone in one step
0.3 = 2/(\lambda_{\max} + \lambda_{\min})−0.50.510Best fixed step: slight zigzag
0.35−0.750.41724Zigzag
0.39−0.950.35135Violent zigzag, barely converging
0.41−1.050.317DivergesEach 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).

Figure 2.2 · Four step sizes on the same bowl: GD paths for η = 0.1, 0.3, 0.39, 0.41, with an inset of steps-to-tolerance vs η and the cliff at 0.4 (computed from the six-order toy)
Figure 2.2 · Four step sizes on the same bowl: GD paths for η = 0.1, 0.3, 0.39, 0.41, with an inset of steps-to-tolerance vs η and the cliff at 0.4 (computed from the six-order toy)

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^* = \frac{2}{\lambda_{\max} + \lambda_{\min}}, \qquad \rho = \max_i |1 - \eta^*\lambda_i| = \frac{\kappa - 1}{\kappa + 1}, \qquad \kappa = \frac{\lambda_{\max}}{\lambda_{\min}}
\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)30.43.02 (2.15 standardized)
km and items, uncentered, with an intercept1830.0468143
Meters and items, centered1.33 \times 10^{6}6.0 \times 10^{-7}—
Meters and items, uncentered, with an intercept8.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:
\kappa = \frac{1 + r}{1 - r} \qquad\text{toy: } r = 0.5 \;\Rightarrow\; \kappa = \frac{1.5}{0.5} = 3
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:

  1. 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.
  2. 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.
  3. 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 \le \rho^k\lVert \mathbf e_0\rVert \quad\Rightarrow\quad k \ge \frac{\ln(1/\varepsilon)}{\ln(1/\rho)}, \qquad \ln\frac{1}{\rho} \ge \frac{2}{\kappa} \;\Rightarrow\; k = \frac{\kappa}{2}\ln\frac{1}{\varepsilon} \text{ steps always suffice}
\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})TheorySimulated
Distance \lVert w - w^*\rVert(\kappa/2)\cdot \ln 10^{6} \approx 990.6991
Loss gap L - L^*half that, ≈ 495.3496

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.

Figure 2.3 · How fast: log₁₀ relative suboptimality vs step for GD (η = 0.1, 0.3), heavy ball and Newton on the toy, the raw 10k model, and sublinear GD on f(w) = w⁴ (computed from the six-order toy and the 10k simulation)
Figure 2.3 · How fast: log₁₀ relative suboptimality vs step for GD (η = 0.1, 0.3), heavy ball and Newton on the toy, the raw 10k model, and sublinear GD on f(w) = w⁴ (computed from the six-order toy and the 10k simulation)

The rates table

MethodSettingRateSteps for accuracy \varepsilon
GDStrongly convex, smoothLinear, \rho = (\kappa - 1)/(\kappa + 1)\approx \kappa\cdot \log(1/\varepsilon)
GDConvex, smooth (\eta = 1/L)L(w_{k}) - L^* \le L\lVert w_{0} - w^*\rVert ^{2}/(2k)\approx 1/\varepsilon
Subgradient methodConvex, non-smooth (MAE, hinge)1/\sqrt{k} with decaying \eta\approx 1/\varepsilon^{2}
SGD, decaying \etaStrongly convex1/k\approx 1/\varepsilon
SGD, decaying \etaConvex1/\sqrt{k}\approx 1/\varepsilon^{2}
Heavy ball, tunedStrongly convex quadraticLinear, \rho = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1)\approx \sqrt{\kappa}\cdot \log(1/\varepsilon)
NesterovStrongly convex, smoothLinear, \rho \approx 1 - 1/\sqrt{\kappa}\approx \sqrt{\kappa}\cdot \log(1/\varepsilon)
NesterovConvex, smooth1/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
NewtonNear a non-degenerate minimumQuadratic\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:

\tfrac12\lVert\nabla L(\mathbf w)\rVert^2 \;\ge\; \mu\,\big(L(\mathbf w) - L^*\big) \quad\text{for all } \mathbf w
\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 = \frac1B\sum_{i\in\mathcal B}\nabla\ell_i(\mathbf w), \qquad \mathbb E[\mathbf g_B] = \mathbf G, \qquad \mathbb E\lVert \mathbf g_B - \mathbf G\rVert^2 = \frac{\operatorname{tr}\Sigma}{B}\cdot\frac{N - B}{N - 1}
\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 BMeasured E\lVert g_B - G\rVert ^{2}Theory \operatorname{tr} \Sigma/B \cdot (N - B)/(N - 1)
18,5848,330
89671,041
32275260
12867.564.3
1,0247.077.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."

Figure 2.4a · Minibatch gradient variance vs batch size on log-log axes: simulated dots on the tr Σ/B line (computed from the 10k simulation)
Figure 2.4a · Minibatch gradient variance vs batch size on log-log axes: simulated dots on the tr Σ/B line (computed from the 10k simulation)

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+1} = (1 - \eta\lambda)\,e_k - \eta\,\xi_k \quad\Rightarrow\quad \mathbb E[L - L^*] = \frac{\eta\,\sigma^2}{2B\,(2 - \eta\lambda)} \;\approx\; \frac{\eta\,\sigma^2}{4B}
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:

\etaBSimulated floorTheory (\eta/4B)\cdot \operatorname{tr} \Sigma
0.0110.2870.273
0.01320.00920.0085
0.180.4070.341
0.118.242.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}\eta_t = \infty \quad(\text{you can still travel any distance}), \qquad \sum_{t=1}^{\infty}\eta_t^2 < \infty \quad(\text{the total noise injected is finite})
\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:

ScheduleShape of \eta(t)Final excess loss
ConstantFlat at the peak0.0373
Step×0.1 at 50% and 75% of training1.00 \times 10^{-4}
CosineHalf a cosine from peak to 01.58 \times 10^{-4}
Linear warmup + cosineRamp up over 5%, then cosine1.65 \times 10^{-4}
Warmup-stable-decay (WSD)Ramp, hold flat to 80%, linear decay to 02.57 \times 10^{-3}
1/tpeak/(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."

Figure 2.4b · Schedules: (b) excess training loss vs step for constant, step, cosine, warmup + cosine, WSD and 1/t, plus a batch-size panel at equal examples seen; (c) the learning-rate shapes (computed from the 10k simulation)
Figure 2.4b · Schedules: (b) excess training loss vs step for constant, step, cosine, warmup + cosine, WSD and 1/t, plus a batch-size panel at equal examples seen; (c) the learning-rate shapes (computed from the 10k simulation)

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:

\text{PyTorch form: }\; \mathbf v_{k+1} = \beta\,\mathbf v_k + \nabla L(\mathbf w_k), \qquad \mathbf w_{k+1} = \mathbf w_k - \eta\,\mathbf v_{k+1}
\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
\text{Classical form: }\; \mathbf w_{k+1} = \mathbf w_k - \eta\,\nabla L(\mathbf w_k) + \beta\,(\mathbf w_k - \mathbf w_{k-1})
(\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):

\mathbf v_k = \nabla L_k + \beta\,\nabla L_{k-1} + \beta^2\,\nabla L_{k-2} + \dots \qquad\Rightarrow\qquad \text{constant } \mathbf g:\;\; \mathbf v \to \frac{\mathbf g}{1 - \beta}
\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^* = \frac{4}{(\sqrt{\lambda_{\max}} + \sqrt{\lambda_{\min}})^2}, \qquad \beta^* = \Big(\frac{\sqrt\kappa - 1}{\sqrt\kappa + 1}\Big)^2, \qquad \rho = \sqrt{\beta^*} = \frac{\sqrt\kappa - 1}{\sqrt\kappa + 1}
\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 = \mathbf w_k + \beta\,(\mathbf w_k - \mathbf w_{k-1}), \qquad \mathbf w_{k+1} = \mathbf u_k - \eta\,\nabla L(\mathbf u_k)
\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.

Figure 2.5 · The optimizer race in the detour valley: GD, heavy ball, Nesterov, Adam and Newton paths on the contours, with steps-to-tolerance (computed from the detour-valley toy)
Figure 2.5 · The optimizer race in the detour valley: GD, heavy ball, Nesterov, Adam and Newton paths on the contours, with steps-to-tolerance (computed from the detour-valley toy)

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

MethodPer-weight update (g = gradient)What it fixesWhat it breaks
SGD + momentumw -= \eta\cdot v, v = \beta v + gZigzag (Section 7)One \eta for all weights
AdaGrad (2011)w -= \eta\cdot g/\sqrt{\Sigma g^{2}}Rare features get big stepsSteps shrink forever: \eta/\sqrt{t}
RMSProp (2012)w -= \eta\cdot g/\sqrt{\text{EMA of } g^{2}}AdaGrad's dying stepNo momentum, biased early
Adam (2015)w -= \eta\cdot \hat m/(\sqrt{\hat v} + \varepsilon)Both, plus bias correctionL2 penalty gets rescaled
AdamW (2019)Adam step, plus w -= \eta\lambda wDecay independent of gradient scaleStill 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

\begin{aligned} \mathbf m_t &= \beta_1\mathbf m_{t-1} + (1-\beta_1)\,\mathbf g_t && \textbf{D}\text{irection (momentum)}\\ \mathbf v_t &= \beta_2\mathbf v_{t-1} + (1-\beta_2)\,\mathbf g_t^2 && \textbf{D}\text{enominator (gradient scale)}\\ \hat{\mathbf m}_t &= \frac{\mathbf m_t}{1-\beta_1^t},\quad \hat{\mathbf v}_t = \frac{\mathbf v_t}{1-\beta_2^t} && \textbf{D}\text{ebias}\\ \mathbf w_t &= \mathbf w_{t-1} - \eta\Big(\frac{\hat{\mathbf m}_t}{\sqrt{\hat{\mathbf v}_t}+\epsilon} + \lambda\,\mathbf w_{t-1}\Big) && \textbf{D}\text{ecay (the W, decoupled)} \end{aligned}
\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[\mathbf m_t] = (1-\beta_1)\sum_{j=0}^{t-1}\beta_1^{\,j}\,\mathbb E[\mathbf g] = (1-\beta_1^t)\,\mathbb E[\mathbf g] \quad\Rightarrow\quad \mathbb E[\hat{\mathbf m}_t] = \mathbb E[\mathbf g]
\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):

\mathbf w \leftarrow \mathbf w - \eta(\nabla L + \lambda\mathbf w) = (1 - \eta\lambda)\,\mathbf w - \eta\,\nabla L
\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:

  1. 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.
  2. 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 = -H^{-1}\nabla L \qquad\Rightarrow\qquad \mathbf w_{k+1} = \mathbf w_k - H^{-1}\,\nabla L(\mathbf w_k)
\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 \approx \tfrac{2}{n}J^\top J \quad(\text{always PSD}), \qquad \Delta = -(J^\top J)^{-1}J^\top\mathbf r
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):

MethodFeaturesIterationsWhy
Our IRLSStandardized10 to \lVert \Delta w\rVert < 10^{-8} (7 to relative suboptimality 10^{-6})Quadratic convergence
statsmodels Logit (Newton) / GLM (IRLS)Standardized10 / 9 (Newton: 10 on raw features too)Same algorithm; affine invariant
GD at \eta = 1/LStandardized553Hessian \kappa = 27.5 at the optimum
GD at the best fixed \eta (14)Standardized70Tuned step
GD at \eta = 1/LRaw km and items49,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

  1. 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.
  2. Too noisy. A minibatch Hessian is a far noisier estimate than a minibatch gradient, and inverting noise amplifies it.
  3. 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

LetterAssumptionIn plain wordsIf it breaks…
SSmooth (L-Lipschitz gradient)Curvature is bounded, so a small enough step is safeNo safe \eta: kinks (Section 12), exploding gradients
CConvexOne bottomMany stationary points; the start matters (Section 11)
UUnbiased gradientEach minibatch is a fair random sampleSGD converges to the wrong point, or cycles
BBounded noise varianceGradient noise is finiteHeavy tails (storms, bad batches) → spikes
AAligned scales (small \kappa)Features centered and scaledCrawl 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

  1. Shuffle every epoch (not once).
  2. Stratify batches so each carries the class or segment mix of the whole.
  3. Sample by importance, then reweight. Draw example i with probability p_{i} and scale its gradient so the estimate stays unbiased:
\mathbf g = \frac{1}{N\,p_i}\,\nabla\ell_i, \qquad \mathbb E[\mathbf g] = \sum_i p_i\,\frac{\nabla\ell_i}{N p_i} = \frac1N\sum_i \nabla\ell_i = \nabla L, \qquad p_i^{\text{opt}} \propto \lVert\nabla\ell_i\rVert
\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).

  1. 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

ShapeNameReal 10k run that made itUsual cause
Shoots up, then NaNNaNRaw 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^*NoiseStandardized, 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 highNothingRaw 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 spikesNotchesStored (sorted) order, no shuffle: 14.7 ↔ 201 each epochOrdered data (U); bad batches; loss spikes

Memory hook: NaN, Noise, Nothing, Notches.

Figure 2.6 · The four N's: NaN, Noise, Nothing and Notches, four small multiples from real runs on the 10k orders, each labeled with its cause (computed from the 10k simulation)
Figure 2.6 · The four N's: NaN, Noise, Nothing and Notches, four small multiples from real runs on the 10k orders, each labeled with its cause (computed from the 10k simulation)

Test it

TestWhat it checksPass mark
Gradient check (central differences)The gradient code (S, and your bugs)Relative error \lesssim 10^{-6} in float64
Overfit one batchModel, loss and optimizer can fit anything at allTraining loss → ≈ 0 on 1–10 examples
LR range testWhere the cliff is (A, S)A clear blow-up point; pick well below it
Shuffle on/off A/BOrdering bias (U)Same final loss either way
Gradient norms per layerExplosions, 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

SymptomLikely causeFix
NaN or loss explodes in the first steps\eta above the cliff; unscaled inputsLower \eta; standardize; warmup
NaN after many stepsfp16 overflow; a bad batch; log(0)bf16 or loss scaling; clip; check data; stable log-sum-exp
Plateau far above a baselineIll-conditioning; dead units; \eta far too smallScale features; Adam; LR range test
Plateau just above the best lossSGD noise floorDecay \eta; bigger batch late; average iterates
Sawtooth with the epoch periodOrdered dataShuffle every epoch
Rare large spikesHeavy-tailed examples; Adam with tiny \hat vClip gradients; skip the batch; lower \eta or \beta_{2}
Train loss falls, validation risesOverfitting, not optimizationRegularize, 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| = (1 + \eta|\lambda|)^k\,|b_0| \quad\Rightarrow\quad k \approx \frac{\ln(|b_{\text{target}}|/|b_0|)}{\ln(1 + \eta|\lambda|)}
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.

Figure 2.7 · Non-convex landscapes: (a) the double well with GD endpoints from five starts; (b) escaping a saddle, GD vs noisy GD; (c) flat vs sharp minima with final-position histograms for η = 0.01 and 0.1 (computed from the non-convex toys)
Figure 2.7 · Non-convex landscapes: (a) the double well with GD endpoints from five starts; (b) escaping a saddle, GD vs noisy GD; (c) flat vs sharp minima with final-position histograms for η = 0.01 and 0.1 (computed from the non-convex toys)

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:

\etaSettles in the sharp holeEnds in the flat basinTrapped 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) = L(\mathbf w) + \mu\,g(\mathbf w), \qquad \mu \ge 0
\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):

ConditionMeaningOn 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 feasibilityThe constraint holdsw_{\text{items}} \le 1.5
Dual feasibilityThe wall can only push, not pull\mu \ge 0
Complementary slackness\mu\cdot g(w) = 0: a price only if the wall is touchedw_{\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).

Figure 2.8 · Constraints: (a) the toy contours with the cap w_items ≤ 1.5, the projected-GD path, the constrained optimum (3.25, 1.5) and −∇L balanced by μ = 1.25; (b) the soft-thresholding function; (c) the toy lasso path vs λ (computed from the six-order toy)
Figure 2.8 · Constraints: (a) the toy contours with the cap w_items ≤ 1.5, the projected-GD path, the constrained optimum (3.25, 1.5) and −∇L balanced by μ = 1.25; (b) the soft-thresholding function; (c) the toy lasso path vs λ (computed from the six-order toy)

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) = \arg\min_w \tfrac12(w - z)^2 + t|w| = \operatorname{sign}(z)\,\max(|z| - t,\, 0)
\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

SettingDefaultWhat bitesFix
torch.optim.Adam(weight_decay=λ)L2 added to the gradientPenalty divided by \sqrt{\hat v} (Section 8)AdamW; exclude biases and norm weights from decay
scheduler.step() placementWhatever your loop doesA per-iteration schedule stepped per epoch barely moves; warmup never endsCount schedules in steps; step once per optimizer step
clip_grad_norm_ placement—Before backward: clips nothing; under fp16 before unscale_: wrong thresholdAfter backward, after unscale, before step
SGDRegressor()penalty='l2', alpha=1e-4, learning_rate='invscaling', eta0=0.01, power_t=0.25It's ridge with a decaying step, and unscaled features wreck itStandardScaler in a pipeline; set penalty and schedule on purpose
LogisticRegression()L2 with C = 1.0, lbfgs, max_iter=100Penalized coefficients ((−4.97, 4.03, 1.75) vs MLE (−5.05, 4.11, 1.78) here); a ConvergenceWarning you might not seeScale; raise max_iter; C=np.inf for no penalty
DataLoader(shuffle=False)No shufflingSorted data → biased gradients (Section 10)shuffle=True for training
Adam \varepsilon = 1e-8—Huge steps for weights with tiny \hat v in low precisionRaise \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

DomainOptimizer in practiceWhy
Credit scoring, insurance pricingIRLS / Newton, L-BFGSSmall convex models; exact, reproducible fits for regulators
Ads and recommendation CTRSGD / AdaGrad / FTRL on sparse featuresBillions of rows, rare features need larger steps
Delivery ETAs, demand forecastingGradient-boosted trees (Ch 10), small nets with AdamNonlinear tabular data
Operations constraints (SLAs, budgets, fairness caps)Projected or constrained solvers, Lagrangian penaltiesKKT and shadow prices are the business answer
LLM pretraining and fine-tuningAdamW + schedule + clippingThe 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."

IngredientTypical valueWhy (section)
AdamW, \beta_{1}0.9Momentum 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.1Clean shrink, independent of \hat v (8, 12)
Linear warmupe.g. 2,000 steps (LLaMA)Early \hat v unreliable; early curvature high (6, 8)
DecayCosine to 10% of the peak, or WSDRemoves the noise floor (6)
Gradient clippingGlobal norm 1.0Insurance against rare huge gradients (10)
Precisionbf16 mixed precision, fp32 master weightsfp32'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 chapterBecomesWhere in the book
Gradient of the MSE on six ordersBackpropagation through billions of weightsCh 16
\kappa and feature scalingLayerNorm, careful initialization, Adam's per-weight scalingCh 17, 21
Minibatch noise, \eta/BBatch-size and learning-rate scaling, critical batch sizeCh 23
Warm up, hold, cool downWarmup + cosine or WSD schedulesCh 22
AdamW's four D'sThe default LLM optimizerCh 22–24
Flat vs sharp, implicit regularizationWhy overparameterized models generalizeCh 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

Figure 2.9 · Memory map of Chapter 2 (schematic)
Figure 2.9 · Memory map of Chapter 2 (schematic)

Every hook in one place

HookWhat it unlocksSection
F-L-O (Family, Loss, Optimizer)The three choices of every model; this chapter is the O1
Slope says which way, curvature says how farGradient vs Hessian; the Taylor picture2
\text{Step} = -\eta \times \text{slope} (Ch 1)Gradient descent2
Chord above, tangent below, curvature upThree convexity tests3
SAM (non-negative Sums, Affine inside, Max)Building convex losses3
1/\lambda smooth, 2/\lambda the cliffThe four learning-rate regimes4
Center, then scale, then worry about correlationConditioning; \kappa = (1 + r)/(1 - r)4
Sublinear crawls, linear adds digits, quadratic doubles themConvergence rates5
Unbiased, cheap, noisyThe minibatch gradient6
Warm up, hold, cool downEvery modern schedule6
\beta is memory: window 1/(1 - \beta), step \eta/(1 - \beta)Momentum and every EMA7
Direction, Denominator, Debias, DecayAdamW's four parts8
Too big, too noisy, too saddle-happyWhy deep learning skips Newton9
SCUBA (Smooth, Convex, Unbiased, Bounded noise, Aligned scales)The five assumptions behind every guarantee10
NaN, Noise, Nothing, NotchesReading a bad loss curve10
Shuffle, Stratify, Sample by importance, ScaleFixing a biased gradient10
Signs of the Hessian: all up, all down, mixedMinimum, maximum, saddle11
Big steps can't sit in sharp holesStability selects flat minima11
Some People Dislike Constraints (KKT)Stationarity, Primal, Dual, Complementary slackness12
Projected: step, then snap. Proximal: step, then shrink.Constrained and lasso updates12
L1 kills, L2 shrinks (Ch 1)Why the prox gives exact zeros12
Zero, Forward, Backward, Clip, Step, ScheduleA PyTorch training step14
Always Warm up, Cool down, Clip, in bf16, with AdamWThe LLM recipe15
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.