Optimisers: SGD, momentum, RMSProp and Adam with bias correction

The gradient descent of lesson 1 works, but on neural networks it is slow: the error surface has narrow, long valleys, flat regions and millions of dimensions with very different curvatures. Modern optimisers modify the update rule to go faster and more stably. This lesson presents mini-batch stochastic gradient descent, momentum, RMSProp and Adam (the most used optimiser), proves why Adam needs the “bias correction”, discusses the choice of step size over time, and compares the methods on the Rosenbrock function, a curved valley that is the classic test bench.

Stochastic gradient descent

The gradient of the loss is an average over all examples. With millions of examples, computing it entirely at each step is a waste. Stochastic gradient descent (SGD) uses at each step a mini-batch of a few tens or hundreds of examples drawn at random:

is an unbiased estimate of the true gradient (the mean of a random sample has the population mean as its expectation), with a variance that decreases as . The noise is not only a defect: it helps to escape flat regions and, according to many studies, leads to “wider” minima that generalise better. A complete pass over the data is called an epoch.

Momentum

In a narrow valley the gradient points mostly towards the walls: descent bounces from side to side and advances slowly along the bottom. Momentum (Polyak, 1964: the “heavy ball method”) accumulates a velocity:

with . The gradient components that change sign at every step (across the valley) cancel in the average; those that always have the same sign (along the valley) add up. If the gradient were constant, the velocity would tend to : with a step ten times longer in the good direction.

RMSProp and Adam: one step size per parameter

Different parameters can have gradients of very different scales (for example the weights of the first layers and those of the last). RMSProp (Hinton, 2012) divides each parameter’s step by a moving average of the root of its squared gradient: those with large gradients take shorter steps, those with small ones longer steps. Adam (Kingma and Ba, 2014) combines this idea with momentum:

with the default values , , . The effective step is almost independent of the gradient’s scale: it is about in the direction of the sign of the average gradient.

Theorem (bias correction). If the gradients have constant mean and , then ; so has mean exactly .

Proof. Unrolling the recursion, . Taking the mean, . The same holds for with . ∎

Without correction, in the first steps and would be very close to zero (they start from zero) and, especially with , the denominator would be underestimated: the first steps would be huge. After a few thousand steps and the correction no longer matters. AdamW (2017) separates the weight penalty (weight decay) from the adaptive step, and is today the default choice for transformers.

The step size over time

A large constant step advances quickly but at the end “dances” around the minimum; a small one is precise but slow. This is why the step size is varied during training:

• warm-up: for a few hundred or thousand steps the step grows from zero to its full value, while Adam’s estimates settle and the network leaves the chaotic initial phase;
• decay: linear, stepwise, or cosine (the step follows half a cosine period down to almost zero). It was also seen in the PPO of the reinforcement learning course: the step decreasing to zero stabilised the final policy.

The comparison on the Rosenbrock function

The function has its minimum at (1, 1) at the bottom of a narrow, curved, banana-shaped valley: at the bottom of the valley the gradient is almost zero, on the walls huge. Starting from (−1.5, 2), we count how many steps are needed to reach :

method	step	steps for
gradient descent	0.001	15,681
momentum ()	0.001	1,484
Adam	0.02	2,286

Momentum is about ten times faster than plain descent (15,681 / 1,484 = 10.6), as predicted by the factor . Adam, with a step that does not depend on the gradient’s scale, arrives in a similar number of steps. On real neural networks the differences are often less clear-cut and depend on the problem: SGD with momentum is still widely used for convolutional networks, Adam and AdamW for transformers and recurrent networks. The choice of step size almost always matters more than the choice of optimiser.

import numpy as np

f = lambda p: (1 - p[0]) ** 2 + 100 * (p[1] - p[0] ** 2) ** 2
grad = lambda p: np.array([-2 * (1 - p[0]) - 400 * p[0] * (p[1] - p[0] ** 2), 200 * (p[1] - p[0] ** 2)])

def optimise(method, eta, steps=20000):
    p, v, m, s = np.array([-1.5, 2.0]), np.zeros(2), np.zeros(2), np.zeros(2)
    for k in range(1, steps + 1):
        g = grad(p)
        if method == "gradient":
            p = p - eta * g
        elif method == "momentum":
            v = 0.9 * v + g
            p = p - eta * v
        else:                                              # Adam
            m = 0.9 * m + 0.1 * g
            s = 0.999 * s + 0.001 * g * g
            m_hat, s_hat = m / (1 - 0.9 ** k), s / (1 - 0.999 ** k)     # bias correction
            p = p - eta * m_hat / (np.sqrt(s_hat) + 1e-8)
        if f(p) < 1e-6:
            return k
    return None

for method, eta in (("gradient", 1e-3), ("momentum", 1e-3), ("adam", 2e-2)):
    print(f"{method:9s}: f < 1e-6 after {optimise(method, eta)} steps")