Chapter 12 // Training at Scale
Optimization
Since chapter 3 you've been typing optimizer names like incantations. This chapter opens the box: why descent gets stuck, why noise helps and hurts, and how five memories turn raw gradients into Adam.
Adapted from the original chapter of Dive into Deep Learning · CC BY-SA 4.0 · the Python shown is the book's; the figures run a faithful TypeScript port in your browser
Training a deep network means one thing: drive a loss downhill, millions of parameters at a time. But the ground is treacherous, the compass is noisy, and the stride length can kill you. Start with the terrain.
Act 1 · The terrain
01/05
First, an honest disclaimer: the loss you minimize is measured on your finite training set. The risk you care aboutis measured on all data. Their minima sit at different places — perfect optimization of the wrong curve. That gap is chapter 3's generalization story; this chapter is about the descent itself.
Every algorithm here descends the gradient. The whole craft is in how far to step, and the simplest possible function is enough to expose the stakes.
Act 2 · One knob: the step size
01/05
epoch 10, x: 0.060466
Gradient descent on f(x) = x², starting at x = 10 with η = 0.2. Each hop is x −= η·f′(x), and each hop shrinks x by the same factor — ten steps land at 0.06, a whisker from the optimum ◎.
Full gradients are a luxury: one update costs the whole dataset. Chapter 3's answer was to sample — and sampling has a price of its own.
Act 3 · Noise, and what it costs
01/05
50 steps · ends (0.11, -0.19) · distance 0.22
Real training can't afford the full gradient, so SGD estimates it from one random example — here simulated as the true gradient plus N(0,1) noise. The estimate is unbiased, so the walk still finds the valley… then jitters around it forever. The noise never runs out.
Noise tamed, one problem remains that no batch size fixes: valleys that curve a hundred times harder in one direction than another.
Act 4 · Memory
01/04
η 0.4 · ends (-0.94, -0.0001) — x₂ done, x₁ crawling
A nastier landscape: f = 0.1x₁² + 2x₂² — a canyon, twenty times flatter along x₁ than across it. At η = 0.4, plain descent tames the steep walls fast but inches along the floor: twenty steps and x₁ is still at −0.94.
Three leaky memories power everything from here — a velocity, a squared-gradient scale, and (for Adam) both at once, bias-corrected:
Act 5 · A learning rate per coordinate
01/05
adagrad η 0.4 · 20 steps · ends (-2.38, -0.16) — stalled
Adagrad's move: give each coordinate its ownlearning rate, η/√s, where s hoards every squared gradient ever seen. The steep x₂ axis racks up s fast and gets tiny careful steps; the flat x₁ axis keeps bigger ones. Smooth — but s only grows, so every step shrinks, and at η = 0.4 it stalls at x₁ = −2.4.
One knob left. Every trick so far chose the direction better; the last one retires the step size on a schedule.
Act 6 · The schedule
01/05
constant η 0.1 · 50 noisy steps · final distance 0.22 — and it will never improve
One last knob — η as a function of time. Hold it constant and the act-3 story replays: fast start, then permanent jitter at whatever radius the noise dictates. The final distance is a property of η, not of how long you train.
The lab
The chapter's claims, yours to stress-test: that momentum survives steps gd cannot, that per-coordinate scaling changes which η works at all, and that when you decay matters as much as whether.
The optimizer race
Pick a surface and an η, wake the optimizers, and race the whole chapter at once.
f = 0.1·x₁² + 2·x₂² (ill-conditioned) · start (−5, −2) · step 0
Momentum uses β 0.5, RMSProp γ 0.9, Adadelta ρ 0.9, Adam 0.9/0.999 — the book's settings. Try η 0.6 on the canyon: plain gd detonates while momentum glides. Then hand Adagrad η 2 — the one optimizer that wants it.
Momentum's extra headroom
Drag η and β — find the steps gd could never survive.
40 steps on the canyon · final distance 0.000
the stability gauge (steep axis, λ = 4)
effective step η/(1−β): 1.20
verdict: inside the momentum region
β widens the stable range from η·λ < 2 to η·λ < 2+2β. Set η 0.6 and slide β to zero to watch the theory come true — then push η past even the momentum limit.
Schedule the step size
Choose a decay policy and a base η, then let 50 noisy steps judge it.
The run is the act-3 simulation: true gradient plus N(0,1) noise on the bowl. Base η 1 or 2 exceeds the bowl's stability limit — see which policies merely delay the explosion and which never survive it.
final distance to the optimum: 0.067
The convex playbook
Deep learning is nonconvex, yet every algorithm above was born and analyzed on convex ground — where the chord test of act 1 guarantees no trap exists, and where SGD provably converges at rate O(1/√T) if η decays like the act-3 polynomial.
The real code
Every act carries its section's PyTorch in the </> chip. The book drives all
of them through one shared harness — a linear regression on the airfoil dataset
where only trainer_fn changes:
def train_ch11(trainer_fn, states, hyperparams, data_iter,
feature_dim, num_epochs=2):
w = torch.normal(mean=0.0, std=0.01, size=(feature_dim, 1),
requires_grad=True)
b = torch.zeros((1), requires_grad=True)
net, loss = lambda X: d2l.linreg(X, w, b), d2l.squared_loss
for _ in range(num_epochs):
for X, y in data_iter:
l = loss(net(X), y).mean()
l.backward()
trainer_fn([w, b], states, hyperparams) # the only variable
d2l.train_ch11(sgd, None, {'lr': 0.05}, data_iter, feature_dim)The concise versions — every optimizer is one import away
trainer = torch.optim.SGD
d2l.train_concise_ch11(trainer, {'lr': 0.01}, data_iter)
trainer = torch.optim.SGD # momentum built in
d2l.train_concise_ch11(trainer, {'lr': 0.005, 'momentum': 0.9}, data_iter)
trainer = torch.optim.Adagrad
d2l.train_concise_ch11(trainer, {'lr': 0.1}, data_iter)
trainer = torch.optim.RMSprop
d2l.train_concise_ch11(trainer, {'lr': 0.01, 'alpha': 0.9}, data_iter)
trainer = torch.optim.Adadelta
d2l.train_concise_ch11(trainer, {'rho': 0.9}, data_iter)
trainer = torch.optim.Adam
d2l.train_concise_ch11(trainer, {'lr': 0.01}, data_iter)
# and the schedule, applied per epoch
scheduler = lr_scheduler.MultiStepLR(trainer, milestones=[15, 30], gamma=0.5)
for param_group in trainer.param_groups:
param_group['lr'] = scheduler(epoch)Three things to remember
- The learning rate is a contract with curvature. Descent shrinks error only while |1 − ηλ| < 1 in every direction — too small crawls, too large detonates, and an ill-conditioned valley can make both true at once.
- Noise is both the discount and the bill. Sampling gradients makes each step cheap and even shakes you out of traps, but it never switches off — so average it (minibatches: noise ∝ 1/√B) and outlast it (decay η, not too fast).
- The modern optimizers are just leaky memories. Momentum remembers direction (÷(1−β) reach); AdaGrad-family remembers scale per coordinate; Adam remembers both, bias-corrected — and a warmup-plus-decay schedule remains the last knob nothing else replaces.
Adapted from Chapter 12 of Dive into Deep Learning by Zhang, Lipton, Li, and Smola (CC BY-SA 4.0). The figures run this site's verified TypeScript port; all computation happens live in your browser.