Optimizers

The step size is the hyperparameter

Plain gradient descent does one thing: move each weight a little way against its gradient. The gradient is the direction in which the loss rises fastest. The learning rate sets how far each step goes. Too small and the run takes forever. Too large and it bounces across the valley, overshoots, and eventually diverges. The stable range depends on the curvature of the loss, which is how sharply it bends. One global rate is therefore always a compromise.

A useful mental model is a ball rolling down a valley. The gradient gives the slope at the current point, not the shape of the valley. A step that is right on a steep wall is tiny on a flat floor. A step that is right on the floor overshoots the wall. Every optimiser after plain descent is an attempt to use more than the local slope.

Curvature is the missing information. A quadratic with a steep direction and a shallow one has a condition number. That number is the ratio of the two curvatures, and gradient descent needs a number of steps proportional to it. The same rate can therefore crawl in one direction and diverge in the other.

Momentum remembers the direction

Momentum keeps a running average of past gradients and steps along that average. Consistent directions accumulate and grow. Directions that alternate cancel out. In a narrow ravine, the gradient points across the ravine on alternating steps and along it only slightly. Momentum damps the zigzag and accelerates the useful part.

The momentum coefficient is usually 0.9, which gives the average a memory of about ten steps. Nesterov momentum evaluates the gradient at the point the momentum is about to reach. That look-ahead often gives a slightly better correction. Both variants add one buffer per parameter, and both are sensitive to the interaction between the coefficient and the learning rate.

Momentum does not remove the need to tune the rate. It changes the effective step to roughly one over one minus beta, times the rate. A coefficient of 0.9 multiplies the step by ten in a consistent direction. The rate usually comes down when momentum is added.

AdamW with a schedule

import torch
from torch import nn

model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(100):
    optimizer.zero_grad()
    loss = model(torch.randn(8, 10)).sum()
    loss.backward()
    optimizer.step()
    scheduler.step()

Adam scales each parameter

Adam tracks two running averages per parameter: the gradient and the squared gradient. The first is momentum. The second estimates the scale of the gradient, and dividing by its square root turns every update into roughly the same size. A parameter with a consistently tiny gradient gets a larger relative step. One with a huge gradient gets a smaller one.

The running averages start at zero, so early estimates are biased toward zero. Adam corrects for that with a bias-correction term. The term is large for the first few steps and fades to one. Without it the first updates are far too small. An implementation that omits the correction then looks like it is warming up for no reason.

Epsilon keeps the division finite when the second moment is zero. Its default is usually 1e-8. That is fine in float32 and too small in float16, so mixed precision often raises it. The detail is boring until a run produces NaN on the first step.

Decay, schedules and what to tune

Weight decay pulls every weight toward zero. That is a prior: a belief, built in before the data arrives, that small weights are more likely. In Adam it was long implemented by adding the penalty to the gradient. That interacts badly with the adaptive scaling. AdamW decouples the two. The decay is applied directly to the weights, and the gradient step is left alone. The difference is small in theory and visible in practice.

A schedule changes the rate over training. Warmup avoids a large step while the estimates are noisy. Cosine decay spends more time near a small rate at the end, and step decay does the same in discrete drops. The batch size interacts with all of it. A larger batch gives a less noisy gradient, so the linear scaling rule raises the rate in proportion.

Tune the learning rate first, on a short run. Look at the loss curve rather than the final number. Then choose a schedule, then the weight decay. Tuning all three at once produces a result you cannot explain and a configuration you cannot transfer.

Early stopping is the other schedule, and it costs nothing. Keep the validation loss, save the weights when it improves, and stop after a patience of epochs without a new best. Patience is the number of epochs you wait for an improvement. The final model is then the best one you saw, not the one that happened to be last.