Gradient descent is "walk downhill." But the loss landscape is a mountain range with ravines, plateaus, and saddle points. SGD walks naively. Momentum builds inertia through ravines. Adam scales each parameter's step by its own gradient history. AdamW is what trains GPT. Pick the wrong optimizer and a model that should train in 4 hours takes 4 days.
Learning Objectives
After this lesson, you will be able to:
Trace the historical evolution of neural network optimizers from full-batch gradient descent through SGD, momentum, AdaGrad, RMSprop, and Adam — and understand why each step solved a real problem in the previous one
Pick the right optimizer for the job — SGD with momentum for vision tasks, AdamW for transformers and NLP — and explain why the same architecture can change accuracy by 10% just by swapping optimizers
Implement SGD, momentum, and Adam from scratch and watch their trajectories on a 2D loss surface to see firsthand how adaptive learning rates change the path the optimizer takes
Avoid the most expensive optimizer mistakes: weight-decay coupling in vanilla Adam, picking a default learning rate that was tuned for a different optimizer, and forgetting that Adam's optimizer state doubles your GPU memory cost
Don't worry if the formulas look intimidating — every modern optimizer is just gradient descent with a couple of running averages bolted on. Once you see what each running average is for, the rest collapses to bookkeeping.
The starting point of everything: take the gradient of the loss across the entire training set, multiply by a learning rate, subtract from the parameters.
θt+1=θt−η∇L(θt)
In 1951 this was the only known way to train a model. Then Robbins and Monro published a paper that quietly reshaped the next 70 years of machine learning.
Instead of computing the gradient over the full dataset, compute it on a small mini-batch (32, 64, 256 examples). Each update is approximate — the gradient is noisy — but you get thousands of updates per epoch instead of one.
θt+1=θt−η∇LBt(θt)
The noise in SGD is not a bug — it is a feature. It helps the optimizer escape sharp local minima and converge to flatter ones, which generalize better. This is why SGD with momentum still wins on vision benchmarks today.
Plain SGD has a problem on real loss landscapes: when the loss surface is a long narrow valley (common in deep networks), SGD oscillates back and forth across the valley walls and barely moves down the valley floor. Each step partially cancels the last.
The fix: add momentum. Maintain a running average of gradients; step in the direction of that average instead of the raw gradient.
vt+1=βvt+∇L(θt)θt+1=θt−ηvt+1
Nesterov momentum is a small but clever twist: instead of computing the gradient at the current position, compute it at where the velocity is about to take you. This look-ahead correction reduces overshoot.
vt+1=βvt+∇L(θt−ηβvt)θt+1=θt−ηvt+1
#Adaptive Methods: One Learning Rate Per Parameter
Momentum helps with direction. But there is a deeper problem: in deep networks, different parameters have wildly different gradient magnitudes. Embeddings get tiny, sparse gradients. Output-layer weights get huge gradients. A single global learning rate η has to be small enough not to explode the noisy params and large enough to actually move the quiet ones.
The breakthrough idea: let every parameter have its own learning rate, derived from its own gradient history.
Accumulate the squared gradient for each parameter; divide the learning rate by the square root of that accumulator. Parameters that have seen big gradients get smaller effective learning rates; quiet parameters get bigger ones.
The flaw: the accumulator only grows, so effective learning rates monotonically shrink to zero. AdaGrad eventually stops learning. Works well for sparse problems (NLP with bag-of-words) but kills deep networks.
Replace AdaGrad's cumulative sum with an exponential moving average of squared gradients. Now old gradient magnitudes get forgotten — the effective learning rate stays alive.
In 2014 Diederik Kingma and Jimmy Ba combined the best of both worlds. Track the first moment of gradients (momentum, m) AND the second moment (RMSprop, v). Apply bias correction so the moments are unbiased estimates even early in training when the EMAs have not warmed up.
Race the optimizers across a 2D loss landscapeInteractive
Loading visualization...
Try this: Watch SGD oscillate in a narrow valley while momentum coasts through. Notice how Adam adapts step size per axis — it takes large steps along the valley floor and tiny ones across the walls.
There is a subtle bug in how Adam handles weight decay (the L2 penalty term used to regularize models). The original Adam paper added the L2 gradient into the gradient before the moment updates — which means the adaptive learning-rate denominator √v + εalso divides the weight decay. Parameters with large historical gradients get less weight decay than they should. The regularization is silently inconsistent.
In 2017, Ilya Loshchilov and Frank Hutter showed that decoupling weight decay — applying it to the parameters directly, outside the adaptive update — fixes the bug and consistently improves generalization. They called the fixed version AdamW.
θt+1=θt−η(v^+ϵm^+λθt)
AdamW is now the default optimizer for nearly every transformer trained today: BERT, GPT-2, GPT-3, GPT-4, Claude, Llama, ViT, and on. If you remember nothing else from this lesson, remember: for transformers, use AdamW; for vision CNNs, use SGD with momentum or AdamW.
Adam's reign as the default has been challenged twice in recent years.
Lion (Google Brain, 2023): instead of maintaining a second-moment buffer, use only the sign of the momentum. Cuts optimizer memory in half and matches AdamW on ViT training. The discovery process itself is interesting — Lion was found by an evolutionary search over optimizer programs.
Sophia (Stanford, 2023): uses a clipped second-order (Hessian-based) preconditioner. Claims 2x speedup over AdamW on language model pre-training, though the picture is more mixed in independent reproductions.
Neither has unseated AdamW yet. Both are worth watching if you train models at frontier scale where optimizer memory and convergence speed compound into millions of dollars in compute.
What Do You Think?
You are training a vision transformer (ViT) and the loss is plateauing after 5 epochs. AdamW is your optimizer. Which change is most likely to help?
ViTs are notoriously sensitive to learning-rate warmup. The original ViT paper used 10000 steps of linear warmup before cosine decay; without warmup the gradients in the first few hundred steps explode and break attention. Higher LR or smaller batch will not fix the warmup problem; switching to SGD on a transformer typically makes things worse, not better. (Full coverage of warmup and schedulers is in Practical Training Tricks, lesson 10.)
Every modern optimizer is gradient descent plus running averages. Momentum tracks the first moment, RMSprop/Adam/AdamW track the second moment, and that is essentially the full toolkit
AdamW is the default for transformers, SGD+momentum for vision CNNs. Picking wrong can cost 5-10% accuracy on the same architecture; the choice is dictated by the loss landscape, not the model size
Bias correction in Adam matters early in training. Without the (1 - β^t) correction, m and v are biased toward zero for the first ~1/(1-β) steps and the optimizer behaves erratically
AdamW's only difference from Adam is decoupled weight decay. It is a one-line code change that consistently improves generalization on every transformer benchmark; never use vanilla Adam with L2 regularization in the loss
Optimizer state is the hidden GPU cost at scale. Adam-family optimizers double your parameter memory; 8-bit optimizers and ZeRO sharding exist precisely to claw that back
Why does AdamW outperform vanilla Adam on transformer training?
The optimizer takes the gradient and turns it into learning. Next: why a randomly-initialized deep network often refuses to train at all — and how Xavier and He fixed it in two lines of math.