M5.5 · LESSON 5 OF 6

Momentum

Perform gradient descent updates and explain learning rate, stochastic gradients, momentum and Adam.

THE IDEA

A ball with memory

Plain gradient descent forgets: each step looks only at the current gradient. On the valley that is a problem. Across the valley the gradient flips sign each step (the zigzag), along it the gradient is small but always the same sign. Momentum keeps a velocity v, a decaying sum of past gradients:

v ← β·v + g, then θ ← θ − η·v

β (between 0 and 1) is how much of the old velocity survives. Example: velocity 2, β = 0.9, new gradient 1 gives v = 0.9·2 + 1 = 2.8. Gradients that keep the same sign pile up, so the ball speeds up. Gradients that alternate cancel inside the sum, so the zigzag is damped. If the gradient stayed at g forever, v would approach g/(1 − β), which is 10 g for β = 0.9.

Roll it down the valley

Learning rate0.95Momentum β0.5

iteration 0

Plain gradient descent: position [-5, 2], loss 4.5. At the start. Press Step.

Momentum β = 0.5: position [-5, 2], loss 4.5. At the start. Press Step.

Hollow dot: start. Big dot: where the optimizer is now.

Blue: plain gradient descent. Orange: momentum. The lower chart is y, the position across the valley: oscillation is the wiggle. Press Run 50, then try the two presets and your own values.

Predict first. At learning rate 0.95 on the valley, what does the ball with β = 0.9 do compared with β = 0.5?

Common mistakes

More momentum is not always better. At lr 0.9 and β = 0.9 the ball runs past the valley floor in x and has to come back. Tune β together with η.
Forgetting that the velocity is state. It starts at zero and must be carried from step to step. Rebuilding it each step turns momentum back into plain descent.
Assuming every library writes it the same way. Some put η inside the velocity (v ← βv + ηg, then θ ← θ − v). With a constant learning rate that is the same path; if you change η during training, the two forms differ. Check the documentation of the one you use.

Chapter context

The gradient points toward the steepest local increase under Euclidean geometry. To lower a differentiable loss, a basic optimizer steps in the opposite direction. A local direction says little about how large a safe step can be.

TRY IT YOURSELF

Part 1. Tune the ball

Learning rate0.3Momentum β0

iteration 0

Plain gradient descent: position [-5, 2], loss 4.5. At the start. Press Step.

Momentum β = 0: position [-5, 2], loss 4.5. At the start. Press Step.

Hollow dot: start. Big dot: where the optimizer is now.

To do: Damp the zigzag: with a learning rate of 0.9 or more (below 1) choose a momentum beta above 0 for which |y| is below 0.01 after 20 steps (plain gradient descent leaves 0.23 or more)
To do: Beat plain gradient descent: reach loss 0.05 or lower within 20 steps (the best plain gradient descent needs about 31)
To do: See the cost of momentum: choose settings where the ball overshoots the valley floor along x (x rises above 0.5 within 40 steps) yet the loss is 0.05 or lower at step 60

Part 2. Write the momentum loop

Write momentum_step(theta, vel, grad, lr, beta) using the rule vel_new = beta * vel + grad, then theta_new = theta - lr * vel_new. Return the pair (theta_new, vel_new) as new lists.
Write momentum_descent(grad_fn, theta0, lr, beta, steps): start with zero velocity and return the whole path of steps + 1 positions.

def momentum_step(theta, vel, grad, lr, beta):
    # vel_new = beta * vel + grad ; theta_new = theta - lr * vel_new
    pass

def momentum_descent(grad_fn, theta0, lr, beta, steps):
    # start with velocity 0; return [theta0, theta1, ..., theta_steps]
    pass

print(momentum_step([1.0], [2.0], [1.0], 0.1, 0.9))

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

If velocity is 2, decay is 0.9 and the new gradient is 1, what is 0.9*v+g?

A2.8.B1.9, because the gradient is subtracted from the decayed velocity.C2.9, because the old velocity is multiplied by 0.9 and then 0.9 is added.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

\theta_{t+1}=\theta_t-\eta\nabla L(\theta_t).

Adam uses $m_t=\beta_1m_{t-1}+(1-\beta_1)g_t$ and $v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2$ , coordinatewise. Divide by $1-\beta_1^t$ and $1-\beta_2^t$ before using $\hat m_t/(\sqrt{\hat v_t}+\epsilon)$ in the update.

For the special quadratic $L(\theta)=(\theta-c)^2$ , the residual follows $\theta_{t+1}-c=(1-2\eta)(\theta_t-c)$ . Starting away from c, convergence requires $|1-2\eta|<1$ , or $0<\eta<1$ . This interval depends on this particular curvature and loss normalization.

Worked example

For L=(theta−3)^2, start at 0 and use eta=0.25. The gradients are −6, −3 and −1.5. Three updates produce theta values 1.5, 2.25 and 2.625. Their losses are 2.25, 0.5625 and 0.140625. A fixed learning rate works here because the example is a simple quadratic.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
theta = 0.0
positions = []
for _ in range(3):
    theta -= 0.25 * 2*(theta-3)
    positions.append(theta)
assert positions == [1.5, 2.25, 2.625]
assert isclose((theta-3)**2, 0.140625)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Subtracting loss instead of its gradient is not gradient descent. Squared gradients in Adam are coordinatewise. A zero gradient may identify a saddle point or maximum, not just a minimum.