Deep Learning Optimisers

SGD, momentum, RMSProp, Adam, AdaGrad — convergence and learning rate schedules

Batch norm and layer norm stabilise what each layer's *inputs* look like — they say nothing about how a weight should actually *move* once its gradient is known. That's the next question. You are training a ResNet on ImageNet with plain SGD at a fixed learning rate. After 30 epochs the loss flattens out. Switch to an optimiser that gives every weight its own effective step size, and in the next 5 epochs the loss drops more than it did in the previous 30. What just happened — and why might plain SGD still win in the end?

The core issue with one fixed learning rate is that different weights have wildly different lives. A weight in the embedding for a rare word might get a gradient once in a thousand steps; a weight in the final layer gets one every step. The same step size cannot suit both — the rare one needs a big push to make progress from its infrequent updates, the busy one needs small steps to avoid bouncing around its target.

Think of each weight as a hiker descending its own private mountain, carrying a cane that feels out how rough the terrain underfoot has been lately. A hiker on a long, gentle, featureless stretch takes big confident strides. A hiker on jagged, jumpy terrain takes small, careful ones. Force every hiker, on every terrain, to use the same fixed stride, and the gentle-terrain hiker crawls while the jagged-terrain one stumbles. Adaptive optimisers give each weight its own cane — its own effective learning rate, derived from its own gradient history.

The lineage: AdaGrad → RMSProp → Adam

AdaGrad keeps, for every weight, a running *sum* of its squared gradients — v_t = v_{t-1} + g_t² — and scales its step by 1/√v_t: w_t = w_{t-1} − (η / √(v_t + ε))·g_t. A weight whose gradient has been small or rare keeps a small v_t and gets a comparatively large step; a weight updated every step accumulates a large v_t and its step shrinks.

Walk two weights through a few steps. A *busy* weight sees a gradient of magnitude 1 every step: v_1 = 1 (step scaled by 1/√1 = 1), v_2 = 1 + 1 = 2 (1/√2 ≈ 0.71), v_3 = 3 (1/√3 ≈ 0.58) — its effective step keeps shrinking. A *rare* weight sees nothing for 99 steps, then a single gradient of the same magnitude 1: its v_t is still just 1 at that point, giving it a step scaled by 1/√1 = 1 — noticeably bigger than the busy weight's step at that same moment (v_t ≈ 100, scale ≈ 1/√100 = 0.1). That is exactly the behaviour the earlier crisis called for: the busy weight self-throttles, the rare one still gets a real push the moment it fires.

But AdaGrad's sum only ever grows — on *dense* problems, where every weight behaves like the busy one, the step size keeps shrinking toward zero and learning eventually stalls entirely. RMSProp fixes this with a *decaying* running average instead of a running sum: v_t = β·v_{t-1} + (1 − β)·g_t² (β typically 0.9–0.99) — old squared gradients fade out geometrically instead of piling up forever, so the step size never collapses.

Adam adds one more piece on top of RMSProp: a momentum term that smooths the *direction*, not just the scale. It keeps m_t = β1·m_{t-1} + (1 − β1)·g_t (a running average of the gradient itself) alongside v_t = β2·v_{t-1} + (1 − β2)·g_t² (RMSProp's adaptive scale), bias-corrects both — m̂_t = m_t/(1 − β1^t), v̂_t = v_t/(1 − β2^t) — and updates with w_t = w_{t-1} − η · m̂_t / (√v̂_t + ε).

That bias correction matters more than it looks. Take the very first step, default β1=0.9, β2=0.999, and a gradient g_1 = 1.0: m_1 = 0.9·0 + 0.1·1.0 = 0.1, and v_1 = 0.999·0 + 0.001·1.0² = 0.001 — both estimates are still mostly their zero initialisation, badly underestimating the true gradient. Bias correction divides each by (1 − β^1): m̂_1 = 0.1/0.1 = 1.0, v̂_1 = 0.001/0.001 = 1.0 — the correction exactly cancels the initialisation bias on step 1, so Adam's first update is already essentially full-size (≈ −η), computed from a single noisy gradient with no history behind it yet. That is exactly why Transformer training needs a learning-rate *warmup*: those first, history-free steps are the least trustworthy ones Adam ever takes, and a full-size η at that moment is the single most common cause of an early training blow-up. The combination converges fast on almost anything, which is why Adam is the default for most deep learning.

So why does SGD sometimes win?

Here is the twist. On image-classification benchmarks, a well-tuned SGD with momentum — which keeps a running *velocity* v_t = β·v_{t-1} + g_t (β typically 0.9, no bias correction, no per-parameter rescaling — just gradient history rolling into the update direction like a ball gathering speed downhill) and updates w_t = w_{t-1} − η·v_t — often *beats* Adam on validation accuracy, even though Adam trains faster. The reason ties back to the sharp-versus-flat-minima idea: Adam's per-weight rescaling lets it slide neatly into the *nearest* minimum, which tends to be a sharp, narrow one. Plain SGD keeps more of its gradient noise, which jostles it toward *flatter, wider* minima — and flat minima generalise better, especially when the test data drifts from training. So the trade is real: Adam gives you speed, SGD (tuned, with patience) can give a slightly better final model. Whichever you use, never leave the learning rate fixed for the whole run — decay it over time, and for Transformers *warm it up* first (start tiny and ramp up over the first few thousand steps), or the early, unreliable gradients will blow training up.

Key points

• Adam (or AdamW) when training speed matters; SGD+momentum for image classification when you want the best generalisation and can afford a full run. Adam is the safe default for Transformers and NLP. For vision with a big compute budget, well-tuned SGD+momentum often edges ahead on validation accuracy. And whenever you add weight decay to Adam, use AdamW: plain Adam's L2 penalty gets distorted by the per-parameter rescaling, while AdamW applies the decay uniformly — always prefer it.
• The trap: running Adam on a Transformer with no learning-rate warmup. At the very start, Adam's running estimates have no history, so the first few hundred steps take confident, full-size steps in noisy, unreliable directions — and the loss spikes or diverges in the first epoch. Ramp the learning rate up from near-zero over the first 1,000–4,000 steps so the estimates can settle before big updates land. Missing warmup is the most common cause of early Transformer training blow-ups.
• The diagnostic: plot training and validation loss per optimiser, not just final accuracy. If Adam reaches a lower *training* loss but the same or worse *validation* loss than SGD, it has found a sharper minimum, not a better one — and sharp minima are exactly what a shift in the test distribution punishes. If your deployment data differs from training, lean toward SGD's flatter optima; if training speed is the bottleneck, take Adam and accept the trade-off.

Recap

• One fixed learning rate can't suit every weight: a rare-word embedding sees a gradient maybe once in 1000 steps and needs big pushes when it does; a final-layer weight sees one every step and needs small ones. A single global step size serves neither well — each weight instead gets its own cane, sized to its own gradient history.
• AdaGrad: v_t = v_{t-1} + g_t² (running *sum* of squared gradients), step = η/√(v_t+ε) · g_t. Toy trace: a busy weight (gradient every step) has v_3=3, step-scale 1/√3≈0.58 and shrinking; a rare weight (silent 99 steps, one gradient) still has v_t=1 at that moment, step-scale 1/√1=1 — bigger than the busy weight's ≈0.1 at step 100. Flaw: v_t only grows, so on dense problems the step decays to zero and training stalls.
• RMSProp: swaps the sum for a decaying average, v_t = β·v_{t-1}+(1−β)·g_t² — old squared gradients fade instead of piling up, so the step never collapses.
• Adam: adds momentum on RMSProp's scale — m_t=β1·m_{t-1}+(1−β1)·g_t, v_t=β2·v_{t-1}+(1−β2)·g_t², bias-corrected m̂_t=m_t/(1−β1^t), v̂_t=v_t/(1−β2^t), update w_t=w_{t-1}−η·m̂_t/(√v̂_t+ε). Adam is the default for most deep learning.
• Bias correction ≈ full step on step 1: with g_1=1.0, β1=0.9, β2=0.999 → m_1=0.1, v_1=0.001 (both mostly zero-init) → correction gives m̂_1=0.1/0.1=1.0, v̂_1=0.001/0.001=1.0 → update ≈ −η already, from one noisy gradient with zero history. That's exactly why Transformers need learning-rate warmup — the earliest steps are the least trustworthy ones Adam takes, and a full-size η then is the most common cause of an early blow-up.
• SGD+momentum sometimes wins: momentum keeps a running velocity v_t = β·v_{t-1} + g_t (β≈0.9) and updates w_t = w_{t-1} − η·v_t — plain gradient history rolled into the step, no bias correction, no per-parameter rescaling. On image classification, well-tuned SGD+momentum often *beats* Adam on validation accuracy — Adam's per-weight rescaling slides neatly into the *nearest* (often sharp) minimum, while SGD keeps more gradient noise and drifts toward flatter, wider minima that generalise better under distribution shift. The trade is real: Adam gives speed, tuned SGD can give a slightly better final model.
• Use AdamW whenever you add weight decay: plain Adam's L2-in-the-loss penalty gets distorted by the per-parameter rescaling (√v̂ divides the penalty down on exactly the weights that need it most — backwards), whereas AdamW applies the decay uniformly and directly to the weights. Always prefer AdamW.
• Never leave the learning rate fixed for the whole run: decay it over time, and for Transformers *warm it up* first (start near-zero, ramp over ~1,000–4,000 steps) — at the start Adam's running estimates have no history, so full-size steps in noisy directions spike or diverge the loss. Missing warmup is the most common cause of early Transformer blow-ups.
• Diagnostic: if Adam reaches a lower *training* loss but the same or worse *validation* loss than SGD, it found a *sharper* minimum, not a better one — and a shift in the test distribution punishes sharp minima. Lean toward SGD's flatter optima when deployment data differs; take Adam when training speed is the bottleneck.

Check your understanding