S099 · Optimizers — SGD, Momentum, Adam, RMSprop
How to actually train the thing.
Module M12: Deep Learning · Session 99 of 130 · Track: ML 🚴
What you'll be able to do after this session
- Explain Optimizers to a friend in one minute, out loud, no notes.
- Recognise it when you see it in real code / systems / papers.
- Complete the hands-on exercise at the end without looking up help.
Prerequisites
If you can't answer these in 30 seconds each, redo the linked session first:
Pre-read (skim before the session)
- 🎥 Intuition (5–15 min) · DeepLearningAI (Andrew Ng): Gradient Descent With Momentum — the 9-min classic on why momentum helps.
- 🎥 Deep dive (30–60 min) · DeepLearningAI: Optimization Algorithms (full week playlist) — Ng walks through momentum, RMSprop, Adam.
- 🎥 Hands-on demo (10–20 min) · StatQuest: Stochastic Gradient Descent, Clearly Explained!!! — clean visual for how SGD steps differ from full-batch GD.
- 📖 Canonical article · Sebastian Ruder: An overview of gradient descent optimization algorithms — the definitive survey of every optimiser worth knowing.
- 📖 Book chapter / tutorial · Deep Learning Book — Chapter 8: Optimization for Training Deep Models — rigorous coverage of SGD, momentum, adaptive methods.
- 📖 Engineering blog · Distill.pub: Why Momentum Really Works — interactive article; play with momentum on a 2D quadratic and it clicks forever.
(a) Intuition · 5 min
The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.
In one sentence: How to actually train the thing.
Backprop tells you which direction the weights should move. An optimiser tells you how far to move and how to smooth out the jitter. Different optimisers make different choices, and the choice matters: pick the wrong one and your loss curve looks like a stock chart in a crash.
Think of a ball rolling down a hilly landscape trying to find the lowest valley. Vanilla SGD is a ball with no mass — it takes exactly the step you give it, wobbles around wherever the ground is bumpy, and gets stuck in shallow potholes. Momentum gives the ball mass — it builds up speed in consistent directions and rolls straight through small potholes. RMSprop gives it a smart set of shocks that shorten strides on rough terrain (large recent gradients) and lengthen strides on smooth terrain. Adam = momentum + RMSprop in one package, and it's what you should reach for by default in 2026.
Gotcha to remember forever: the learning rate is the single most important hyperparameter in all of deep learning. A well-tuned learning rate with plain SGD often beats a badly-tuned Adam. And no matter which optimiser you pick, add a learning rate schedule (warmup then decay) — flat LR is almost always suboptimal. If your loss goes NaN in the first 100 steps, the learning rate is too high; if it barely moves after 1000, it's too low.
(b) Visual walkthrough · 15 min
Diagrams, worked examples, step-by-step visuals.
The optimiser family tree
Update rules cheat sheet (per parameter w, gradient g)
| Optimiser | Update rule | Hyperparameters |
|---|---|---|
| SGD | w -= lr * g | lr |
| SGD + Momentum | v = mu*v + g; w -= lr * v | lr, mu (~0.9) |
| RMSprop | s = rho*s + (1-rho)*g^2; w -= lr * g / sqrt(s + eps) | lr, rho (~0.9) |
| Adam | m = b1*m + (1-b1)*g; v = b2*v + (1-b2)*g^2; w -= lr * m_hat / (sqrt(v_hat) + eps) | lr, b1=0.9, b2=0.999 |
| AdamW | Adam, but subtract lr * wd * w for weight decay separately from adaptive step | lr, wd (~1e-2), b1, b2 |
Worked example: which optimiser gets there faster?
Imagine a 2D loss surface shaped like a long narrow valley (very stretched ellipse). The gradient at any point mostly points across the valley, not along it.
- SGD zigzags wildly across the valley walls, making tiny progress along its length. Slow.
- Momentum builds up velocity along the valley's length while the across-valley oscillations cancel out. Fast.
- Adam does something similar but automatically per parameter — no need to hand-tune momentum for each direction.
This is exactly the geometry of most deep networks near a minimum. It's why plain SGD without momentum is essentially never used for deep learning.
Learning rate schedules
- Warmup: ramp LR from 0 to peak over the first ~1000 steps. Prevents early divergence in transformers.
- Cosine decay: smoothly decay LR to zero over training. Standard for LLM pretraining.
- Step decay: drop LR by 10× at fixed epochs. Old-school CNN training (ResNet paper).
- ReduceLROnPlateau: halve LR whenever validation loss stops improving. Great for supervised fine-tuning.
(c) Hands-on · 20 min
Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.
# pip install torch matplotlib
import torch, torch.nn as nn
import numpy as np
torch.manual_seed(42)
# Regression: y = 2x + 1 + noise, but hard because of scale mismatch on features
N = 2000
x1 = torch.randn(N, 1)
x2 = torch.randn(N, 1) * 100 # deliberately huge scale to stress SGD
X = torch.cat([x1, x2], dim=1)
y = 2.0 * x1 + 0.01 * x2 + 1.0 + 0.1 * torch.randn(N, 1)
def make_model():
torch.manual_seed(0)
return nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
def train(opt_cls, name, lr, **kw):
model = make_model()
opt = opt_cls(model.parameters(), lr=lr, **kw)
lossf = nn.MSELoss()
losses = []
for step in range(300):
idx = torch.randperm(N)[:64]
xb, yb = X[idx], y[idx]
opt.zero_grad()
loss = lossf(model(xb), yb)
loss.backward()
opt.step()
losses.append(loss.item())
print(f"{name:>18} final loss={losses[-1]:.4f} best={min(losses):.4f}")
return losses
l_sgd = train(torch.optim.SGD, "SGD", lr=1e-4)
l_mom = train(torch.optim.SGD, "SGD+momentum", lr=1e-4, momentum=0.9)
l_rms = train(torch.optim.RMSprop, "RMSprop", lr=1e-2)
l_adam = train(torch.optim.Adam, "Adam", lr=1e-2)
l_adw = train(torch.optim.AdamW, "AdamW", lr=1e-2, weight_decay=1e-2)
# Add a LR scheduler to Adam
model = make_model()
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=300)
lossf = nn.MSELoss()
losses = []
for step in range(300):
idx = torch.randperm(N)[:64]
opt.zero_grad()
loss = lossf(model(X[idx]), y[idx])
loss.backward()
opt.step()
sched.step()
losses.append(loss.item())
print(f"{'Adam + cosine LR':>18} final loss={losses[-1]:.4f} best={min(losses):.4f}")What to observe when you run it:
- Plain SGD at
lr=1e-4on the badly-scaled featurex2is painfully slow. Trylr=1e-2and it diverges. - Adding
momentum=0.9to SGD improves things but still isn't robust. - RMSprop and Adam handle the scale mismatch automatically because they use per-parameter adaptive learning rates.
- AdamW (weight decay decoupled) is the modern default for anything non-trivial.
- Adam + cosine schedule usually beats flat-LR Adam by a noticeable margin, especially on longer runs.
Try this modification: replace torch.optim.Adam(..., lr=1e-2) with lr=1.0 and watch the loss explode into NaN. Then try lr=1e-8 and watch it barely move. That is why the LR-tuning meme ("3e-4 is a safe default for Adam") exists.
(d) Production reality · 10 min
How this shows up in real systems, gotchas, war stories, common bugs.
War story #1: Adam vs SGD on ImageNet. Facebook AI Research (2017) famously trained ResNet-50 on ImageNet in 1 hour using plain SGD + momentum + linear scaling rule (LR scales with batch size). Adam converges faster in the first few epochs but generalises slightly worse on some vision tasks. Rule of thumb in 2026: Adam / AdamW for transformers and NLP; SGD + momentum with heavy tuning for large-scale image classification. Both work; Adam is the safer default.
War story #2: the LR warmup revelation. Early transformer papers (attention is all you need, BERT) needed a warmup phase to work at all — without it, the first few Adam updates would push weights so far that training never recovered. Every LLM training script today includes a warmup schedule (typically 1–10% of total steps) followed by cosine decay. Skip warmup at your peril.
Common bugs:
- Learning rate too high. Loss goes to
NaNwithin tens of steps. Halve LR and try again. - Learning rate too low. Loss decreases smoothly but takes forever to reach a good value. Multiply LR by 3×–10× and try again.
- Forgot to zero gradients. Loss oscillates wildly or diverges within a few steps.
- Using Adam's
weight_decayon parameters where you don't want L2 (biases, layer norms). Use AdamW and provide a param group that excludes them. - Loading a checkpoint without loading the optimiser state. Momentum / Adam moments reset to zero, causing a spike in loss for a few hundred steps after every resume.
How top teams handle it:
- LLMs (GPT, LLaMA, Mistral, Claude) use AdamW with cosine decay, warmup, gradient clipping, and bf16 mixed precision. Almost universal.
- Vision (ResNet, ViT, EfficientNet) — varies: SGD+momentum for CNNs, AdamW for transformers.
- RL (PPO, SAC) — Adam with
eps=1e-5(slightly larger to stabilise). - Every serious training script does an LR-range test (Leslie Smith's method) once at the start of a project to identify the right LR band, then uses that as the peak of a schedule.
(e) Quiz + exercise · 10 min
- In one sentence, what does momentum add to plain SGD, and what problem does it solve?
- Write the Adam update rule for a single parameter using its gradient
gand moment estimatesm,v. - What is the difference between
AdamandAdamW, and why is AdamW usually preferred? - Why do transformer training scripts almost always include a warmup phase for the learning rate?
- Your training loss goes to
NaNafter 50 steps. Name three things to try, in order of what to check first.
Stretch (connects to S098): Backprop computes dL/dw. An optimiser uses that gradient to update w. If backprop is buggy (say, off by a sign for one layer), which optimiser will most obviously reveal the bug — SGD, momentum, or Adam? Justify by describing what the loss curve would look like in each case.
Explain-out-loud test
If you can't teach these 3 things to a friend without notes, redo the session:
- What is Optimizers? (one sentence, no jargon)
- When would you use it? (one concrete example)
- When would you NOT use it? (one anti-pattern)
What comes next
- Next session (S100): PyTorch Fundamentals — Tensors, Autograd, nn.Module
- Previous (S098): Backpropagation — Derived by Hand on a 2-Layer Net
- Hub: The 6-Month Learning Plan
Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.
More from M12 · Deep Learning
all modules →