Momentum, Adam, and what the optimiser is doing for you

What plain SGD struggles with

Gradient descent follows the slope where it currently stands. Two situations make that inefficient.

Ravines. A surface that is steep across and shallow along: the long narrow valley from module 2. The gradient points mostly across the valley, so the optimiser zigzags from wall to wall and makes slow progress along the floor.

Uneven scales. One parameter needs large steps and another needs small ones, and a single learning rate has to serve both.

Momentum

Keep a running average of recent gradients and step in that direction.

v = beta * v + gradient          # beta typically 0.9
w = w - learning_rate * v

The zigzag components across the ravine alternate in sign and cancel in the average. The consistent component along the floor accumulates. The result is faster progress in the useful direction and damped oscillation, and it also carries the optimiser through small bumps and flat regions where the raw gradient is near zero.

With beta = 0.9, the effective step is roughly ten times the current gradient once the average has built up, which is why momentum with a given learning rate is considerably more aggressive than plain SGD with the same rate.

Nesterov momentum evaluates the gradient at the position the momentum is about to take you rather than where you are: a lookahead that damps overshoot slightly. A small, real improvement.

Adam

Adam adds a per-parameter learning rate on top of momentum. It keeps two running averages: one of the gradient, and one of the gradient squared. Then it divides the step by the square root of the second.

m = b1 * m + (1 - b1) * g          # average gradient
v = b2 * v + (1 - b2) * g * g      # average squared gradient
w = w - lr * m / (sqrt(v) + eps)

The effect: a parameter that consistently receives large gradients has a large v, so its steps are scaled down; one receiving small gradients gets steps scaled up. Every parameter ends up with a step size suited to its own history, which removes most of the sensitivity to feature scaling and to how the layers happen to be sized.

Defaults b1 = 0.9, b2 = 0.999, lr = 0.001 work across an enormous range of problems, which is why Adam is what almost everyone uses by default.

Adam’s costs

Memory. Two extra values per parameter. For a model with a billion parameters, the optimiser state is twice the size of the model, and it is a large part of why training needs so much more memory than inference.

Generalisation. On image classification specifically, well-tuned SGD with momentum often generalises slightly better than Adam. The gap is small and the tuning effort is not, and Adam remains the right default when you do not have time to tune.

Weight decay is broken in plain Adam. L2 regularisation added to the loss gets divided by the same sqrt(v) term, so parameters with large gradients receive weaker decay, which is the opposite of intended. AdamW applies the decay directly to the weights, outside the adaptive scaling, and it is what you should use.

torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

AdamW is the standard for training transformers and should be your default over Adam. The fix is one letter and it matters.

Schedules

The learning rate should not stay constant, for the reasons in module 2.

Warmup: increase from near zero over the first few hundred to few thousand steps. Essential for transformers, where early gradients on freshly initialised weights are unreliable and a large early step does lasting damage.

Cosine decay: fall smoothly to near zero over training. The current standard.

One-cycle: warm up to a high rate, then decay below the starting point. Trains surprisingly fast and is built into fastai and PyTorch.

sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-3,
                                            total_steps=n_steps)

Gradient clipping

One more safety measure. If the gradient’s norm exceeds a threshold, scale it down.

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

This prevents a single unusual batch from taking a catastrophic step and destroying weights that took hours to train. It is standard for recurrent networks and transformers and costs essentially nothing. If your loss occasionally spikes to NaN after training normally for a while, this is usually the fix.

What to actually use

AdamW at 3e-4, cosine decay with warmup, gradient clipping at 1.0. That combination trains almost anything without drama. Reach for SGD with momentum only if you are chasing the last fraction of a point on a vision benchmark and have the compute to tune it.