optimization

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,

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

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, $\alpha = 2/(\lambda_1 + \lambda_2)$, shrinks the error by a factor of $(\kappa - 1)/(\kappa + 1) \approx 0.82$ per step.

Adding momentum #

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

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

What to try #

• Set $\beta = 0$ and raise $\alpha$ until the path zig-zags, then diverges past $\alpha = 0.2$.
• Keep that $\alpha$ and raise $\beta$: the zig-zag damps, and the path reaches the centre sooner.
• Push $\beta$ towards $1$: 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.