Optimisers I: SGD, Momentum and Nesterov Acceleration

In the mathematics track we analysed gradient descent and saw its Achilles' heel: on ill-conditioned problems — long, narrow valleys — it oscillates across the steep walls while creeping along the floor. Deep-learning loss surfaces are full of such ravines, plateaus and noisy gradients. Momentum is the simplest and most important fix, and it remains part of almost every optimiser used today.

Mini-batch SGD recap#

Problems:

1. Ravines — a learning rate small enough to be stable along the steep direction is tiny along the shallow direction.
2. Noise — mini-batch gradients fluctuate, causing jittery progress.
3. Plateaus and saddle points — tiny gradients mean tiny steps.

Momentum (heavy ball)#

Polyak's heavy-ball method (1964) keeps a velocity — an exponentially decaying accumulation of past gradients:

with momentum coefficient $\beta$ typically 0.9.

Physical intuition: a ball rolling downhill gathers speed in consistent directions and is not deflected much by small bumps.

Why it fixes ravines: across the valley, gradients alternate in sign, so they cancel in the velocity; along the valley floor they point the same way every step, so they accumulate. Unrolling the recursion, the effective step in a consistent direction approaches

— with $\beta = 0.9$, ten times the plain step. Momentum also averages out gradient noise over roughly $1/(1 - \beta)$ steps.

Nesterov accelerated gradient#

Nesterov's idea (1983): since we are going to move by roughly $\beta\mathbf{v}_t$ anyway, compute the gradient at that look-ahead point rather than the current one:

If momentum is about to carry us past the minimum, the look-ahead gradient already points back and applies the brakes earlier. Nesterov momentum reduces overshooting and oscillation, and for smooth convex problems it achieves the optimal $O(1/t^2)$ convergence rate among first-order methods. In practice it is a small but reliable improvement over classical momentum.

Seeing the difference#

import numpy as np

H = np.diag([1.0, 50.0])                  # ill-conditioned quadratic, kappa = 50
grad = lambda th: H @ th
start = np.array([10.0, 1.0])

def run(method, lr, beta=0.9, steps=100):
    th, v = start.copy(), np.zeros(2)
    for _ in range(steps):
        if method == "sgd":
            th = th - lr * grad(th)
        elif method == "momentum":
            v = beta * v + grad(th); th = th - lr * v
        elif method == "nesterov":
            v = beta * v + grad(th - lr * beta * v); th = th - lr * v
    return np.linalg.norm(th)

for m in ["sgd", "momentum", "nesterov"]:
    print(f"{m:<9} distance from optimum after 100 steps: {run(m, lr=0.035):.2e}")

With the same learning rate, momentum methods reach the optimum orders of magnitude faster than plain gradient descent.

Momentum in PyTorch#

import torch
model = torch.nn.Linear(10, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, nesterov=True, weight_decay=5e-4)

(PyTorch's formulation folds the learning rate slightly differently from the equations above, but the behaviour is equivalent for a constant learning rate.)

Tuning SGD with momentum#

• $\beta = 0.9$ is a robust default; 0.95–0.99 for very noisy or large-batch settings.
• The effective learning rate is about $\eta/(1 - \beta)$ — if you increase $\beta$, reduce $\eta$ accordingly.
• Use a learning-rate schedule (step decay, cosine) — SGD with momentum depends heavily on it.
• Add weight decay for regularisation.

SGD + momentum versus adaptive optimisers#

Adaptive methods such as Adam (next lecture) usually converge faster with less tuning, and dominate transformer training. But SGD with momentum remains competitive — and sometimes generalises slightly better — for convolutional networks in image classification, where many well-known results were obtained with it. It also uses less memory (one state vector per parameter versus two for Adam).