Optimizers

It is December 2014. Deep networks are getting deeper, batches are getting noisier, and every lab has its own folklore about learning-rate schedules — momentum for this architecture, per-parameter tricks for that one, hours of manual tuning before a single real training run begins. An eight-page paper out of the University of Amsterdam and the University of Toronto proposes a single update rule that, its authors claim, needs almost no tuning at all. Within a few years it is the default optimizer for nearly every deep learning model trained anywhere.

Where We Were

Stochastic gradient descent, in the form Robbins and Monro described in 1951, is the oldest idea in the room: take a noisy estimate of the gradient from a mini-batch, step against it, repeat. It works, eventually, but it treats every parameter identically — one global learning rate applied to millions of weights that may need very different step sizes. A weight near a saddle point wants a large step; a weight already near a good value wants a small one. Plain SGD cannot tell the difference.

Momentum, the heavy-ball method Polyak described in 1964, was the first fix: accumulate a running average of past gradients and step in that direction instead of the instantaneous one. This smooths out noise and accelerates progress through the long, shallow ravines that loss surfaces are full of. But momentum still applies the same learning rate to every parameter — it changes the direction of the step, not its per-parameter scale.

The second fix arrived from a different angle. AdaGrad, published by Duchi, Hazan, and Singer in 2011, divided each parameter’s learning rate by the square root of the sum of its squared gradients seen so far. Parameters that received large or frequent gradients got smaller effective steps; parameters that rarely moved got larger ones. This was exactly the per-parameter adaptivity SGD lacked, and it worked well for the sparse, high-dimensional features common in the NLP models of the day. But the fix carried its own flaw: because AdaGrad accumulates squared gradients over the entire training history without ever forgetting, the denominator only grows, and the effective learning rate decays monotonically toward zero. On short convex problems this was a feature. On the long, non-convex training runs deep networks required, it was a slow death — AdaGrad’s learning rate could vanish long before training converged.

Geoffrey Hinton and Tijmen Tieleman’s fix, presented in Hinton’s 2012 Coursera lecture and never formally published as a paper, was RMSprop: replace AdaGrad’s ever-growing sum of squared gradients with an exponential moving average, so old gradient history decays instead of accumulating forever. This solved the vanishing-learning-rate problem and made per-parameter adaptivity viable for long non-convex training runs — but RMSprop had no momentum term of its own, and, being an unpublished lecture slide, no formal convergence analysis or agreed-upon default hyperparameters.

By the end of 2014, the field had two half-solutions circling the same problem: momentum gave direction-smoothing but no per-parameter scale; RMSprop gave per-parameter scale but no momentum. Nobody had combined them into one default that a practitioner could reach for without extensive tuning.

The Breakthrough

Adam (Kingma and Ba, 2014) is exactly the combination the field was missing: it keeps an exponentially decaying running average of the gradient itself — momentum’s contribution, giving direction — and a separate exponentially decaying running average of the squared gradient — RMSprop’s contribution, giving per-parameter scale. Both estimates are, in the language of the paper’s title, “adaptive estimates of lower-order moments”: the first moment is the mean gradient, the second is its uncentered variance.

The detail that made Adam more than a mechanical merger was bias correction. Because both moving averages start at zero, they are biased toward zero in the first steps of training, especially with decay rates close to one. Kingma and Ba corrected for this analytically, dividing each moment estimate by a term that shrinks toward one as training proceeds. The resulting method, per the paper’s own description, is “invariant to diagonal rescaling of the gradients,” has “little memory requirements,” and is “well suited for problems that are large in terms of data and/or parameters” and for “very noisy and/or sparse gradients.” Critically, the abstract adds that the hyperparameters “have intuitive interpretations and typically require little tuning” — a direct answer to the tuning burden that plain SGD and momentum imposed. (The same paper introduced AdaMax, an infinity-norm variant, which saw far less adoption.)

Adam’s adoption was near-immediate and durable — it remains, a decade later, the starting point for nearly every training recipe in this book. But two structural gaps surfaced as models and batches grew.

The first gap was regularization. Standard practice for SGD was L2 regularization: add a penalty proportional to the weight’s magnitude directly into the gradient before the update. For plain SGD this is mathematically identical to weight decay — shrinking every weight by a fixed fraction each step. Loshchilov and Hutter showed in November 2017 that the equivalence breaks for adaptive methods like Adam: because the L2 penalty is folded into the gradient, it gets divided by the same per-parameter adaptive denominator as the rest of the gradient, so parameters with large historical gradients are decayed less than parameters with small ones — the opposite of what uniform weight decay is supposed to do. AdamW (Loshchilov and Hutter, 2017) fixes this by decoupling weight decay from the gradient path entirely, applying it as a direct, uniform shrinkage of the weights each step. AdamW “substantially improves Adam’s generalization performance,” letting it compete with SGD-with-momentum on tasks where Adam had previously trailed. It is, today, the more commonly used of the two names.

The second gap was scale. As accelerator clusters grew, labs wanted to train on ever-larger batches to use more chips in parallel per step, but pushing Adam’s or SGD’s global learning rate up to match a larger batch destabilized training or hurt final accuracy — a problem this book’s chapter on data parallelism and large batches covers in depth for the SGD case, including LARS, the layerwise-adaptive-rate technique for large-batch SGD. LARS solved the problem for convolutional image models but, per the LAMB paper, “performs poorly for attention models like BERT.” LAMB (You et al., April 2019) generalizes LARS’s layerwise trust-ratio idea to Adam’s adaptive update itself, and the empirical result was striking: LAMB trained BERT with a batch size of 32,868 “without any degradation of performance,” cutting training time “from 3 days to just 76 minutes” on a TPUv3 Pod.

A third, more structural direction addressed a different limitation entirely. Adam and its adaptive relatives treat every individual parameter independently — a diagonal approximation to the true curvature of the loss surface, ignoring how parameters within the same weight matrix interact. Shampoo (Gupta, Koren, and Singer, February 2018) instead maintains a small preconditioning matrix for each dimension of a tensor-shaped parameter, capturing correlations that diagonal methods discard, while remaining cheap enough that, per the paper, its “runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.”

The Mathematics

Adam maintains two running averages per parameter. Given the gradient $g_t$ at step $t$:

$$\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \end{aligned}$$

Symbol	Meaning
$g_t$	Gradient of the loss with respect to a parameter at step $t$
$m_t$	Exponential moving average of the gradient, the momentum term
$v_t$	Exponential moving average of the squared gradient, the scale term
$\beta_1, \beta_2$	Decay rates for the two moving averages, both close to 1

Because $m_0 = v_0 = 0$, the raw averages are biased toward zero in early steps. Adam corrects for this before using them:

$$\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1-\beta_2^t}$$ $$\theta_t = \theta_{t-1} - \alpha \, \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

Symbol	Meaning
$\hat{m}_t, \hat{v}_t$	Bias-corrected moment estimates
$\alpha$	Global learning rate
$\theta_t$	Parameter value after step $t$
$\epsilon$	Small constant preventing division by zero

AdamW changes only where weight decay enters:

$$\theta_t = \theta_{t-1} - \alpha \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \, \theta_{t-1} \right)$$

Symbol	Meaning
$\lambda$	Weight-decay coefficient
$\theta_{t-1}$	Parameter value before the current step
$\alpha$	Global learning rate
$\hat{m}_t, \hat{v}_t$	Bias-corrected moment estimates, as in plain Adam

LAMB applies a per-layer trust ratio on top of the AdamW update. For a layer $l$ with weights $\theta_l$:

$$r_l = \frac{\| \theta_l \|}{\left\| \dfrac{\hat{m}_l}{\sqrt{\hat{v}_l} + \epsilon} + \lambda \, \theta_l \right\|}$$ $$\theta_l \leftarrow \theta_l - \alpha \, r_l \left( \frac{\hat{m}_l}{\sqrt{\hat{v}_l} + \epsilon} + \lambda \, \theta_l \right)$$

Symbol	Meaning
$r_l$	Trust ratio for layer $l$
$\theta_l$	Weight tensor of layer $l$
$\hat{m}_l, \hat{v}_l$	Bias-corrected moment estimates for layer $l$‘s weights
$\lambda$	Weight-decay coefficient
$\alpha$	Global learning rate

Shampoo takes a different route: instead of one scalar variance estimate per parameter, it keeps a small matrix per side of a tensor-shaped weight. For a matrix weight $W$ with gradient $G_t$ at step $t$:

$$L_t = L_{t-1} + G_t G_t^{\top}, \qquad R_t = R_{t-1} + G_t^{\top} G_t$$ $$W_t = W_{t-1} - \alpha \, L_t^{-1/4} \, G_t \, R_t^{-1/4}$$

Symbol	Meaning
$G_t$	Gradient matrix at step $t$ for a matrix-shaped parameter
$L_t, R_t$	Kronecker-factor preconditioning matrices, one per tensor side
$W_t$	Weight matrix after step $t$
$\alpha$	Global learning rate

What It Enabled

Adam and AdamW’s most direct consequence was that training a deep network stopped requiring a bespoke learning-rate schedule for every new architecture. The same optimizer, with the same default hyperparameters, trained convolutional networks, LSTMs, and — starting with the transformer — every major language model this book’s sister volume covers. That portability came with a cost the field paid for later: Adam’s two moment tensors, $m$ and $v$, are each the same size as the parameters themselves, so an Adam-family optimizer roughly triples the memory a model’s weights alone would need (parameters, plus $m$, plus $v$) before gradients or activations are counted. For a 7-billion-parameter model stored in 32-bit precision, that is roughly 28 GB of weights plus another 56 GB of optimizer state — a bill so large it directly motivated the optimizer-state sharding this book covers in ZeRO and the memory tricks, and interacts closely with the loss-scaling techniques in mixed precision, since $v_t$‘s small values are exactly the kind of number that underflows in 16-bit formats.

Optimizer	Year	Extra per-parameter state	Designed for	Verified detail
SGD with momentum	1964	1x (velocity)	general-purpose baseline	oldest adaptive-direction method still in wide use
AdaGrad	2011	1x (cumulative squared gradient)	sparse features	learning rate decays monotonically to zero
RMSprop	2012	1x (EMA of squared gradient)	non-stationary objectives	never formally published, Hinton Coursera lecture only
Adam	2014	2x (m and v)	general deep learning	”invariant to diagonal rescaling of the gradients”
AdamW	2017	2x (m and v)	generalization, transformer pretraining	decouples weight decay from the adaptive step
LAMB	2019	2x (m and v) plus per-layer norms	very large batch distributed training	BERT at batch 32,868: 3 days cut to 76 minutes on TPUv3 Pod
Shampoo	2018	Kronecker factors per tensor side, larger than 2x	faster convergence per step	runtime per step comparable to SGD, AdaGrad, and Adam

The scale strand had its own legacy. LAMB’s trust-ratio mechanism made batch sizes in the tens of thousands practical for transformer pretraining, feeding directly into the data-parallelism and large-batch techniques that let a training run spread across hundreds of accelerators without a proportional loss in final quality. Shampoo’s structure-aware preconditioning, meanwhile, seeded an ongoing line of work on distributed second-order optimization for the largest training runs, trading extra per-step compute and engineering complexity for fewer total steps — a trade-off this book’s chapter on training runs in practice revisits when weighing wall-clock time against accelerator-hour cost. And because Adam’s default hyperparameters interact with a model’s initialization and normalization scheme, tuning them well is inseparable from the techniques in keeping training alive — an unstable loss curve is as often an optimizer-hyperparameter problem as it is a normalization problem.

Key Papers

Paper	Authors	Year	Link	Why it matters
Adam: A Method for Stochastic Optimization	Kingma & Ba	2014	arXiv:1412.6980	Combined momentum and per-parameter adaptive scale with bias correction; became the default deep learning optimizer
Decoupled Weight Decay Regularization	Loshchilov & Hutter	2017	arXiv:1711.05101	Fixed the L2-versus-weight-decay mismatch in adaptive methods; AdamW became the transformer-training default
Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes	You et al.	2019	arXiv:1904.00962	Layerwise trust-ratio scaling for Adam; enabled BERT pretraining at batch size 32,868
Shampoo: Preconditioned Stochastic Tensor Optimization	Gupta et al.	2018	arXiv:1802.09568	Structure-aware, Kronecker-factored preconditioning; a practical path toward second-order training

PM Lens