Optimisers: SGD, Momentum & Adam (Deep Learning Tutorial)

By the end of this lesson you will be able to explain and write the three optimisers you will use most: plain SGD, SGD with momentum and Adam. You will see why they behave differently, compare them on a real network, and set up a learning rate schedule.

An optimiser is the rule that turns gradients into weight updates. So far we used the simplest one: w = w - lr * grad. It works, but it can be slow. Better optimisers reach a good answer in fewer steps and with less tuning.

The problem with plain SGD

Many loss surfaces look like a long, narrow valley: steep across, gentle along. Think of a ghat road in the Western Ghats: the slope down to the river is steep from side to side, but the road itself descends gently.

Plain gradient descent struggles here. The steep sides give big gradients, so it zig-zags from wall to wall. The gentle direction gives tiny gradients, so it barely moves forward. If you raise the learning rate to go faster along the valley, it bounces out of the walls.

Momentum

Momentum keeps a running memory of past gradients, called the velocity. Think of a cycle rolling down that ghat road. It builds up speed in the direction it keeps going, and side-to-side pushes cancel out.

The update in the form PyTorch uses is:

v = momentum * v + grad (usually momentum = 0.9)
w = w - lr * v

Gradients that point the same way every step add up. Gradients that flip sign every step (the zig-zag) mostly cancel.

Adam

Adam (adaptive moment estimation, Kingma and Ba, 2014) combines momentum with a per-weight step size. It keeps two running averages for every weight:

m: the average gradient (like momentum);
v: the average squared gradient, a measure of how large that weight's gradients usually are.

Each weight's step is m / sqrt(v), scaled by the learning rate. A weight with consistently huge gradients gets a smaller step; a weight with tiny gradients gets a larger one. Early on, both averages start at zero, so Adam divides them by a bias correction factor to fix that. PyTorch's defaults are lr=0.001, betas=(0.9, 0.999) and eps=1e-8.

Step 1: race three optimisers across a narrow valley. The loss is f(x, y) = 0.05*x² + 10*y², steep in y and very gentle in x. The best point is (0, 0).

PythonRun it yourself

import numpy as np
import matplotlib.pyplot as plt

def grad(p):
    return np.array([0.1 * p[0], 20 * p[1]])      # slopes of 0.05*x^2 + 10*y^2

def run(method, lr, steps=60):
    p = np.array([-8.0, 1.0])
    v = np.zeros(2)
    m, s = np.zeros(2), np.zeros(2)
    path = [p.copy()]
    for t in range(1, steps + 1):
        g = grad(p)
        if method == "sgd":
            p = p - lr * g
        elif method == "momentum":
            v = 0.9 * v + g
            p = p - lr * v
        elif method == "adam":
            m = 0.9 * m + 0.1 * g
            s = 0.999 * s + 0.001 * g ** 2
            m_hat, s_hat = m / (1 - 0.9 ** t), s / (1 - 0.999 ** t)
            p = p - lr * m_hat / (np.sqrt(s_hat) + 1e-8)
        path.append(p.copy())
    return np.array(path)

for method, lr in [("sgd", 0.09), ("momentum", 0.03), ("adam", 0.2)]:
    path = run(method, lr)
    print(f"{method:<9} lr {lr:<4}: after 60 steps at ({path[-1][0]:7.3f}, {path[-1][1]:7.3f})")
    plt.plot(path[:, 0], path[:, 1], "o-", markersize=3, label=method)

xs, ys = np.meshgrid(np.linspace(-9, 2, 100), np.linspace(-1.5, 1.5, 100))
plt.contour(xs, ys, 0.05 * xs ** 2 + 10 * ys ** 2, levels=15, colors="lightgrey")
plt.legend()
plt.title("Paths across a narrow valley")
plt.show()

Output

sgd       lr 0.09: after 60 steps at ( -4.651,   0.000)
momentum  lr 0.03: after 60 steps at ( -0.945,   0.026)
adam      lr 0.2 : after 60 steps at (  0.001,  -0.035)

Plain SGD has to use a small learning rate or it would bounce out of the steep y walls, so it creeps along x and is still far from 0. Momentum builds up speed along the valley and gets much closer. Adam scales each direction separately, so it moves steadily on both. In the chart, look at the zig-zags in the SGD path.

Step 2: write optimisers as small classes

To compare them on a real network, let us write each optimiser as a class with one step method. Each one updates a dictionary of parameters in place, just like torch.optim does.

PythonRun it yourself

import numpy as np

class SGD:
    def __init__(self, params, lr):
        self.params, self.lr = params, lr

    def step(self, grads):
        for k in self.params:
            self.params[k] -= self.lr * grads[k]

class Momentum:
    def __init__(self, params, lr, beta=0.9):
        self.params, self.lr, self.beta = params, lr, beta
        self.v = {k: np.zeros_like(p) for k, p in params.items()}

    def step(self, grads):
        for k in self.params:
            self.v[k] = self.beta * self.v[k] + grads[k]
            self.params[k] -= self.lr * self.v[k]

class Adam:
    def __init__(self, params, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params, self.lr, self.b1, self.b2, self.eps = params, lr, beta1, beta2, eps
        self.m = {k: np.zeros_like(p) for k, p in params.items()}
        self.v = {k: np.zeros_like(p) for k, p in params.items()}
        self.t = 0

    def step(self, grads):
        self.t += 1
        for k in self.params:
            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * grads[k]
            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * grads[k] ** 2
            m_hat = self.m[k] / (1 - self.b1 ** self.t)
            v_hat = self.v[k] / (1 - self.b2 ** self.t)
            self.params[k] -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

# quick test on one weight: minimise (w - 5)^2
for opt_class, lr in [(SGD, 0.1), (Momentum, 0.05), (Adam, 0.1)]:
    params = {"w": np.array([0.0])}
    opt = opt_class(params, lr=lr)
    for _ in range(200):
        opt.step({"w": 2 * (params["w"] - 5)})
    print(f"{opt_class.__name__:<8}: w = {params['w'][0]:.4f}")

Output

SGD     : w = 5.0000
Momentum: w = 5.0000
Adam    : w = 5.0001

All three find w = 5 on this easy bowl. The differences show up on harder surfaces, like real networks.

Step 3: compare them on handwritten digits

Now we train the same digits network three times, changing only the optimiser. SGD and momentum both use a learning rate of 0.05. Adam uses 0.003, a little above its usual default of 0.001, because this network is small.

PythonRun it yourself

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X = X / 16.0
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=42, stratify=y)

def init(seed=42):
    rng = np.random.default_rng(seed)
    return {"W1": rng.normal(0, 1, (64, 64)) * np.sqrt(2 / 64), "b1": np.zeros(64),
            "W2": rng.normal(0, 1, (64, 10)) * np.sqrt(2 / 64), "b2": np.zeros(10)}

def loss_and_grads(p, X, y):
    Z1 = X @ p["W1"] + p["b1"]
    A1 = np.maximum(0, Z1)
    Z2 = A1 @ p["W2"] + p["b2"]
    P = np.exp(Z2 - Z2.max(axis=1, keepdims=True))
    P /= P.sum(axis=1, keepdims=True)
    loss = -np.log(P[np.arange(len(y)), y]).mean()
    dZ2 = P.copy()
    dZ2[np.arange(len(y)), y] -= 1
    dZ2 /= len(y)
    dZ1 = (dZ2 @ p["W2"].T) * (Z1 > 0)
    return loss, {"W1": X.T @ dZ1, "b1": dZ1.sum(axis=0), "W2": A1.T @ dZ2, "b2": dZ2.sum(axis=0)}

def make_optimizer(name, p):
    state = {k: np.zeros_like(v) for k, v in p.items()}
    m = {k: np.zeros_like(v) for k, v in p.items()}
    t = [0]
    def step(g):
        t[0] += 1
        for k in p:
            if name == "SGD":
                p[k] -= 0.05 * g[k]
            elif name == "Momentum":
                state[k] = 0.9 * state[k] + g[k]
                p[k] -= 0.05 * state[k]
            else:  # Adam
                m[k] = 0.9 * m[k] + 0.1 * g[k]
                state[k] = 0.999 * state[k] + 0.001 * g[k] ** 2
                p[k] -= 0.003 * (m[k] / (1 - 0.9 ** t[0])) / (np.sqrt(state[k] / (1 - 0.999 ** t[0])) + 1e-8)
    return step

for name in ["SGD", "Momentum", "Adam"]:
    p, rng = init(), np.random.default_rng(0)
    step = make_optimizer(name, p)
    results = []
    for epoch in range(1, 11):
        order = rng.permutation(len(X_train))
        for start in range(0, len(X_train), 32):
            idx = order[start:start + 32]
            _, g = loss_and_grads(p, X_train[idx], y_train[idx])
            step(g)
        if epoch in (2, 5, 10):
            val_loss, _ = loss_and_grads(p, X_val, y_val)
            results.append(f"epoch {epoch}: {val_loss:.3f}")
    print(f"{name:<8} validation loss  " + "   ".join(results))

Output

SGD      validation loss  epoch 2: 1.240   epoch 5: 0.577   epoch 10: 0.329
Momentum validation loss  epoch 2: 0.218   epoch 5: 0.164   epoch 10: 0.106
Adam     validation loss  epoch 2: 0.453   epoch 5: 0.203   epoch 10: 0.127

With the same small learning rate, plain SGD is clearly the slowest. Momentum and Adam both reach a much lower validation loss within a few epochs. On real projects Adam is the usual first choice because it works well with its default learning rate. Well-tuned SGD with momentum is still popular for image models such as ResNets, where it often generalises slightly better.

Learning rate schedules

A learning rate schedule changes the learning rate during training. The common idea: take big steps early to make fast progress, then small steps later to settle precisely. It is like driving from Jaipur to a friend's house: fast on the highway, slow in the lanes near the house.

Popular schedules are step decay (divide by 10 every few epochs), cosine decay (a smooth curve down to almost zero) and warm-up (start tiny for the first few hundred steps, common with transformers).

Step 4: compute a warm-up plus cosine schedule.

PythonRun it yourself

import math
import matplotlib.pyplot as plt

def lr_at(step, total=1000, warmup=100, peak=0.001):
    if step < warmup:
        return peak * (step + 1) / warmup                      # straight line up
    progress = (step - warmup) / (total - warmup)
    return peak * 0.5 * (1 + math.cos(math.pi * progress))     # smooth curve down

for s in [0, 50, 99, 100, 300, 550, 800, 999]:
    print(f"step {s:>3}: lr = {lr_at(s):.6f}")

plt.plot([lr_at(s) for s in range(1000)])
plt.xlabel("step")
plt.ylabel("learning rate")
plt.title("Warm-up then cosine decay")
plt.show()

Output

step   0: lr = 0.000010
step  50: lr = 0.000510
step  99: lr = 0.001000
step 100: lr = 0.001000
step 300: lr = 0.000883
step 550: lr = 0.000500
step 800: lr = 0.000117
step 999: lr = 0.000000

Optimisers in PyTorch

In PyTorch every optimiser lives in torch.optim and every schedule in torch.optim.lr_scheduler. Here they minimise the same narrow valley from Step 1. This is plain code with its real output.

Python

import torch

def valley(p):
    return 0.05 * p[0] ** 2 + 10 * p[1] ** 2

for name, make in [("SGD", lambda ps: torch.optim.SGD(ps, lr=0.09)),
                   ("SGD+momentum", lambda ps: torch.optim.SGD(ps, lr=0.03, momentum=0.9)),
                   ("Adam", lambda ps: torch.optim.Adam(ps, lr=0.2))]:
    p = torch.tensor([-8.0, 1.0], requires_grad=True)
    opt = make([p])
    for _ in range(60):
        opt.zero_grad()
        valley(p).backward()
        opt.step()
    print(f"{name:<13} ends at {[round(v, 3) for v in p.tolist()]}")

model_params = [torch.zeros(3, requires_grad=True)]
opt = torch.optim.AdamW(model_params, lr=0.001, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
for epoch in range(10):
    opt.step()          # (after a real epoch of training)
    sched.step()        # once per epoch
    if epoch in (0, 4, 9):
        print(f"after epoch {epoch}: lr = {sched.get_last_lr()[0]:.6f}")

Output

SGD           ends at [-4.651, 0.0]
SGD+momentum  ends at [-0.945, 0.026]
Adam          ends at [0.001, -0.035]
after epoch 0: lr = 0.000976
after epoch 4: lr = 0.000500
after epoch 9: lr = 0.000000

The final positions match our NumPy versions, which shows our hand-written optimisers follow the same rules. AdamW is Adam with "decoupled" weight decay, explained in Overfitting & Regularisation. It is the standard choice for transformers.

Common mistakes

Using Adam's default learning rate with SGD. lr=0.001 is tiny for plain SGD. Each optimiser needs its own range.
Using SGD's large learning rate with Adam. lr=0.1 with Adam usually makes the loss explode.
Calling the scheduler at the wrong frequency. CosineAnnealingLR(T_max=10) stepped once per epoch is not the same as stepped once per batch.
Creating the optimiser before moving the model to the GPU. Create it after model.to(device), so it holds the right parameters.

Key points

Plain SGD zig-zags in narrow valleys and is slow along gentle directions.
Momentum keeps a velocity of past gradients, which speeds up steady directions and cancels zig-zags.
Adam keeps averages of the gradient and squared gradient, giving each weight its own step size.
Adam (or AdamW) with lr=0.001 is a strong default; SGD with momentum 0.9 is common for image models.
Learning rate schedules such as warm-up and cosine decay take big steps early and small steps late.

Exercise

In Step 1, try momentum values of 0.5, 0.9 and 0.99 with lr=0.03. Print where each run ends after 60 steps. What happens when the momentum is too high?

Show solution

Python

import numpy as np

def grad(p):
    return np.array([0.1 * p[0], 20 * p[1]])

for beta in [0.5, 0.9, 0.99]:
    p, v = np.array([-8.0, 1.0]), np.zeros(2)
    for _ in range(60):
        v = beta * v + grad(p)
        p = p - 0.03 * v
    print(f"momentum {beta}: ends at ({p[0]:7.3f}, {p[1]:7.3f})")

Output

momentum 0.5: ends at ( -5.597,  -0.000)
momentum 0.9: ends at ( -0.945,   0.026)
momentum 0.99: ends at (  5.913,  -0.350)

With 0.5 the memory is short and progress along the valley is slow. With 0.9 it gets much closer to (0, 0). With 0.99 the velocity is so strong that the point overshoots and keeps swinging back and forth like a heavy pendulum, and after 60 steps it is still far from the bottom.

Quick Quiz

1. What does momentum add to plain SGD?

2. What does Adam track for each weight?

3. What is a sensible first learning rate for Adam?

4. Why do many training runs use a learning rate warm-up?

The problem with plain SGD

Momentum

Adam

Step 2: write optimisers as small classes

Step 3: compare them on handwritten digits

Learning rate schedules

Optimisers in PyTorch

Common mistakes

Key points

Exercise

Quick Quiz

Learn Deep Learning with a mentor