Search Tech Journey

Find topics, journeys and posts

6-month learning plan99 / 130
back to blog
mladvanced 15m read

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)


(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)

OptimiserUpdate ruleHyperparameters
SGDw -= lr * glr
SGD + Momentumv = mu*v + g; w -= lr * vlr, mu (~0.9)
RMSprops = rho*s + (1-rho)*g^2; w -= lr * g / sqrt(s + eps)lr, rho (~0.9)
Adamm = 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
AdamWAdam, but subtract lr * wd * w for weight decay separately from adaptive steplr, 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:

  1. Plain SGD at lr=1e-4 on the badly-scaled feature x2 is painfully slow. Try lr=1e-2 and it diverges.
  2. Adding momentum=0.9 to SGD improves things but still isn't robust.
  3. RMSprop and Adam handle the scale mismatch automatically because they use per-parameter adaptive learning rates.
  4. AdamW (weight decay decoupled) is the modern default for anything non-trivial.
  5. 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 NaN within 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_decay on 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

  1. In one sentence, what does momentum add to plain SGD, and what problem does it solve?
  2. Write the Adam update rule for a single parameter using its gradient g and moment estimates m, v.
  3. What is the difference between Adam and AdamW, and why is AdamW usually preferred?
  4. Why do transformer training scripts almost always include a warmup phase for the learning rate?
  5. Your training loss goes to NaN after 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:

  1. What is Optimizers? (one sentence, no jargon)
  2. When would you use it? (one concrete example)
  3. When would you NOT use it? (one anti-pattern)

What comes next


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 →
  1. 096Perceptron & Activation Functions
  2. 097Multi-Layer Perceptron — Forward Pass
  3. 098Backpropagation — Derived by Hand on a 2-Layer Net
  4. 100PyTorch Fundamentals — Tensors, Autograd, nn.Module
  5. 101Regularization in DL — Dropout, BatchNorm, Weight Decay
  6. 102CNNs — Convolution, Pooling, ImageNet Architectures