Momentum, step by step: give gradient descent a memory, and it rolls past the traps that stop it.
Plain gradient descent looks only at the slope under its feet. Every step is
w ← w − η ∇L(w)
and nothing else. That makes it simple, but it also makes it forgetful: the moment the ground goes flat, it stops. Momentum fixes this with one extra number per weight, a velocity that remembers the slopes it has already rolled down. The examples below follow the CS878 lectures on gradient descent and scaling the optimization, and every number comes from running the update rule.
The trap
Take one weight and a loss with two dips:
L(w) = 0.12 (w − 2.5)² − 0.55 · exp(−(w + 1)² / 0.12)
Start at A, w = −3.5, where the loss is 4.32. There is a shallow dip B at w = −0.904 (loss 0.881) and the real minimum C at w = 2.5 (loss 0). Between them sits a bump only 0.13 higher than B.
Plain gradient descent with η = 0.1 takes big steps on the steep slope, smaller and smaller ones as it nears B, and by step 21 it has arrived. At B the gradient is zero, so every step after that is zero too. At step 50, at step 100, it is still in B, at a loss of 0.88.
Plain gradient descent stops in the shallow dip B: the gradient is zero there, so learning stops.
Add a velocity
Momentum keeps a velocity v and moves along it:
- v_t = β v_(t−1) + ∇L(w_t)
- w_(t+1) = w_t − η v_t
With β = 0.9, each step keeps 90% of the old velocity and adds the new gradient. (This is the form PyTorch uses. Lecture notes often write v = βv + (1 − β)∇L, which only rescales the learning rate.)
Run it from A with the same η = 0.1. By step 6 the ball is near B and its velocity is about 6. The gradient at B is zero, but the velocity isn't, so the ball keeps going: it clears the bump at step 7, overshoots all the way to w = 4.71, swings back without ever returning to the bump, and settles in C by step 89.
Momentum carries the ball over the bump; plain GD's ghost is still stuck in B.
What v remembers
Unroll the velocity and it becomes a weighted sum of every gradient so far:
v_t = ∇L_t + β ∇L_(t−1) + β² ∇L_(t−2) + …
The latest gradient counts fully, the one before 90%, then 81%, 73%, 66% and so on. The weights add up to 1/(1 − β) = 10, so momentum remembers about ten steps. With β = 0.1 they would be 1, 0.1, 0.01: it forgets almost at once. In practice β is usually around 0.9.
Eight bars, heights β^k, sum 1/(1 − β) = 10: gradients that agree stack toward 10; gradients that flip cancel, staying within ±0.53.
Agree or cancel
That memory does two different things depending on the gradients:
- When they agree (+1, +1, +1 …), they stack: v grows 1, 1.9, 2.71, 3.44 … toward 1/(1 − β) = 10. The step grows up to ten times.
- When they flip sign (+1, −1, +1 …), they cancel: v goes 1, −0.1, 0.91, −0.18, 0.84, −0.25 … and settles at about ±1/(1 + β) ≈ 0.53. The step stays small.
A narrow valley
Now two weights and a valley that is steep across and almost flat along: L = ½(w₁² + 50 w₂²), starting at (−4, 1), where the loss is 33. Count the steps until the loss drops below 0.001, with η = 0.015:
- Plain GD: 298 steps. A bigger rate doesn't help: at η = 0.045, each step multiplies the across direction by 1 − 0.045 · 50 = −1.25 and it blows up. So plain GD has to crawl along the floor.
- Momentum (β = 0.9): 73 steps. The along direction keeps agreeing, so it builds speed there. But the across direction overshoots and swings from side to side.
The narrow valley: plain GD crawls along the floor; momentum is faster but swings across.
Nesterov: look before you step
Nesterov momentum changes one thing: it takes the gradient where the velocity is about to carry the weight, not where the weight is now.
v_t = β v_(t−1) + ∇L(w_t − η β v_(t−1))
If that look-ahead point is already uphill, the gradient there points back, and the step brakes before it overshoots. On the same valley, Nesterov needs 18 steps, with almost no swing.
Nesterov looks ahead: the gradient at the look-ahead point brakes the swing. 298 · 73 · 18.
Try it in PyTorch
Every weight in a network gets its own velocity, with the same shape as the weight. In PyTorch it is one line:
opt = torch.optim.SGD(model.parameters(), lr=0.015,
momentum=0.9, nesterov=True)
And the velocity itself, with gradients that flip sign:
v, out = 0.0, []
for g in [1, -1, 1, -1, 1]:
v = 0.9 * v + g
out.append(round(v, 2))
print(out)
It prints [1.0, -0.1, 0.91, -0.18, 0.84]. Change every g to 1 and v climbs toward 10.
Key takeaways
- Plain gradient descent has no memory: it stops wherever the slope is zero, even in a shallow dip.
- Momentum keeps a velocity, v = βv + ∇L, and steps along it, so it can roll through flat spots and over small bumps.
- v is a sum of past gradients weighted by β^k, about 1/(1 − β) = 10 steps of memory at β = 0.9.
- Agreeing gradients stack (up to 10×); flipping ones cancel. That is why momentum speeds up along valleys.
- Nesterov takes the gradient at the look-ahead point, so it brakes before overshooting: 298 → 73 → 18 steps in the narrow valley.
This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.
Top comments (1)
This is a nice example of why a simple model needs a stated validity range. I’d make the dimensionless threshold and assumptions visible next to the animation, then show one counterexample where the approximation fails. That helps readers learn when the intuition transfers rather than treating the visual as a universal rule.