Optimizers: from SGD to Adam

Backpropagation tells us, for every parameter, which way is downhill. It does not tell us how far to step. That decision belongs to the optimizer — the rule that turns gradients into parameter updates. The choice of optimizer and its learning rate can make the difference between a model that trains in an hour and one that never converges.

The trouble with one step size

Plain (stochastic) gradient descent uses the same learning rate for every parameter:

\theta \leftarrow \theta - \eta \, \nabla L(\theta)

Real loss surfaces are rarely round bowls. They are full of long, narrow valleys: steep across, gentle along. A learning rate small enough to be stable across the steep walls is far too small for the gentle floor, so progress is a slow zig-zag.

AnimationThree ways to walk down a narrow valley

Loading animation…

A long, narrow valleyThe contours show a loss that is 20 times steeper across the valley than along it. Every optimizer starts at the same point and takes 30 steps.
Gradient descent zig-zagsThe gradient points mostly across the valley, so each step overshoots the floor and bounces back. After 30 steps it has covered little distance along the valley.
Momentum builds speedMomentum averages recent gradients. The back-and-forth components cancel while the consistent downhill component accumulates, so the ball rolls along the valley — overshooting a little at the end.
Adam rescales each directionAdam divides each parameter’s step by the recent size of its gradients. The steep direction is tamed and the gentle direction sped up, so the path heads straight for the minimum.
Same gradients, different journeysAll three used exactly the same gradients. The optimizer’s rule for turning gradients into steps is what made the difference.

Four optimizers you will meet

SGD: Step against the gradient with a fixed learning rate. Simple and, with a good schedule, still excellent for many vision models.
Momentum: Keep a velocity v ← βv + ∇L (β ≈ 0.9) and step along it. Speeds up consistent directions and damps oscillations.
RMSProp: Keep a running average of squared gradients and divide each step by its square root, so parameters with large gradients take smaller steps.
Adam: Momentum plus RMSProp-style scaling, with a correction for the first few steps. The default choice for most deep-learning work.
AdamW: Adam with weight decay applied separately from the gradient update. The standard optimizer for training transformers.

Race them in 3D

The surface below is a curved valley — a version of the Rosenbrock function, a classic test for optimizers. The minimum sits at the far end of a banana-shaped trough. Start the race, then change the learning-rate multiplier and the starting point.

3D modelOptimizer race on a curved valleyDrag to rotate · start the race

Loading 3D model…

Each optimizer has its own base learning rate, and the multiplier scales all four. The loss readouts update live; an optimizer that jumps off the edge of the map has diverged.

Things to look for:

SGD slides down the steep wall quickly, then crawls along the curved floor.
Momentum carries speed around the bend, overshooting and correcting.
RMSProp and Adam adapt their step sizes and reach the minimum first.
Start from the steep wall: plain SGD copes at 1×, but at 2× its first step is so large that it flies off the map — the 3D version of a loss that suddenly explodes. Adam, meanwhile, builds up momentum in the wrong direction and arrives last: no optimizer wins everywhere.

Learning-rate schedules

Even with a good optimizer, the best learning rate changes during training. Early on, large steps make fast progress. Later, the same large steps bounce around the bottom of the valley and keep the loss from settling. So we schedule the learning rate.

InteractiveLearning-rate schedules and their effect on training

Loading visual…

A simplified simulation: higher rates learn faster but settle on a noisier, higher floor, and very high rates at the start cause spikes.

Step decay divides the rate by 10 at fixed epochs — you can see the loss drop at each step.
Cosine decays smoothly to zero and is a strong default.
Warmup + cosine starts small, ramps up, then decays. Set the peak rate above 0.6 and compare cosine with and without warmup: warmup avoids the early spike.

Reading the loss curve

The shape of the training-loss curve tells you what to change:

Flat from the start — the learning rate is too small, or gradients are not flowing (check for bugs and dead activations).
Spiky or exploding — the rate is too large; lower it, add warmup or gradient clipping.
Falls, then plateaus high — try decaying the rate; a schedule usually unlocks another drop.
Training loss falls but validation loss rises — the optimizer is working fine; the model is overfitting. That is the subject of the next lecture.

In PyTorch, an optimizer and a schedule take three lines:

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)

for batch in loader:
    loss = loss_fn(model(batch.x), batch.y)
    optimizer.zero_grad()
    loss.backward()        # backpropagation fills the gradients
    optimizer.step()       # the optimizer turns gradients into updates
    scheduler.step()       # the schedule adjusts the learning rate

In this lecture you will