Optimisers — SGD, Adam, AdamW
Momentum, adaptive learning rates, and weight decay done right. Why AdamW replaced Adam as the default and when SGD still wins.
Backpropagation tells you which direction to move each weight. The optimiser decides how far to move — and how to move smarter than just "subtract the gradient."
After backpropagation you have a gradient for every weight — the direction in which each weight should change to reduce the loss. The simplest possible update: subtract a small fraction of the gradient. That fraction is the learning rate. This is plain SGD. It works, but it has two major problems in practice.
First: the same learning rate for every weight. A weight that receives large, consistent gradient signals needs smaller steps to avoid overshooting. A weight that receives rare, tiny gradients needs larger steps to make any progress. Treating all weights the same wastes most of the gradient signal.
Second: gradient noise. Mini-batch gradients are noisy estimates of the true gradient. A single step in a noisy direction wastes a step. Accumulating direction from many past steps — momentum — filters noise and accelerates progress. Modern optimisers (Adam, AdamW) solve both problems simultaneously.
SGD is like hiking downhill in thick fog with one step at a time — you only see the slope directly under your feet right now. SGD with momentum is like a ball rolling downhill — it accumulates speed in consistent directions and is slowed less by small bumps. Adam is a smart hiker with a map of the terrain history — they take big steps on flat ground and small careful steps on steep or unpredictable terrain.
AdamW is Adam who also carries a light backpack that gets heavier the further they walk — gently pulling them back toward the origin (weight decay) to prevent them from wandering too far.
SGD and momentum — from naive update to direction accumulation
Plain SGD is the simplest possible optimiser: subtract learning_rate × gradient from each weight every step. Momentum extends this by accumulating a velocity — a weighted average of all past gradients. Instead of updating directly from the current gradient, you update from the velocity, which smooths out noise and accelerates in consistent directions.
Consistent gradients accumulate — speed builds up.
Noisy gradients cancel — oscillation dampened.
Adam — per-weight adaptive learning rates via first and second moments
Adam (Adaptive Moment Estimation) maintains two running statistics per weight: the first moment (exponential moving average of gradients — like momentum) and the second moment (exponential moving average of squared gradients — measures how large gradients have been historically). The effective learning rate for each weight is lr / √(second moment) — weights with large past gradients get a smaller effective step size automatically.
Adam vs AdamW — why weight decay was broken in Adam
In standard SGD, L2 regularisation (adding λ||W||² to the loss) and weight decay (subtracting λW from the weight directly) are mathematically equivalent. In Adam they are not — and this caused Adam's weight decay to be effectively much weaker than intended for years before anyone noticed.
The problem: in Adam, the L2 gradient λW gets divided by √v̂ just like any other gradient — weights with large historical gradients get a smaller effective weight decay than weights with small gradients. The regularisation strength varies per weight in an uncontrolled way. AdamW (Loshchilov and Hutter, 2019) fixes this by decoupling weight decay from the gradient update — applying it directly to the weight before the adaptive gradient step.
Learning rate schedules — warmup, cosine decay, and ReduceLROnPlateau
The learning rate is the single most important hyperparameter. A fixed learning rate is always a compromise — too high early on causes divergence, too low late in training means slow progress. Schedules give you the best of both: a high rate for fast early exploration and a low rate for precise final convergence.
Linear warmup is especially important for Adam-based optimisers. In the first steps, the second moment estimate v is near zero — the bias correction denominator (1−β₂ᵗ) is small, making v̂ small, making the effective learning rate very large. Warmup starts with a tiny learning rate and gradually increases it, preventing unstable large updates in the first steps.
When SGD+momentum beats Adam — and why generalisation differs
Adam converges faster in almost every setting. But on large-scale image classification (ImageNet-scale CNNs) and some NLP tasks, SGD+momentum often achieves better final test accuracy despite slower convergence. This is a known phenomenon with a theoretical explanation: Adam finds sharp minima (narrow valleys in the loss landscape) while SGD tends to find flat minima. Flat minima generalise better because small perturbations to weights — which happen naturally when data distribution shifts slightly — do not change the loss much. Sharp minima are sensitive to such perturbations.
Every common optimiser mistake — explained and fixed
The optimiser line almost nobody rewrites — and the schedule bug that does real damage
Open the training script for almost any production deep learning project — a HuggingFace Trainer config, a PyTorch Lightning module, an internal training framework at a mid-size company — and the optimiser line is nearly always the same: AdamW, a learning rate somewhere between 1e-5 and 1e-3 depending on whether the run is pretraining or fine-tuning, weight_decay around 0.01, paired with a warmup-then-decay schedule. This is not laziness. AdamW's adaptive, per-weight step sizes make it forgiving of imperfect learning rate choices across a huge range of architectures and datasets, so teams standardise on it and spend their tuning effort on data, architecture, and regularisation instead of relitigating the optimiser choice on every new project.
SGD with momentum still gets deliberately chosen in a narrower set of cases — mainly teams training large vision backbones from scratch on ImageNet-scale data, where the flatter minima it tends to find translate into a real generalisation advantage worth the slower convergence and extra learning-rate tuning. Outside of large-scale vision pretraining, reaching for SGD instead of AdamW as a default is unusual enough that it should have a specific reason behind it, not habit.
A team training a fraud model watches validation loss improve steadily for twenty-five epochs, then go completely flat for the next fifteen. The instinctive read is "the model has hit its capacity — time to add layers or more features." Before touching the architecture, a more experienced engineer asks a cheaper question first: what did the learning rate actually do during those fifteen flat epochs?
The scheduler was a StepLR configured to halve the learning rate every ten epochs — but scheduler.step() had been placed inside the batch loop instead of the epoch loop. With roughly two hundred batches per epoch, ten "steps" of decay happened before the first epoch even finished. By epoch fifteen the effective learning rate had been halved dozens of times over and was effectively zero. The model had not run out of capacity — it had simply stopped being allowed to move.
Five things people get wrong about optimisers
Faster convergence and better final performance are different properties, and this module's own comparison section shows them diverging: on large-scale image classification and some NLP tasks, SGD with momentum reaches worse training loss more slowly but ends up generalising better, because Adam's adaptive per-weight steps tend to settle into sharp, narrow minima while SGD's noisier trajectory tends to find flatter ones. "Always use Adam" is a reasonable starting default, not a universal law — the right answer depends on the architecture and how much the task rewards generalisation over raw convergence speed.
They share a name and half a mechanism, not the whole idea. SGD's momentum is only a first moment — an exponential moving average of the gradient direction — used to smooth out noise and build speed in consistent directions, but every weight still gets the same learning rate. Adam adds a second moment — an exponential moving average of the squared gradient — and divides the step by its square root, giving each individual weight its own adaptive effective learning rate based on how large its gradients have historically been. Momentum changes the direction of the step; Adam's second moment changes the size of the step, per parameter, independently of direction.
They are mathematically equivalent under plain SGD, which is exactly why the terms get used interchangeably — but that equivalence breaks under Adam. Adding an L2 penalty to the loss in Adam means the resulting gradient term λW gets divided by the same adaptive √v̂ as every other gradient, so weights with a history of large gradients receive weaker effective decay than weights with small gradients — an uncontrolled, per-weight regularisation strength nobody actually intended. AdamW exists specifically to restore the "same fraction shrinks every weight equally" behaviour that L2 and weight decay share under plain SGD but silently lose under Adam.
Past a certain threshold, a higher learning rate does not train faster — it stops training at all. As this module's errors section shows, too-large a step causes updates to overshoot the minimum and bounce between opposite sides of the loss valley, producing a loss that oscillates indefinitely rather than converging, sometimes even diverging to NaN. The actual fix for "training feels slow" is usually a schedule — a higher rate early for fast exploration, decayed down for precise convergence later — not a single higher static value applied for the whole run.
For Adam-family optimisers specifically, warmup addresses a real instability, not a refinement. In the first few steps the second moment estimate v is still close to zero, so its bias-corrected version v̂ is small, which makes the effective step size lr/√v̂ very large right when the network's weights are least trained and most sensitive to a bad update. Warmup starts the learning rate near zero and ramps it up over the first steps specifically to avoid this early-training instability — skipping it is a common, hard-to-diagnose cause of runs that diverge or perform erratically in exactly the first few hundred steps.
Optimisers — 5 questions interviewers actually ask
Plain SGD subtracts learning_rate × gradient every step — simple, but every weight gets the same step size, and noisy mini-batch gradients cause the update direction to wobble. Momentum adds an exponential moving average of past gradients (velocity) and updates from that instead of the raw gradient — this smooths out the noise and builds speed in consistent directions, addressing the wobble problem. Adam goes further by also tracking a second moment — a moving average of squared gradients — and dividing the step by its square root, so each individual weight gets its own adaptive effective learning rate: large, frequently-updated gradients get smaller steps, and rare, small gradients get larger ones.
Because in Adam, weight_decay is implemented as an L2 penalty folded directly into the gradient before the adaptive division by √v̂ — which means the "decay" ends up scaled unevenly per weight, exactly like any other gradient term, rather than shrinking every weight by the same intended fraction. AdamW decouples the two steps: it applies weight decay directly to the weight (W = W − lr×λ×W) separately from the adaptive gradient update, restoring the property that every weight decays by the same proportion regardless of its gradient history. This was identified by Loshchilov and Hutter in 2019, and AdamW has since replaced Adam as the practical default almost everywhere.
Adam's moment estimates m and v are both initialised to zero and updated as exponential moving averages, which means in the very first steps they are biased heavily toward zero — far below the true gradient magnitude, since there has not been enough history yet to average over. Bias correction divides each estimate by (1 − βᵗ), a factor that is small early on (correcting the underestimate strongly) and approaches 1 as t grows, so the correction fades out naturally once enough steps have accumulated real history. Without it, the first updates would be artificially tiny, slowing down early training for no good reason.
This is an empirically observed property of the loss landscape, not a bug in Adam. Adam's per-weight adaptive steps tend to converge into sharp minima — narrow valleys where the training loss is very low but small perturbations to the weights (which happen naturally when the data distribution shifts even slightly at test time) cause the loss to spike. SGD's noisier, non-adaptive trajectory tends to settle into flatter minima, which are more robust to those perturbations and therefore generalise better on unseen data. This is especially pronounced on large-scale image classification, which is why ImageNet-scale CNNs are still frequently trained with SGD+momentum rather than Adam.
Bias correction fixes the fact that m and v start at zero, but it does not fix the fact that v itself is estimated from only a handful of gradient samples in the first few steps — a small, noisy sample of squared gradients can make v̂ artificially small, which makes the effective step size lr/√v̂ artificially large right when the network is least trained and most vulnerable to a bad update. Warmup adds a second, independent safeguard on top of bias correction: it explicitly scales the learning rate from near-zero up to its target value over the first several hundred steps, giving the second-moment estimate time to stabilise before the optimiser is allowed to take full-sized steps.
Optimisers are chosen. Next: make deep networks stable and prevent them from overfitting.
You now have the complete training loop: forward pass, loss, backprop, optimiser step. Module 45 adds the two techniques that make deep networks stable and generalisable at scale — Batch Normalisation (stabilise activations between layers) and Dropout (prevent co-adaptation and overfitting). These are not optional extras — they are standard components of every production deep learning model.
Internal covariate shift, running statistics, and why model.eval() is not optional when BatchNorm is in your network.
🎯 Key Takeaways
- ✓SGD updates every weight by the same learning rate times the gradient. SGD with momentum accumulates a velocity — a weighted average of past gradients. Momentum smooths noisy gradient directions and accelerates in consistent directions. β=0.9 is the standard default.
- ✓Adam maintains per-weight adaptive learning rates using two moment estimates: the first moment (running mean of gradients — like momentum) and the second moment (running mean of squared gradients — measures gradient magnitude). Weights with large past gradients get smaller effective steps automatically.
- ✓Bias correction in Adam is essential in the first training steps. Without it, m and v start at zero and underestimate the true moments — producing unstable first updates. The correction terms 1/(1−β₁ᵗ) and 1/(1−β₂ᵗ) fix this and become negligible after ~100 steps.
- ✓AdamW decouples weight decay from the gradient update. In Adam, L2 regularisation is scaled by the adaptive learning rate — making it weaker for frequently-updated weights. AdamW applies weight decay directly to the weight before the gradient step — uniform across all weights. Always prefer AdamW over Adam.
- ✓Default starting point for any new deep learning project: AdamW with lr=1e-3 and weight_decay=0.01. Pair with CosineAnnealingLR or ReduceLROnPlateau. Only switch to SGD+momentum when you have evidence it generalises better — primarily large-scale image classification.
- ✓The mandatory training step order: optimizer.zero_grad() → forward pass → loss → loss.backward() → optimizer.step(). Never rearrange these four lines. zero_grad() must come before backward() — PyTorch accumulates gradients by default and calling zero_grad() after backward() clears the gradients before they are used.
Discussion
0Have a better approach? Found something outdated? Share it — your knowledge helps everyone learning here.