Weblog Writing & Notes
All posts
1 min read

Why momentum helps gradient descent

Sample Distill-style article: prose and maths around an explorable figure built with Distill's figure element and d3.

sampleoptimizationdistill

Gradient descent is the workhorse of machine learning, and it has a well-known weakness. On a long, narrow valley it bounces from wall to wall while creeping along the floor. Adding momentum, a running memory of past gradients, is a one-line change that often makes it several times faster. This article shows why, on the simplest valley there is.

A quadratic valley #

Take a bowl stretched along one axis,

f(w)=12(λ1w12+λ2w22),λ1=1, λ2=10. f(w) = \tfrac12 \left( \lambda_1 w_1^2 + \lambda_2 w_2^2 \right), \qquad \lambda_1 = 1,\ \lambda_2 = 10.

Its gradient is ∇f(w)=(λ1w1, λ2w2)\nabla f(w) = (\lambda_1 w_1,\ \lambda_2 w_2), and the ratio κ=λ2/λ1=10\kappa = \lambda_2 / \lambda_1 = 10, the condition number, measures how narrow the valley is. Plain gradient descent takes steps of size α\alpha:

wk+1=wk−α∇f(wk). w_{k+1} = w_k - \alpha \nabla f(w_k).

A step large enough to make progress along the shallow direction overshoots in the steep one, and a step small enough for the steep one crawls along the shallow one. The best fixed step, α=2/(λ1+λ2)\alpha = 2/(\lambda_1 + \lambda_2), shrinks the error by a factor of (κ−1)/(κ+1)≈0.82(\kappa - 1)/(\kappa + 1) \approx 0.82 per step.

Adding momentum #

Momentum keeps a velocity zz that accumulates gradients and decays by a factor β\beta each step:

zk+1=βzk+∇f(wk),wk+1=wk−αzk+1. z_{k+1} = \beta z_k + \nabla f(w_k), \qquad w_{k+1} = w_k - \alpha z_{k+1}.

Oscillations across the valley point in alternating directions and cancel in zz; the small, consistent pull along the floor adds up. With the best choice of α\alpha and β\beta the error shrinks by (κ−1)/(κ+1)≈0.52(\sqrt\kappa - 1)/(\sqrt\kappa + 1) \approx 0.52 per step: about 11 steps instead of 35 to reduce it a thousandfold.

Contours of a narrow quadratic valley, with a red plain gradient descent path and a blue path with momentum
Gradient descent without momentum (red) and with it (blue), on the valley above. With JavaScript, this figure has sliders for the step size and momentum.

What to try #

  • Set β=0\beta = 0 and raise α\alpha until the path zig-zags, then diverges past α=0.2\alpha = 0.2.
  • Keep that α\alpha and raise β\beta: the zig-zag damps, and the path reaches the centre sooner.
  • Push β\beta towards 11: too much memory and the path overshoots and circles.

Further reading #

This sample follows the style of Distill articles and borrows its example from Gabriel Goh’s Why Momentum Really Works, which analyses the same valley in far more depth.