Dive into Deep Learning · §25.2
From the theory to the default optimizer
SGD without convexity · AdaGrad to AdamW · schedules, warmup, and the preconditioning ladder.
Motivation
Every modern network trains with AdamW, warmup, and a decaying schedule, not the plain gradient descent we analyzed. Two gaps to close:
Every optimizer below is under ten lines of NumPy, written by hand.
01
SGD without convexity
the descent lemma survives noise, taxed
Ghadimi–Lan
Feed an unbiased gradient with variance \sigma^2 through the quadratic ceiling and take expectations at the right moment:
\mathbb{E}\left[f(\mathbf{x}_{k+1}) \mid \mathbf{x}_k\right] \;\le\; f(\mathbf{x}_k) - \eta\bigl(1 - \tfrac{L\eta}{2}\bigr)\,\|\nabla f(\mathbf{x}_k)\|^2 + \tfrac{L\eta^2\sigma^2}{2}.
One new term in the descent lemma: a noise tax paid every step, gradient large or not.
Telescoping (Ghadimi–Lan, 2013) and balancing the two terms with \eta \propto 1/\sqrt{K}:
\mathbb{E}\bigl[\|\nabla f(\mathbf{x}_R)\|^2\bigr] \;\le\; \frac{2L\Delta}{K} + 2\sigma\sqrt{\frac{2L\Delta}{K}}, \qquad R \sim \mathrm{Uniform}\{0, \ldots, K-1\}.
Ghadimi–Lan
At \sigma = 0 the deterministic O(1/K) returns; with noise the K^{-1/2} term rules: 10\times smaller gradients cost 100\times the budget. Measured on a nonconvex toy, 20 seeds per budget:
K eta_K E|grad(x_R)|^2 min-so-far
125 0.0447 2.134e-01 9.385e-03
500 0.0224 8.485e-02 1.398e-03
2000 0.0112 3.774e-02 2.973e-04
8000 0.0056 1.729e-02 5.428e-05
log-log slope, random iterate: -0.60 (theory: -1/2)
log-log slope, min-so-far: -1.23The guarantee is for a randomly selected iterate: with noisy evaluations you can never identify the best one, and the min-so-far column (slope -1.23) is exactly the luck you cannot bank on.
02
Per-coordinate step sizes
AdaGrad → RMSProp → Adam
The thesis
On a diagonal quadratic, per-coordinate steps \eta_i = 1/\lambda_i solve the problem in one step: a diagonal Newton’s method, and \kappa simply disappears.
The zig-zag comes from using one number where n are called for. Adaptive methods estimate those n numbers from the gradients themselves.
The family
\mathbf{s}_t = \mathbf{s}_{t-1} + \mathbf{g}_t^2, \qquad \mathbf{x}_{t+1} = \mathbf{x}_t - \frac{\eta}{\sqrt{\mathbf{s}_t} + \epsilon}\,\mathbf{g}_t.
Two derivations, one rule: steepest descent in the metric \mathrm{diag}(\sqrt{\mathbf{s}_t}) (a coordinate is expensive in proportion to the evidence its gradients have been large) and the regret-optimal step for sparse features: a rare word’s step decays with its own activity, not wall-clock time.
\mathbf{s}_t never forgets, so steps decay like \eta/(\sigma\sqrt{t}): Robbins–Monro hard-wired in. On some nonconvex problems the resulting decay can make progress impractically slow because the accumulator remembers too much. RMSProp forgets on purpose: an EMA with memory \approx 1/(1-\beta_2) steps (10 at its standard \beta_2 = 0.9).
The family
Unroll \mathbf{v}_t = (1-\beta_2)\sum_{s\le t} \beta_2^{\,t-s}\mathbf{g}_s^2 and take expectations under a stationary scale \bar{\mathbf{g}^2}:
\mathbb{E}[\mathbf{v}_t] = \left(1 - \beta_2^{\,t}\right)\bar{\mathbf{g}^2}
by the geometric series. Dividing by 1-\beta_2^t is exactly unbiased at every t: the correction cancels the zero initialization identically, no approximation.
The transient is large: the raw ratio mis-scales early steps by (1-\beta_1^t)/\sqrt{1-\beta_2^t}, already 3.16 at t=1 and peaking above 6\times near t=12.
The family
\mathbf{m}_t = \beta_1\mathbf{m}_{t-1} + (1-\beta_1)\,\mathbf{g}_t, \quad \mathbf{v}_t = \beta_2\mathbf{v}_{t-1} + (1-\beta_2)\,\mathbf{g}_t^2, \quad \mathbf{x}_{t+1} = \mathbf{x}_t - \eta\,\frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}
RMSProp’s per-coordinate scale, momentum’s averaged direction, and the exact startup correction for both (\hat{\mathbf{m}}_t, \hat{\mathbf{v}}_t).
Near a diagonal quadratic minimum, \sqrt{\hat{v}_i} \approx |g_i| = \lambda_i|x_i|: the update is \approx \eta\,\mathrm{sign}(x_i), sign descent whose per-coordinate steps \eta/(\lambda_i|x_i|) carry the 1/\lambda_i ratio, reconstructed from first-order information alone.
03
When Adam fails
a convex counterexample, and the valley revisited
The counterexample
Reddi–Kale–Kumar (2018): on x \in [-1,1], cycle the convex losses
f_t(x) = \begin{cases} C\,x, & t \bmod 3 = 1,\\ -x, & \textrm{otherwise,}\end{cases} \qquad C > 2,
whose gradients sum to C - 2 > 0 per period, so the best point is x^\star = -1.
The mechanism in one sentence: the effective step on the informative gradient shrinks faster than its information accrues. The rare +C lands in \mathbf{v} quadratically and throttles its own step by \approx C; then \beta_2 forgets, and the two -1 steps run at full size, the wrong way. Signal is throttled by its own magnitude; noise is not.
The counterexample
Running the construction at C = 4, against AMSGrad and projected SGD with the same 1/\sqrt{t} decay:
x_t at t = 30 300 3000 15000 (x* = -1)
adam 0.745 1.000 1.000 1.000
amsgrad -0.953 -0.985 -0.995 -0.998
sgd -0.954 -0.986 -0.995 -0.998AMSGrad’s one-line fix, replacing \hat{\mathbf{v}}_t by the running maximum, makes per-coordinate steps nonincreasing, so the large gradient is never forgotten.
The theorem shows that vanilla Adam has no general convergence guarantee, even on convex objectives, without additional assumptions or modification. Practice keeps Adam anyway: real noise usually breaks up consistently-informative rare gradients.
The payoff
The \kappa = 10^3 quadratic of the gradient-based-optimization section, optimally tuned GD versus hand-rolled Adam:
GD, optimal single eta = 2.0e-03: f < 1e-8 at k = 6160
Adam, eta = 0.01:
k = 1: effective steps [1.e-02 1.e-05]
k = 100: effective steps [1.6534e-02 1.7000e-05]
k = 1000: effective steps [6.5755e-02 6.6000e-05]
Adam reaches f < 1e-8 at k = 344From the first iteration Adam’s per-coordinate steps sit in the ratio 10^{-2} : 10^{-5}, the eigenvalue ratio, reconstructed from gradient magnitudes with no Hessian. GD 6160, Adam 344.
The counterpoint: \sqrt{\hat{\mathbf{v}}} conflates curvature with noise, the diagonal misses every off-axis correlation (rotate the quadratic 45° and the rescaling stops helping), and Reddi steers it adversarially. A well-chosen point partway along the cost–fidelity curve.
04
Decoupled weight decay
coupled versus decoupled regularization
AdamW
Under SGD, “penalize the loss” and “shrink the weights” are the same update. Under Adam the penalty gradient rides through the preconditioner, and the shrinkage on coordinate i becomes
\underbrace{\frac{\eta\,\lambda}{\sqrt{\hat{v}_{t,i}} + \epsilon}\; w_{t,i}}_{\ell_2\ \textrm{through Adam}} \qquad \textrm{versus} \qquad \underbrace{\eta\,\lambda\, w_{t,i}}_{\textrm{decoupled}}.
The regularization strength is no longer a constant of the problem: noisy-gradient weights are barely decayed, quiet ones are decayed hard. Whatever \lambda meant (the Gaussian prior of MAP, the maximum-likelihood section; the norm constraint it multiplies, the constrained-optimization-and- duality section) assumed one \lambda for all coordinates. Adam’s preconditioner silently discards it.
AdamW
Decouple: the loss gradient goes through the preconditioner, the decay does not. Two pure-noise weights at scales \sigma = (10, 0.1), where decay is the only systematic force:
per-step decay rate, l2 through Adam (alpha*lam/sigma_i): [1.e-05 1.e-03]
per-step decay rate, AdamW (uniform): 0.0001
|w| after 4000 steps, Adam + l2: [0.9677 0.0235]
|w| after 4000 steps, AdamW : [0.6681 0.6598]
AdamW prediction (1 - alpha*lam)^T = 0.6703Coupled decay gave the two weights effective rates 100\times apart, a disparity nobody chose; AdamW shrinks both by the uniform (1-\eta\lambda)^T, within 2\% of prediction.
In the major libraries, AdamW and Adam with a weight_decay flag implement two different regularizers: the decoupled update and the coupled one.
05
Schedules and warmup
what decay does, and why ramps come first
Schedules
The mathematics a schedule negotiates was already proved: a constant step parks on a noise floor \propto \eta; Robbins–Monro decay reaches the optimum; beyond convexity no theorem ranks decay shapes.
Schedules
On the noisy quadratic (the miniature where the floor is exact), at 80\% of budget and at the end:
E f at 80% budget E f at end (floor ~ 0.088)
constant 9.53e-02 9.19e-02
c/k decay 2.82e-03 1.94e-03
cosine 8.71e-03 1.58e-03
WSD 80/20 9.53e-02 5.48e-03WSD sits on the constant’s floor at 80\% (same trajectory, to every digit), then 400 decay steps drop it 17\times: the loss-curve cliff of WSD runs, reproduced by a one-line quadratic. Cosine’s longer tail wins this pure-quadratic endgame; WSD’s case lives off the toy, in re-decayable checkpoints and time spent at large steps.
Warmup
For adaptive methods warmup is close to structural, and both reasons are this chapter’s instruments:
Ramp \eta while the preconditioner accumulates data and the curvature meets its step size. Every large run converged on this independently.
06
Beyond diagonals
the preconditioning ladder
The ladder
Which matrix B_t multiplies the gradient? GD says I; Newton says (\nabla^2 f)^{-1} at O(d^3); Adam says a diagonal at O(d). The rungs between exploit one structural fact: parameters come in matrices.
diagonal (Adam) → Kronecker-factored Fisher (K-FAC, natural gradient) → two-sided full-matrix roots (Shampoo) → spectral whitening by polar factor (Muon) → full Newton (unreachable)
K-FAC’s whole trick in one display, a curvature matrix with (mn)^2 entries preconditioned by two small inverses:
(A \otimes G)^{-1}\,\mathrm{vec}(V) \;=\; \mathrm{vec}\left(G^{-1}\, V\, A^{-1}\right).
Wrap-up