Dive into Deep Learning · §26.2
Stochastic and adaptive optimization
nonconvex SGD · coordinate scaling · AdamW · schedules · structured preconditioners.
Motivation
Many modern networks are trained with AdamW, warmup, and a decaying schedule rather than plain gradient descent. Two questions follow:
Every optimizer below is under ten lines of NumPy, written by hand.
01
SGD without convexity
the descent lemma with stochastic gradients
Ghadimi–Lan
Condition on \mathbf{x}_k: unbiasedness removes the linear cross term, and the conditional variance bound contributes L\eta^2\sigma^2/2. A biased gradient estimator leaves an additional term and falls outside this result.
Telescoping the resulting inequality 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 not directly selectable from noisy observations.
02
Per-coordinate step sizes
AdaGrad → RMSProp → Adam
Motivation
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.
A global step size is limited by the largest curvature. Adaptive methods estimate separate coordinate scales from the observed gradients.
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.
The update can be derived as 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 is cumulative, 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 accumulated denominator is large. RMSProp uses exponential averaging: 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 effective step on the large positive gradient is smaller than those on the two negative gradients. The rare +C contributes quadratically to \mathbf{v} and reduces its own effective step by \approx C. The \beta_2 decay then reduces that contribution, so the two -1 gradients receive larger effective steps toward the suboptimal boundary.
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 replaces \hat{\mathbf{v}}_t by a running maximum, making the denominator nondecreasing and addressing this construction’s failure mechanism.
This is an existence result for a periodic convex online problem. It shows that vanilla Adam has no general convergence guarantee without further assumptions; it does not predict failure on every finite-data training problem. AMSGrad’s guarantee likewise depends on its theorem assumptions.
Diagonal scaling
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 have ratio 10^{-2} : 10^{-5}, the eigenvalue ratio, reconstructed from gradient magnitudes with no Hessian. GD 6160, Adam 344.
\sqrt{\hat{\mathbf{v}}} conflates curvature with noise, and a diagonal preconditioner cannot represent off-axis correlations. Rotating the quadratic by 45° removes the benefit of this coordinate scaling; the Reddi et al. construction supplies a separate adversarial failure case.
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}}.
Coupling makes shrinkage coordinate- and time-dependent. The MAP and norm-constraint interpretations assume a uniform \lambda. AdamW keeps decay outside the preconditioner and restores uniform shrinkage.
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, an induced disparity; 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
On the persistent-noise quadratic, a constant step has a nonzero stationary error \propto \eta. Under Robbins–Monro assumptions, a suitable decaying step can converge to the optimum; beyond convexity no theorem ranks decay shapes.
Schedules
On the noisy quadratic, where the stationary error is known exactly, 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 matches the constant schedule at 80\% (the trajectories agree to every printed digit), then 400 decay steps reduce it by 17\times. Cosine’s longer tail gives the lower final error on this quadratic; WSD remains useful for re-decayable checkpoints and time spent at large steps.
Warmup
Warmup can address two early transients:
Ramping \eta while the preconditioner accumulates data can reduce the magnitude of early updates. The useful duration and shape remain model- and data-dependent.
06
Beyond diagonals
structured preconditioning choices
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). Intermediate methods 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
K-FAC preconditions a curvature matrix with (mn)^2 entries using two small inverses:
(A \otimes G)^{-1}\,\mathrm{vec}(V) \;=\; \mathrm{vec}\left(G^{-1}\, V\, A^{-1}\right).
Wrap-up