Optimizers, step by step


Optimizer
Landscape
Learning rate
Batch size (SGD)
Move
The loss for every choice of w and b: darker is worse, × is the best line. Each dot is one step. Gray arrow: straight downhill. Colored arrow: the step the optimizer will take next. Click anywhere to start from there.
The 20 data points and the line for the current w and b. Dashed: the best line. With SGD, the ringed points are the batch for the next step.
The loss curve: how far the loss (MSE) is above the best line's MSE, after every step. The scale is logarithmic, so each gridline is 10× smaller than the one above. Dashed: the bottom, 0.01 above the best.

The whole idea.
The model on the page is a straight line, y = w·x + b. It has two numbers you can change:
• w is the slope, how steep the line is.
• b is the intercept, where the line crosses the y axis.
Choose a value for w and a value for b and you get exactly one line. Change either one and you get a different line. The loss is the average squared miss over the 20 points. Training means finding the lowest spot, the ×.

An optimizer cannot see the whole map. All it knows is where it stands and the gradient: which way is uphill there, and how steeply. Every optimizer on this page is a rule for turning that one piece of information into a step.

The four rules.
Gradient descent: step = −learning rate × gradient, with the gradient worked out from all 20 points.
SGD (stochastic gradient descent): the same rule, but each step's gradient comes from a small random batch of points. It is cheaper and noisier. Real datasets have millions of points, so this is what is actually used.
Momentum: keeps a running velocity, 0.9 × the old velocity + the gradient, and steps along that. Directions that agree step after step build up speed; directions that flip back and forth cancel out.
Adam: momentum, plus a separate step size for each weight. It divides each weight's averaged gradient by the typical size of that weight's recent gradients, so a steep direction and a shallow one both move about one learning rate per step.
Momentum and Adam here use all 20 points, so you see only their rule. In real training they use batches too.

Things worth trying:
1. Gradient descent on the Round bowl. At the default learning rate of 0.03 it takes 54 steps to reach the bottom (a loss within 0.01 of the best). Drag the learning rate to 0.5 and it gets there in exactly one step. At 1 it jumps clean over the bottom to the far side, then back, forever.

2. Gradient descent on the Stretched bowl. This is the same data with x measured on a five times bigger scale, so the loss is far steeper in w than in b. The path zigzags across the narrow direction while creeping along the long one. The steep direction sets the speed limit: 0.04 is the highest learning rate that does not blow up here, against 1 on the round bowl.

3. Adam on the same stretched bowl. Its first step moves exactly 0.2 in both w and b, however steep each one is, so it does not zigzag across the narrow direction. Try any learning rate on the slider, on any of the three landscapes: Adam never blows up. That forgiveness, more than raw speed, is why Adam is the usual first choice.

4. The Tilted valley: the same data shifted to sit around x = 3. Now w and b are tangled, and the valley runs at a slant. Gradient descent drops to the valley floor quickly, then crawls along it: 455 steps. Momentum builds up speed along the floor and takes 120. Adam takes 165: its per-weight step sizes only line up with valleys that run along an axis, and this one does not.

5. SGD on the round bowl, with the batch size changed. The ringed points are the batch for the next step, and the step arrow wobbles away from straight downhill. SGD never settles exactly, it jitters around the bottom. Measured over steps 2001 to 3000, the typical distance above the best loss was 0.0050 with 1 point, 0.0013 with 4, and 0.0002 with 10. A batch of 1 on the stretched bowl blows up after 433 steps: one far-out point can be much steeper than the average.

6. Race all four from the same start, each at its own learning rate. Steps to reach the bottom: round bowl 54, 56, 14, 15 (gradient descent, SGD, momentum, Adam); stretched bowl 32, 35, 44, 45; tilted valley 455, 472, 120, 165. No optimizer wins everywhere.

7. Watch the loss curve during the race on the round bowl. Over 500 steps, counting only while the loss is still more than 0.0001 above the best: gradient descent's loss never goes up once. Momentum's goes up on 34 steps and Adam's on 32; each bump is an overshoot past the × on the map. SGD's goes up on 225 steps and never settles: after step 100 it keeps bouncing just under the dashed line, as high as 0.0087, while the other three sit on the floor. That is the batch noise.