We now have models worth training: linear regressors (Chapter 1), classifiers (Chapter 2), and multilayer perceptrons (Chapter 3). We have losses that tell each of them how bad they are, and we know the direction of improvement is the negative gradient. What remains is the question every practitioner faces daily: how, exactly, do you run the descent when the dataset has a million examples, the model has millions of knobs, and the loss landscape is no longer a friendly bowl?
This chapter is about the workhorse answer, stochastic gradient descent, and the small set of refinements (learning rates, batch sizes, momentum, Adam) that turn it from a theoretical idea into the algorithm that trains essentially everything in this book.
4.1 Where losses come from, one more time
Before optimizing a loss, remember why it is that loss. In Chapter 1 we saw that assuming Gaussian noise on a linear process makes maximum likelihood collapse into least squares; in Chapter 2 the same argument with a categorical distribution produced cross-entropy. This is the general pattern: a loss function is a noise model in disguise. Choose how you believe the data deviates from your model, and maximum likelihood hands you the loss. The loss is how we tell the model it is doing badly \(\rightarrow\) choosing it well is not a detail, it is the supervision itself.
Look at the cost. One step touches every example; modern datasets have millions to billions of them, modern models have millions to billions of parameters, and training needs thousands of steps. One exact gradient per step is a luxury we cannot afford \(\rightarrow\) we need to change strategy to scale up.
4.3 Stochastic gradient descent
The idea is almost cheeky: we do not need the exact gradient. A good approximation is enough. Instead of the expensive sum over all \(n\) examples, average over a small random minibatch\(\mathcal{B}\) — think 32 examples against a million:
Why does this work? At a fixed parameter vector, a uniformly sampled minibatch gradient is unbiased: the expected contribution of example \(i\) to a batch equals its contribution to the full average, so
This statement holds for independent draws with replacement and also for one uniformly chosen subset without replacement. In expectation, that cheap gradient points where the expensive gradient points. Each estimate is noisy; its average direction is right.
In practice the algorithm is:
Shuffle the training data.
Split it into minibatches of size \(B\).
For each batch: compute the average gradient, update the parameters.
One pass through all batches is an epoch; repeat for multiple epochs.
NoteHonest fine print
There are two different facts hiding under the word “without replacement.” A single uniform subset, evaluated at a fixed \(\vect{w}\), is still unbiased and has less variance than with-replacement sampling. Its variance includes the finite-population correction \((n-B)/(n-1)\), which reaches zero at \(B=n\).
Random reshuffling is subtler. After the first batch changes \(\vect{w}\), the next batch comes from the remaining examples and is statistically tied to the earlier path. It is not an independent unbiased draw of the full gradient at this new iterate. Dedicated random-reshuffling theory handles that dependence; the one-line proof in Equation 4.4 does not. Keep the caveat in mind; keep using the shuffle.
4.4 The two zones of SGD
Strang’s treatment of SGD describes semi-convergence: rapid progress early, then oscillation near the solution. His convergence discussion also isolates the assumption behind Equation 4.4: the stochastic gradient must be unbiased at the parameter vector where it is evaluated. Put those ideas together on our convex least-squares bowl. Write the batch gradient as
Far from the optimum, the full gradient is large relative to the batch noise, so most batch directions make useful progress and their mean is the full descent direction. Near the optimum, the full gradient shrinks while batches can still disagree (“go left” — “go right”). A finite learning rate then makes the iterate bounce in what we will call the region of confusion:
Generate one least-squares problem and solve for its empirical optimum.
Partition the examples into equal batches and define their gradients.
Compare batch descent directions far from and at the optimum.
Verify that their fixed-point mean equals the full descent direction, then plot both zones.
Figure 4.1: The two zones at fixed parameter vectors. Left: equal-size batch directions vary, but their mean overlaps the full descent direction and points down the convex bowl. Right: at the empirical optimum the full gradient is essentially zero while individual batches still disagree, so finite steps create a region of confusion.
Here is the surprise: the bouncing is not merely tolerable; in some regimes it is a feature. A noisy step may perturb the iterate away from a saddle, a narrow basin, or a shallow trap that exact descent would follow more predictably. Minibatch noise can also bias training toward different solutions, and smaller batches sometimes improve generalization. None of these outcomes is guaranteed: the effect depends on the model, data, learning rate, and batch size. We will give that evidence its proper treatment in Chapter 6. For now, treat noise as part of the algorithm, not automatically as a defect or a cure.
4.5 The learning rate
One knob dominates all others. The learning rate \(\alpha\) scales every step, and its failure modes are asymmetric: too small wastes your compute budget crawling; too large overshoots the valley and diverges.
Define the reusable losses_for helper.
Prepare the inputs and fixed settings for the example.
Implement the learning-rate triptych.
Report or visualize the measured result.
# [1]def losses_for(lr: float, steps: int=60) ->list[float]: w, b, out = torch.tensor(-0.5), torch.tensor(2.0), []for _ inrange(steps): err = (w * x1 + b) - y1 out.append(float((err **2).mean())) w = w - lr *2* (err * x1).mean() b = b - lr *2* err.mean()return out# [2]plt.figure(figsize=(6.2, 3.4))# [3]for lr, style, color in [(0.005, "-", "#5379AA"), (0.12, "-", "#E57200"), (1.1, "--", "#722F37")]: plt.semilogy(losses_for(lr), style, color=color, label=f"$\\alpha = {lr}$")plt.xlabel("step"); plt.ylabel("loss (log scale)")# [4]plt.legend(); plt.tight_layout(); plt.show()
Figure 4.2: Same problem, same steps, three learning rates. Too small crawls; too large overshoots back and forth and climbs; the middle one converges quickly.
TipPractical rule of thumb
For the normalized toy problems here, plain SGD often starts around \(\alpha = 0.1\). That is a scale-calibrated starting point, not a universal constant; Adam, unnormalized features, and very deep models can require very different values. Read the training curves: crawling \(\rightarrow\) consider raising it; oscillating or exploding \(\rightarrow\) lower it. Later in training it often pays to decay the learning rate so the fine-tuning steps get smaller; that is a learning-rate schedule. One exotic-sounding relative, warmup (starting tiny and ramping up), will matter enormously when we train Transformers in Chapter 14.
4.6 The batch size
The batch size \(B\) sets where you live on the noise–efficiency trade-off:
Batch size
Character
Trade-off
Small (1–32)
Noisy, exploratory
Sometimes a generalization benefit; slow and less hardware-efficient
Medium (32–256)
Balanced
The usual default; hardware-efficient
Large (256+)
Smooth, high throughput
Less gradient noise; often needs learning-rate and schedule retuning
With independent sampling, gradient standard deviation scales like \(1/\sqrt{B}\): quadrupling the batch roughly halves the jitter. For a uniform subset without replacement, multiply by \(\sqrt{(n-B)/(n-1)}\); when \(B=n\), the noise is exactly zero. For \(B \ll n\) the correction is near one, which is why the simpler rule is useful. There is no universal best setting here, only a dial linking statistics and hardware.
What a batch may estimate — and what it may not
The license behind everything above deserves one paragraph of fine print, because the rest of the book will lean on it in three different ways.
Case 1 — decomposable objectives. When the loss is a mean of per-example terms, \(\loss(\vect{w}) = \E_{x}[\ell(x; \vect{w})]\), a minibatch average is an unbiased estimator of the full objective, and its gradient an unbiased estimator of the full gradient. That is the license SGD runs on, and it is why \(B\) is a compute-and-noise dial, not a correctness dial. Cross-entropy, squared error, and every loss in Parts I–III are Case 1.
Case 2 — nonlinear functionals of aggregates. Some quantities are a nonlinear function of an expectation, \(R\bigl(\E_x[s(x)]\bigr)\). A batch plug-in estimate \(R(\bar{s}_B)\) is generally biased (Jensen’s inequality gives the direction when \(R\) is convex or concave), and averaging per-batch values of \(R\) does not converge to the true quantity as training runs longer — only larger batches shrink the bias. Log-of-mean, KL between aggregate distributions, and correlation-style diagnostics live here. When the book meets one (Chapter 19’s aggregate-posterior trap is the flagship), it names the case and either restructures the estimand or declares the bias.
Case 3 — batch-defined objectives. Sometimes the batch is not estimating anything: it is part of the objective’s definition. A contrastive loss whose denominator ranges over the batch’s own candidates (Chapter 20) is a different objective at \(B=64\) than at \(B=256\) — neither is a biased version of the other, and \(B\) becomes protocol, to be reported like any other design choice. Decoding-side analogues exist too: beam width in Chapter 11 defines the search, it does not estimate it.
The audit habit, whenever a loss line reduces over anything: name the target, the estimator, the reduction axes, the case, and the boundary where the license fails. Chapters that meet cases 2 and 3 repeat this in one sentence; the full statement lives only here. Appendix E places these cases beside the population quantity, distribution, and uncertainty that complete the statistical contract.
4.7 Momentum: give the ball some mass
Vanilla SGD has exactly one control, \(\alpha\), and in narrow, curved valleys that is not enough: the iterate zigzags across the steep direction while inching along the shallow one. The fix is to give the update a memory. Keep a running velocity that accumulates gradients, and move along the velocity instead of the raw gradient:
with \(\beta \in [0.9, 0.99]\). This is the ball rolling downhill: in directions where gradients keep agreeing, speed builds; in directions where they flip sign every step, they cancel. The zigzag damps itself, and small bumps in the landscape get rolled through rather than obeyed.
4.8 Adam: adaptive steps per knob
Momentum treats every parameter alike, but some knobs sit on steep cliffs while others sit on plains. Adam (adaptive moment estimation) combines three ideas from the same optimization family:1 momentum (accumulate past gradients), per-parameter adaptive learning rates (scale each knob’s step by its own gradient history), and bias correction (account for both accumulators starting at zero). Formally, with elementwise operations:
where \(\hat{\vect{m}}, \hat{\vect{s}}\) are the bias-corrected accumulators. When to use what: SGD with momentum is simple, reliable, and strong for large models; Adam is the good default that works out of the box.
4.9 The race, on a problem where we know the truth
Claims about optimizers deserve a test you can rerun. We build a least-squares problem with a deliberately ill-conditioned landscape; the second feature is ten times the scale of the first, so the loss surface is a long narrow valley. We give all three optimizers the same 120-step, 16-example-batch budget, while using learning rates chosen for this toy: \(0.003\) for SGD and momentum, \(0.1\) for Adam. This is a mechanism comparison, not a controlled claim that one optimizer wins at a shared learning rate.
The cell also uses backward() as the framework instrument previewed in Chapter 1. It supplies gradients for the race; Chapter 5 opens that box next.
Prepare the inputs and fixed settings for the example.
Figure 4.3: The optimizer race on an ill-conditioned valley (feature scales 1 and 10), with an equal step budget and optimizer-appropriate learning rates (\(0.003\) for SGD/momentum, \(0.1\) for Adam). Plain SGD crawls along the shallow direction; momentum damps the zigzag; Adam’s per-parameter scaling largely neutralizes the conditioning.
Read the result honestly. This is one problem, chosen to expose conditioning; it is not proof that Adam always wins (it does not — plain SGD with momentum, well tuned, still trains many of the best vision models). What the race does show is mechanism: when different directions of the landscape have wildly different steepness, per-direction memory (momentum) and per-parameter scaling (Adam) buy you orders of magnitude. It also whispers a practical lesson: much of this valley’s narrowness came from unscaled features \(\rightarrow\)normalize your inputs and you fight the optimizer less (Exercise 5).
4.10 Three regularizers that live inside training
Regularization can act on the update, the hidden computation, or the duration of training. Ridge from Chapter 1 returns as a training-loop citizen; the other two mechanisms make their contracts explicit here.
Weight decay. Under plain SGD, adding an \(L_2\) penalty is equivalent (up to whether the coefficient is written as \(\lambda\) or \(\lambda/2\)) to one extra shrinkage term: \(\vect{w} \leftarrow (1 - \alpha\lambda)\vect{w} - \alpha\nabla\loss\). This is the ridge connection from Chapter 1. Many PyTorch optimizers expose a weight_decay argument, but the equivalence needs a boundary: coupled \(L_2\) inside an adaptive optimizer such as Adam is rescaled by that optimizer and is not the same trajectory as multiplicative shrinkage. AdamW decouples the decay step explicitly. Also, training recipes often exempt bias and normalization parameters rather than shrinking every parameter indiscriminately.
Dropout. Weight decay changes the update. Dropout changes the network that receives the update. Let \(h_i\) be one hidden activation, let \(p\) be its drop probability, and write \(q=1-p\) for the probability that it survives. During training, inverted dropout draws an independent mask \(r_i\sim\operatorname{Bernoulli}(q)\) and sends
\[
\widetilde{h}_i = \frac{r_i}{q}h_i
\tag{4.7}\]
to the next layer. A dropped activation becomes zero; a surviving activation is scaled by \(1/q\). The scaling is not cosmetic. Conditional on the unmasked activation,
The training-time signal therefore has the same first moment as the full signal, while each step sees a different thinned network. At evaluation time we remove the masks and use \(h_i\) itself. PyTorch’s nn.Dropout(p) implements exactly this convention; model.train() turns masking on and model.eval() turns it off throughout the model. This train/evaluation asymmetry is part of dropout’s definition, not an optional speed setting.
Here is the mechanism written directly, followed by a small seeded check. The empirical training mean is close to the input, while evaluation is deterministic and leaves the input unchanged:
Define the reusable inverted_dropout helper.
Prepare the inputs and fixed settings for the example.
Implement inverted dropout and verify its scale.
# [1]def inverted_dropout( h: torch.Tensor, p: float, training: bool=True) -> torch.Tensor:ifnot training:return hifnot0.0<= p <1.0:raiseValueError("p must satisfy 0 <= p < 1") q =1.0- p mask = (torch.rand_like(h) < q).to(h.dtype)return mask * h / q# [2]torch.manual_seed(6050)h = torch.tensor([1.0, -2.0, 4.0, 0.5])draws = torch.stack([inverted_dropout(h, p=0.25) for _ inrange(20_000)])mean_output = draws.mean(0)# [3]print("training mean:", [round(v, 4) for v in mean_output.tolist()])print(f"max |mean - h|: {(mean_output - h).abs().max():.4f}")print("evaluation:", inverted_dropout(h, p=0.25, training=False).tolist())
training mean: [1.0025, -2.0057, 4.0064, 0.5024]
max |mean - h|: 0.0064
evaluation: [1.0, -2.0, 4.0, 0.5]
There is a useful ensemble intuition here, with an important calibration. Training shares weights across many masked subnetworks; evaluation uses one full network whose activation scale matches their mean. That can resemble averaging many related predictors. It is not an exact ensemble identity, because nonlinear layers and shared weights prevent the output of the full network from being the literal average of all masked-network outputs.
Early stopping. Track the validation loss during training and stop when it turns upward, even though the training loss is still falling:
Figure 4.4: Early stopping as model selection along one training trajectory. In this schematic, both losses initially improve; after the validation minimum, training loss keeps falling while validation loss worsens. Stop at the turn and retain that checkpoint.
All three are common, inexpensive relative to training the model, and previews: the full story of why constraining or randomizing training can improve generalization is the business of Chapter 6.
NoteCheck yourself
Close the book for one minute and reconstruct one update.
Why is a minibatch gradient correct in expectation but noisy in any one draw?
What do momentum and Adam retain from earlier gradients?
Which operations belong inside the training loop, and which belong only at evaluation?
4.11 Okay, so — the training loop
Assemble the chapter into the loop you will run, in some form, for the rest of the book:
Start from (well-initialized) random parameters.
For each epoch: shuffle, split into minibatches of size \(B\).
For each batch: training-mode forward pass (including dropout, if used) \(\rightarrow\) loss (your noise model in disguise) \(\rightarrow\) backward pass for the gradients \(\rightarrow\) optimizer step (\(\alpha\), momentum or Adam, weight decay).
Switch to evaluation mode for validation; watch the curve, schedule the learning rate, and stop early when it turns.
That recipe, as code. The book keeps it as a tested function — Listing 4.1 — printed here once and imported by later chapters, which then show only what they change (their model builder and their budget):
Declare one reusable interface for the model factory, data, and training budget.
Seed before constructing the model and optimizer.
Reshuffle the examples and visit every minibatch each epoch.
Predict, measure cross-entropy, backpropagate, and update.
Return the trained model as the experiment artifact.
"""Listing 4.1 — the supervised training loop, importable.Chapter 4 derives this loop and prints it; Chapters 6, 8, and 9 import it andprint only their deltas (the model builder and the budget). The loop is thebook's canonical minibatch recipe: seed, build, then repeatpredict -> measure -> step over reshuffled minibatches."""from collections.abc import Callableimport torchimport torch.nn.functional as Ffrom torch import nn# [1]def fit_supervised( model_fn: Callable[[], nn.Module], X: torch.Tensor, y: torch.Tensor,*, epochs: int, batch: int=64, lr: float=1e-3, seed: int=6050,) -> nn.Module:"""Train a fresh model on (X, y) with cross-entropy and Adam. Seeding precedes construction, so a given (model_fn, seed) pair always starts from the same tensors; each epoch reshuffles example order. """# [2] torch.manual_seed(seed) model = model_fn() opt = torch.optim.Adam(model.parameters(), lr=lr)# [3]for _ inrange(epochs): perm = torch.randperm(len(X))for i inrange(0, len(X), batch): idx = perm[i : i + batch]# [4] loss = F.cross_entropy(model(X[idx]), y[idx]) opt.zero_grad() loss.backward() opt.step()# [5]return model
Two reading notes. Seeding precedes construction, so one (model_fn, seed) pair always trains from the same starting tensors — later chapters’ paired comparisons lean on that. And the loss line is the estimator licensed earlier in this chapter: a per-example objective averaged over a shuffled minibatch is an unbiased estimate of the full-data loss, which is precisely why the minibatch size is a compute knob rather than a correctness knob (Section 4.6.1).
One box in that loop is still magic: “backward pass for the gradients.” Computing millions of partial derivatives at the cost of roughly one extra forward pass is the subject of Chapter 5.
Sources and further reading
Tieleman and Hinton, “Lecture 6.5 — RMSProp,” COURSERA: Neural Networks for Machine Learning (2012): the original public provenance for RMSProp’s running average of squared gradients.
Gilbert Strang, Linear Algebra and Learning from Data, §VI.5, especially pp. 361–365, develops SGD’s fast start, later oscillation, unbiased-gradient assumption, and convergence-in-expectation view. The assigned excerpt is Resources/Gil Strang/SGD.pdf.
(Pencil.) Prove Equation 4.4 for sampling with replacement: if \(i\) is drawn uniformly from \(\{1,\dots,n\}\), show \(\E[\nabla \ell_i(\vect{w})] = \frac{1}{n}\sum_j \nabla \ell_j(\vect{w})\). Where exactly does the argument use uniformity?
(Code.) In the two-zones figure, sweep \(B \in \{1, 4, 16, 80\}\) and plot the final 30 steps of each path. How does the radius of the region of confusion scale with \(B\)? Compare against \(1/\sqrt{B}\) first, then include the finite-population correction. Why must the noise vanish at \(B=80\)?
(Code.) Add a learning-rate schedule to the race: halve \(\alpha\) every 30 steps for plain SGD. How much of the gap to momentum does scheduling close? What does that tell you about what momentum is really fixing here?
(Code.) Sweep momentum’s \(\beta \in \{0, 0.5, 0.9, 0.99\}\) in the race. Explain the failure mode at the top end using the ball analogy.
(Code.) Standardize the race’s features (divide each column of Xr by its standard deviation) and rerun all three optimizers. How much of Adam’s advantage evaporates? State the practical moral in one sentence.
(Pencil.) (a) From Equation 4.7, show that \(\operatorname{Var}(\widetilde{h}_i\mid h_i)=p h_i^2/(1-p)\). (Code.) (b) Estimate that variance for \(p\in\{0.1,0.5,0.9\}\) using the seeded check. What happens to signal noise as \(p\) approaches one, and what evaluation bug appears if masking is left on?
The per-parameter denominator is RMSProp. Adam adds bias-corrected first moments and names RMSProp directly.↩︎
