Optimisation and gradient descent variants

From batch descent to Adam, what momentum and adaptive scaling really change, how schedules help, and how to diagnose a diverging run.

Batch, stochastic and mini-batch

Full-batch gradient descent takes an exact step per epoch; stochastic descent takes one per example and is far noisier. Mini-batch sits between: enough examples for a useful GPU kernel and a gradient that points roughly downhill, plus noise that helps escape saddles.

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(1000, 5))
true_w = rng.normal(size=(5,))
y = X @ true_w + 0.1 * rng.normal(size=1000)

def grad(w, Xb, yb):
    r = Xb @ w - yb
    return 2 * Xb.T @ r / len(yb)

w = np.zeros(5)
lr = 0.05
batch = 32

for epoch in range(50):
    order = rng.permutation(len(X))          # shuffle every epoch
    for i in range(0, len(X), batch):
        idx = order[i:i + batch]
        w -= lr * grad(w, X[idx], y[idx])

print(np.round(w, 3), np.round(true_w, 3))

Gradient noise falls as 1 / sqrt(batch_size), so quadrupling the batch halves the noise — and you can roughly double the learning rate in exchange.
Larger batches give diminishing returns past a point; the wall-clock gain is limited by how many steps you can afford, not by how clean each step is.
Always shuffle. Contiguous sorted batches (all class A, then all class B) produce large gradient swings that look like a learning-rate problem but are actually a data-ordering problem.
The loss you track during training is a running estimate. Compare a smoothed curve, or evaluate on a fixed held-out batch, before concluding that training has stalled.

Momentum, RMSProp and Adam

# momentum: accumulate velocity, damped
def sgd_momentum(w, grads, lr=0.1, beta=0.9, state=None):
    state = state or {"v": np.zeros_like(w)}
    state["v"] = beta * state["v"] + grads
    return w - lr * state["v"], state

# Adam: keep a decaying average of the gradient and of its square
def adam(w, grads, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, t=1, state=None):
    state = state or {"m": np.zeros_like(w), "v": np.zeros_like(w)}
    state["m"] = b1 * state["m"] + (1 - b1) * grads
    state["v"] = b2 * state["v"] + (1 - b2) * grads ** 2
    m_hat = state["m"] / (1 - b1 ** t)          # bias correction
    v_hat = state["v"] / (1 - b2 ** t)
    return w - lr * m_hat / (np.sqrt(v_hat) + eps), state

Optimiser	Keeps	Strength	Watch out for
SGD	nothing	Simple, often best final generalisation	Slow on ill-conditioned problems
SGD + momentum	velocity	Fast through ravines and plateaus	Momentum overshoot if the rate is too high
RMSProp	squared gradient average	Per-parameter scaling, works without tuning	Sensitive to the decay and epsilon
Adam	first and second moments	Reliable default on sparse and noisy gradients	Can generalise slightly worse; decays badly without weight decay
AdamW	moments plus decoupled decay	Correct weight decay, the modern default	Two hyperparameters instead of one to set

⚠️

Weight decay in plain Adam is added to the gradient, so it is scaled by the same adaptive denominator and behaves inconsistently across parameters. AdamW decouples it. If you intend real L2 regularisation, use AdamW, not Adam with a weight_decay flag.

Schedules and diagnosing divergence

# linear warmup then cosine decay — the common transformer schedule
def lr_at(step, total, base_lr=1e-3, warmup=100):
    if step < warmup:
        return base_lr * step / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    return 0.5 * base_lr * (1 + np.cos(np.pi * progress))

import matplotlib.pyplot as plt
steps = np.arange(2000)
plt.plot(steps, [lr_at(s, 2000) for s in steps])
plt.xlabel("step"); plt.ylabel("learning rate"); plt.yscale("log")

Symptom	Diagnosis	Fix
Loss is `nan` within a few steps	Rate far too high, or a log/sqrt of a non-positive value	Cut the rate 10x, clip gradients, guard the loss
Loss spikes then recovers	Occasional bad batch	Gradient clipping, warmup, skip the batch
Loss oscillates around a floor	Rate near the stability limit	Lower the rate or increase the batch
Loss decreases then worsens	Rate too high late in training	Add decay, or use a plateau schedule
Loss perfectly flat	Rate too low, or dead units	Raise the rate 10x, check for zero gradients

Do a learning-range test before a long run: sweep the rate upward over a few hundred steps and record the loss. Loss falls, bottoms out, then rises sharply — pick a value an order of magnitude below the divergence point.

FAQ

Which optimiser should I start with?

AdamW at 1e-3 with a small weight decay for most problems; SGD with momentum and a schedule when you are squeezing out the last fraction of accuracy on a vision model. Change the learning rate before changing the optimiser.

Why does my loss go down then up?

Almost always the learning rate is too high for the region the model has reached, or the schedule never decayed. Confirm by plotting the loss against step and the learning rate against step on the same axis.

Multivariable calculus for training Numerical stability and floating point

Last refreshed 2026-09-18.

Optimisation and gradient descent variants

Batch, stochastic and mini-batch

Momentum, RMSProp and Adam

Schedules and diagnosing divergence

FAQ

Related