How do you choose an optimiser and a learning-rate schedule?

The rate matters more than the optimiser

If you are given one hyperparameter to tune, tune the learning rate. It sets the size of every step, and being wrong about it by an order of magnitude ruins a run regardless of which optimiser produced the direction. Most reported wins from swapping optimisers are wins from the new optimiser happening to suit the rate that was left unchanged, so a candidate who leads with the rate is already answering better than one who leads with a list of names.

SGD with momentum

Plain stochastic gradient descent steps directly down the gradient of the current mini-batch. Because that gradient is a noisy estimate, successive steps disagree, and the path oscillates across directions where the loss surface is steep while crawling along directions where it is shallow.

Momentum accumulates an exponentially weighted average of past gradients and steps along that instead. Consistent components reinforce and inconsistent ones cancel, so progress along a long shallow valley accelerates and the oscillation across it damps. It is one extra hyperparameter, conventionally around 0.9, and it is close to free.

What you get for the tuning effort is a well-understood method with a single global step size, and a long-standing observation in the vision literature that carefully tuned SGD with momentum can generalise better than adaptive methods on convolutional architectures. Treat that as a reason to try it when you have the compute for a proper sweep, not as a law.

Adaptive methods

Adam maintains, per parameter, a running average of the gradient and of the squared gradient, and divides the step by the square root of the second. A parameter whose gradients have been consistently large gets smaller steps; one whose gradients have been small gets larger ones. The effect is that the per-parameter step size becomes roughly scale-invariant, which is why Adam tolerates a badly scaled problem, sparse features that receive gradient only occasionally, and architectures where different blocks have very different gradient magnitudes.

The practical value is that it works acceptably at its defaults, so it is the right first choice when you have one shot rather than a sweep. AdamW is the version to reach for, because it applies weight decay directly to the parameters rather than adding an L2 term to the loss. In plain Adam that L2 term flows through the adaptive scaling, so the effective amount of decay ends up depending on each parameter's gradient history — decoupling it makes the regularisation mean what the number says. Adaptive methods also cost memory, holding two extra tensors per parameter, which is a real constraint at large model sizes.

The rough rule is that transformers and anything with embeddings are trained with AdamW almost universally, convolutional vision models are a genuine contest, and if you are fine-tuning rather than training from scratch you want a small constant-ish rate more than you want a clever optimiser.

Warmup, then decay

Two schedule ingredients do most of the work, and each answers a different problem.

Warmup ramps the rate from near zero to its peak over the first few hundred or few thousand steps. Its justification is specific: at initialisation the adaptive optimiser's second-moment estimate is built from a handful of batches and is therefore unreliable, so early steps can be enormous in exactly the directions the estimate has underestimated. One such step can push the network into a region it never recovers from. Warmup also protects a fine-tune from having its pretrained weights destroyed before the head has learned anything, and it becomes more important as batch size and peak rate rise.

Decay lowers the rate as training proceeds. Early on you want large steps to cover ground; later you want small ones to settle rather than skate around the minimum. Cosine decay to near zero over the planned number of steps is the common default and needs no tuning beyond knowing the total. Step decay, dropping the rate by a factor at fixed milestones, is older and produces the characteristic staircase where the loss drops sharply at each cut. Reducing on plateau is reactive and useful when you cannot predict the run length, at the cost of two more hyperparameters.

Note that a cosine schedule is defined against a total step count, so shortening or extending a run silently changes the schedule, and comparing two runs of different length under it compares two different experiments.

Read the failure before changing anything

The loss curve distinguishes the two main failures, and they demand opposite responses.

A diverging run has a loss that rises, oscillates with growing amplitude, or becomes NaN. The rate is too high, or warmup is absent, or the gradients are exploding. The response is to lower the peak rate by a factor of three to ten, add or lengthen warmup, and clip the gradient norm. Lowering the rate is the first move because it is one number and it either fixes the run or rules the cause out.

A plateauing run has a loss that falls and then goes flat. Flat is much more ambiguous than rising. It may be a rate too low, in which case raising it moves the curve immediately. It may be a rate too high in a subtler way, where the steps are large enough that the run bounces around a basin without descending — the tell is that the loss is noisy around its plateau rather than smoothly flat, and decaying the rate produces a visible drop. It may be a vanishing gradient, in which case the per-layer gradient norms will show it and no optimiser change will help. Or it may be a model at its capacity, which the training loss cannot tell you at all: check whether training and validation loss have converged to each other, because a plateau with a large gap is a regularisation problem and a plateau with no gap is a capacity or data problem.

# One warmup phase then cosine decay, composed rather than hand-rolled.
sched = torch.optim.lr_scheduler.SequentialLR(
    opt,
    schedulers=[
        torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.01, total_iters=500),
        # T_max is the remaining steps, so changing run length changes the schedule.
        torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps - 500),
    ],
    milestones=[500],
)

Pick AdamW with a short warmup and cosine decay, then spend your budget finding the peak rate. And before you change the optimiser, decide from the shape of the curve whether you are looking at instability or at a model that has nothing left to learn — the two look nothing alike and are treated as if they do.