Unit 05.01: SGD, momentum, Adam, and practical trade-offs

Three optimisers cover almost everything you will do. The differences are real but smaller than the learning rate's.

What each one adds

SGD steps directly along the gradient. Simple, and sensitive to the learning rate.

Momentum accumulates a running average of past gradients, which carries the update through flat regions and damps oscillation across narrow valleys.

Adam adapts a separate step size per parameter from the recent gradient history. It usually converges fastest with the least tuning.

import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(400, 4)
y = (X[:, 0] * 2 + X[:, 1] * X[:, 2]).unsqueeze(1)


def run(make_opt, steps=300):
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 1))
    opt = make_opt(model.parameters())
    history = []
    for step in range(steps):
        opt.zero_grad()
        loss = nn.MSELoss()(model(X), y)
        loss.backward()
        opt.step()
        if step % 100 == 0:
            history.append(round(loss.item(), 3))
    return history, round(nn.MSELoss()(model(X), y).item(), 4)


for name, factory in [
    ("SGD", lambda p: torch.optim.SGD(p, lr=0.01)),
    ("SGD+momentum", lambda p: torch.optim.SGD(p, lr=0.01, momentum=0.9)),
    ("Adam", lambda p: torch.optim.Adam(p, lr=0.01)),
]:
    history, final = run(factory)
    print(f"{name:14} checkpoints {history} -> final {final}")

# Adam usually converges fastest with least tuning, which is why it is the
# sensible default. SGD with momentum often generalises slightly better on
# large vision problems -- worth knowing, not worth agonising over here.

Adam reaches a low loss in far fewer steps than plain SGD on identical data. Momentum sits between them.

That ordering is typical but not universal: SGD with momentum often generalises slightly better on large vision problems, which is why you still see it in published image models.

The mistake this prevents

Switching optimisers to fix a training problem caused by the learning rate. Each optimiser has its own sensible range — Adam around 1e-3, SGD around 1e-2 — and changing one without the other confuses the comparison.

Takeaway

Start with Adam. Change the learning rate before changing the optimiser.

Mark complete