Adam and AdamW: A Step Size for Every Coordinate

Chapter 12 · Mathematics for AI and Machine Learning · Advanced · 2:14

An interactive live film: on a narrow bowl, gradient descent is capped by the steep direction while Adam divides each coordinate by its own gradient scale. Set alpha and beta_1 and beat the best gradient descent in forty steps.

Narration

Here is a long, narrow bowl: the loss is one half of the sum of x one squared and twenty five times x two squared. Across the valley it curves twenty five times more than along it.

Gradient descent uses one step size for every coordinate, and the steep direction caps it: above two over twenty five, the steps across the valley grow and the iterates diverge. So the red path, at the best safe step, crosses the valley back and forth and creeps along the floor.

The rest of the narration is in the book's Companion.