OptimisationMODEL-SPECIFICCONTESTEDSIMULATED

Optimisers: SGD, Momentum, Adam

Three ways to decide what multiplies the gradient. None is universally best: Adam converges fast and sometimes generalises worse; SGD with momentum is still the default in much of vision; the choice interacts with the learning rate and with weight decay.

The problem, the obvious approach, and why it breaks

Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.

The question

A colleague says "just use Adam". Another says the vision team's models are all SGD with momentum and they generalise better. Who is right, and what would it take to know for this model?

The problem

A team building an image quality model for a marketplace has two candidate training configurations. The Adam run reaches a low training loss in a third of the steps; the SGD-with-momentum run takes longer and ends with a slightly better held-out score. Compute is limited and they must pick one recipe for the retraining pipeline.

The obvious approach

Adam works out of the box on almost everything, needs less learning-rate tuning, and converges faster. Use Adam with its default settings everywhere and spend the tuning budget on the architecture.

Why it breaks

The Adam run converges fast to a minimum that is slightly worse on held-out images. Over the retraining pipeline's life that small gap is a permanent tax on every model shipped, paid to save a few hours per training run.

How it breaks — usually after the offline metric looked fine
  • The Adam run converges fast to a minimum that is slightly worse on held-out images. Over the retraining pipeline's life that small gap is a permanent tax on every model shipped, paid to save a few hours per training run.
  • The "default settings" include an L2 penalty that, under Adam's per-parameter scaling, no longer acts as weight decay: parameters with large historical gradients are barely regularised. The team believes the model is regularised and it is not.
  • When the vision team's SGD recipe is copied for a transformer-based text model, it fails to train at all: the per-layer gradient scales differ so much that no single rate works, which is the problem Adam's per-parameter scaling exists to fix.
ProblemTargetDataRepresentationSplitModelTrainingEvaluationValidationDeploymentInferenceMonitoringDriftRetraining

What is being predicted, and from what data

This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.

Target
  • Find parameters θ with low loss that also generalise: the point reached matters as much as how quickly it is reached, because the model is judged on held-out images and served for months.
  • The surrounding system scores listing photos for a moderation queue; the decision downstream is which photos a human reviews first.
Data
  • Gradients from mini-batches of images, noisy and with very different magnitudes across layers: early convolutional filters see gradients that differ in scale from the final linear layer by orders of magnitude.
  • The optimiser sees only these gradient vectors, step after step. Its state — a velocity per parameter, or running moments per parameter — is derived entirely from their history.

How it actually works

Precisely enough to predict its behaviour — not a framework API.

  • SGD with momentum keeps a velocity v ← β·v + ∇L and steps along v. Consecutive gradients that agree accumulate into a larger step and gradients that alternate sign cancel, so the iterate accelerates along the valley floor and damps oscillation across the walls. The effective step along a consistent direction is roughly η / (1 − β) — with β = 0.9, ten times the nominal rate.
  • Adam keeps two running moments per parameter — a mean of gradients (momentum) and a mean of squared gradients — and divides the step by the square root of the second. Each parameter gets its own effective learning rate, small where gradients have been large and large where they have been small. That is why it tolerates layers with wildly different gradient scales and why it is the default for transformers (Transformer Fundamentals).
  • The per-parameter scaling is also why Adam and plain L2 regularisation interact badly: the penalty gradient is scaled down for exactly the parameters with large gradient history. Decoupled weight decay — subtracting a fraction of the weight directly, outside the adaptive step — restores the intended behaviour, and is what "AdamW" names.
  • The visualiser at /ml/descent makes the behaviours visible: on the valley, momentum overshoots and recovers, Adam walks the floor directly; on the ravine, the three optimisers can land in different minima from the same start, because the path, not only the destination, depends on the rule.

What multiplies the gradient

All three optimisers keep the update rule of gradient descent and change one thing: the vector that is actually subtracted. SGD subtracts the gradient. Momentum subtracts a running average of gradients, so consistent directions compound. Adam subtracts a running average divided, per parameter, by a running root-mean-square of past gradients, so each parameter is stepped in units of its own typical gradient size.

That last division is the important one. It makes Adam nearly invariant to the scale of each parameter's gradient, which is why it needs less tuning across layers with different scales — and it is also why the same division rescales any L2 penalty into something that is no longer weight decay.

The three updates side by side
1def momentum_step(p, g, state, lr, beta=0.9):
2 state["v"] = beta * state["v"] + g
3 return p - lr * state["v"]
4
5def adam_step(p, g, state, lr, b1=0.9, b2=0.999, eps=1e-8, wd=0.0):
6 state["t"] += 1
7 state["m"] = b1 * state["m"] + (1 - b1) * g # first moment
8 state["s"] = b2 * state["s"] + (1 - b2) * g * g # second moment
9 m_hat = state["m"] / (1 - b1 ** state["t"]) # bias correction
10 s_hat = state["s"] / (1 - b2 ** state["t"])
11 p = p - lr * wd * p # decoupled decay: outside the scaling
12 return p - lr * m_hat / (s_hat ** 0.5 + eps) # per-parameter step size

The decay is applied to p directly, before the adaptive step. Folding it into g instead would divide it by sqrt(s_hat) and shrink it for exactly the parameters with a large gradient history — the coupled-L2 mistake.

Choosing by family and evidence

The decision is not "which is best" but "which has been shown to work on models like this one, and can we afford to check". The table below scores the three on the axes an engineer actually trades, and its caveat is the most important cell.

Three optimisers on the axes that decide
OptionQualityCostOperationalData neededNote
SGD with momentumStrong record on convolutional vision; needs a tuned rate and schedule; one state tensor per parameter
Adam (coupled L2)Tolerant of gradient-scale heterogeneity; regularisation behaves unexpectedly; three tensors per parameter
Adam with decoupled decayThe transformer default; decay does what it says; same memory cost as Adam

caveat The quality column is a summary of contested evidence that flips by model family, and the scores say nothing about the interaction with the learning-rate schedule, which matters more than the optimiser choice on most tasks. Read the column as "worth testing", never as a ranking.

What the recipe assumes after it ships

A training recipe lives in a retraining pipeline for years and is re-run on data that changes. The optimiser choice was justified once, by a comparison on one dataset and one architecture; every later edit — a wider model, a new loss term, a different batch size — inherits the choice without the evidence.

The optimiser state is also part of the run. A checkpoint that saves the weights and not the moment estimates resumes as a different, worse optimiser for the first thousand steps, and the loss spike that follows is blamed on the data that arrived that day.

must stay trueThe comparison still applies

The optimiser and decay configuration chosen for this pipeline was compared under a matched budget on this model family, and the model has not changed family since.

holds when The architecture, loss and rough scale are those of the original comparison, and the learning-rate sweep was re-run for the current optimiser.

breaks when A different backbone is swapped in; a new loss term with a different gradient scale is added; the pipeline resumes from a weights-only checkpoint; a framework upgrade changes the default decay coupling.

how you would know A training-smoke test that pins the loss trajectory over the first few hundred steps against a reference run (Training Smoke Tests); a per-layer update-to-gradient ratio log with an alert on collapse; a checkpoint test that asserts optimiser state round-trips.

respond Re-run the learning-rate sweep and, if the family changed, the matched comparison. Do not change optimiser mid-pipeline on the basis of a single run's held-out score.

How to build it

Most important first.

  • Choose by model family and by evidence, not by reputation. For transformers and other architectures with heterogeneous gradient scales, start with Adam with decoupled weight decay. For convolutional vision models, SGD with momentum and a decay schedule remains a strong default and has a long record of generalising well.
  • Tune the learning rate for whichever optimiser is chosen; they have different stable ranges, and a rate that suits one is meaningless for the other (Gradient Descent).
  • If regularisation matters, use decoupled weight decay under Adam and verify the effective decay on the parameters that receive large gradients, rather than assuming an L2 term does what its name says (Regularisation).
  • When two recipes differ on held-out quality, run both to convergence with the same budget on the same split more than once with different seeds before deciding; a single run's gap is within seed noise more often than it looks (Random Seeds).

What to measure

Which number actually maps to the decision — and which numbers look relevant and are not.

  • Held-out loss at the end of a matched compute budget, averaged over seeds. That is the number the choice is about; training loss after a fixed number of steps favours whichever optimiser moves fastest early and says little about where it ends.
  • The effective per-parameter step sizes under Adam (the ratio of update to gradient) across layers, sampled during training — the diagnostic for whether the scaling is doing something useful or has collapsed in a layer whose second moment is tiny.
  • Do not compare optimisers by steps to reach a training-loss threshold. Adam wins that comparison almost by construction and it is not the decision.

What must stay true after deployment

The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.

Assumptions
  • The optimiser state — velocities, moment estimates — is saved and restored together with the weights when a run is checkpointed and resumed; resuming with fresh state is a different optimiser for the first thousand steps (Checkpointing).
  • The chosen recipe was compared against the alternative on this model family under a matched budget; a recipe inherited from another family carries no evidence for this one.
  • Weight decay under Adam is decoupled and its effective strength has been checked on the parameters with the largest gradient history.
How to verify — offline, online, and over time
  • Offline: for a new model family, run both candidates to convergence on the same split under a matched compute budget, with more than one seed, and compare held-out loss with its spread — not training loss at a fixed step.
  • During training: log the update-to-gradient ratio per layer; a layer with an effective step orders of magnitude off the others is either a scaling win or a collapse, and the log tells you which.
  • Over time: a retraining pipeline that changes optimiser or weight-decay handling must re-run the learning-rate sweep and the matched comparison; the previous evidence no longer applies.

What can go wrong

Failure modes in production
  • Adam's second-moment estimate is near zero for a parameter that has received almost no gradient, so its first real gradient produces an enormous step. The bias correction and warm-up exist for this; skip them and the early steps can wreck a layer.
  • Momentum's accumulated velocity carries the iterate through a sharp region it should have slowed for, and the loss spikes tens of steps after the gradient that caused it — the cause and the symptom are separated in time.
  • A recipe tuned under one optimiser is migrated to another with the same rate and schedule; the run is slower or diverges, and the optimiser is blamed rather than the untuned rate.
What the recommended approach costs
  • Adam adds two state tensors per parameter, tripling optimiser memory — a real constraint on large models where VRAM is the budget (Memory Bandwidth & VRAM).
  • SGD with momentum needs a carefully tuned rate and schedule and is less forgiving of gradient-scale heterogeneity, which is a tuning cost paid per model family.
  • A matched, multi-seed comparison of two recipes is several full training runs. Most teams cannot afford it for every model, and the honest position is that they are then choosing on prior, not on evidence.
Misreads
  • "Adam converged faster, so it is the better optimiser." It reached a low training loss faster. The model is judged on held-out data at the end of the budget, and that comparison can go either way.
  • "We have weight decay, the model is regularised." Under Adam with a coupled L2 term, the parameters with the largest gradients — often the ones most in need of it — are barely decayed. Check the effective decay or decouple it.
  • "SGD generalises better, so we should switch the transformer to SGD." The evidence for that claim comes mostly from convolutional vision; on transformers plain SGD often fails to train at all, because of gradient-scale heterogeneity, not because of generalisation.

Where this applies

ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.

  • MODEL-SPECIFICThe Adam-for-transformers, SGD-with-momentum-for-convolutional-vision split reflects where each has been shown to work; on tabular networks, recurrent models and small models the picture is muddier and either can win.
  • CONTESTEDPractitioners genuinely disagree about the generalisation gap. The strongest form of the "Adam generalises worse" position is a body of vision results where SGD with momentum reaches flatter minima and better test accuracy under matched budgets; the strongest reply is that most of those comparisons used coupled L2 rather than decoupled decay and an untuned Adam rate, and that with both fixed the gap shrinks or vanishes on many tasks. The honest answer is that it depends on the family and must be measured.
  • SIMULATEDThe behaviours described on the bowl, valley and ravine are produced by the two-parameter visualiser at /ml/descent; they illustrate the mechanism and are not measurements on any real model.

Where the depth lives

This domain teaches the model and hands the rest off by name.

Computer Architecturememory-hierarchy