Momentum and Adam: gradient descent with a memory

Plain gradient descent has no memory: each step measures the slope underfoot, moves against it and forgets it. In the gradient descent lesson that made it zigzag across a narrow valley while it crept along the floor. A network's loss has millions of directions, some far steeper than others, so this is the normal case.

A ball rolling down the valley does better because it keeps its speed: pushes along the floor add up, and pushes from side to side cancel. That is momentum. Shrinking the steps of directions that have been steep lately is RMSProp. Doing both is Adam, the usual default for training neural networks.

velocity v

Momentum's memory: past gradients added up, each shrunk by β for every step since. Momentum steps against v, as plain descent steps against g.

decay β

How much of the memory survives each step. At β = 0.9 it spans about 1/(1 − β) = 10 recent gradients.

running average

An average that forgets slowly: m ← βm + (1 − β)g. RMSProp keeps one of g², Adam one of g and one of g².

bias correction

Adam's averages start at 0, so they are too small at first. Dividing by 1 − βᵗ at step t scales them back up.

The race below runs four optimisers from one start for 60 steps. Each step, each one measures the gradient $g$ where it stands (derivatives and gradients shows how) and turns it into a step. Plain descent steps against it, $x \leftarrow x - \eta g$. Momentum shrinks its velocity by $\beta$, adds $g$, and steps against the velocity:

RMSProp keeps a running average $s$ of $g^2$, coordinate by coordinate (weight by weight, in a network), and divides each coordinate's step by $\sqrt{s}$:

Adam keeps a running average $m$ of $g$ as well as RMSProp's $s$, corrects both for starting at 0, and steps by their ratio:

with $\hat m = m / (1 - \beta_1^t)$ and $\hat s = s / (1 - \beta_2^t)$ at step $t$; the tiny $\epsilon$ only prevents a division by 0. The race uses $\beta = \beta_1 = \rho = 0.9$ and $\beta_2 = 0.999$, and its own $\eta$ per optimiser and surface, so read it as a picture of behaviour, not a benchmark. On the narrow valley $f = x^2/20 + 5y^2$ the gradient, $(x/10,\ 10y)$, changes 100 times faster across the valley than along it.

Step once: plain descent ($\eta = 0.18$) jumps 2.7 across the valley and 0.063 along it, while Adam moves exactly its $\eta$, 0.1, in each coordinate: at $t = 1$ the corrections give $\hat m = g$ and $\hat s = g^2$, so the step is $\eta \, g/|g|$.

Momentum's velocity is a sum, not an average: on a steady slope it grows towards $g/(1 - \beta)$, ten times $g$ at $\beta = 0.9$, so momentum gets a smaller $\eta$ (0.03 here). Along the floor the gradient keeps its direction, so momentum builds up speed there: after 60 steps it is at $x = -0.41$, plain descent at $-1.18$. A gradient that flips sign every step adds $+g$ and $-g$ in turn, so the velocity settles at about $g/(1 + \beta)$. The price is inertia: momentum overshoots the floor (its $f$ rises at step 3), then swings across, a little less each time.

RMSProp divides by how big each coordinate's gradients have been lately, so while a gradient holds steady, $g/\sqrt{s}$ is about $\pm 1$ and the step about $\eta$, steep or gentle. Adam's $m$ is momentum written as an average, about the size of $g$, so dividing by $\sqrt{\hat s}$ again makes each step about $\eta$: Adam walks down the steep wall in steps of about 0.1 instead of leaping across, though its momentum still carries it past the floor in slow swings.

On the two-valleys surface, $f = (x^2 - 1)^2 + x/3 + y^2$, the left valley is lower than the right, with a ridge near $x = 0$ between them. All four start at (1.8, 1.2), on the steep outer slope of the right-hand valley. Predict who will be lowest after step 1 and after every tenth step: think about each first step, and who could gather enough speed to cross the ridge.

Here is momentum on $f(x) = x^2$, whose gradient is $2x$. With $\beta = 0$ the velocity is just the latest gradient, so one function runs both: momentum with $\beta = 0$ is plain descent. With $\eta = 0.25$, $\beta = 0.5$ and a start at $x = 8$, every $x$ stays a whole number.

def descend(label, beta, steps):
    x = 8
    v = 0
    for t in range(steps):
        g = 2 * x
        v = beta * v + g
        x = x - 0.25 * v
        print(label + ' ' + str(int(x)))

descend('plain', 0, 3)
descend('momentum', 0.5, 3)

On a round bowl plain steps already point straight at the bottom, so momentum has no zigzag to cancel, and its speed only carries it past.

Real gradients come from small random batches, so they are noisy, and the running averages smooth that too. SGD with momentum is still widely used beside Adam, and sometimes generalises a little better. One setting never goes away.

• Plain gradient descent uses only the current gradient; momentum, RMSProp and Adam add a memory of past ones.
• Momentum steps against $v \leftarrow \beta v + g$: steady gradients build it up to about $1/(1 - \beta)$ times the gradient (10 at $\beta = 0.9$), flipping ones cancel, and its speed makes it overshoot.
• RMSProp divides each coordinate's step by the root of its running average of $g^2$, so steep and gentle directions both step about $\eta$.
• Adam is momentum as an average, plus RMSProp, plus bias correction.
• None of them sees the surface: they can settle in a local minimum, and $\eta$ still has to be chosen.
• Next: neural networks and the neural network playground, which trains with momentum. Backpropagation shows where the gradients come from.