Optimizers that Train Faster: Momentum, Adam & Schedules

Plain SGD's struggles: ravines, plateaus, and zig-zags

In the previous guide you met the training loop: forward pass, loss, gradients, and a step downhill. The engine doing that stepping was stochastic gradient descent — at each step it looks at one mini-batch, computes the gradient (the direction of steepest increase), and moves a small amount in the opposite direction. That small amount is set by the learning rate. It is simple, honest, and it works. But on the bumpy, twisty landscapes of real neural networks it can be painfully slow and shaky. This section makes you feel exactly why, so the rest of the guide lands as relief rather than as more formulas.

Picture the loss surface as a physical landscape: every setting of the weights is a location, and the height is the loss. Training is rolling a ball downhill to find the lowest valley. The trouble is the shape of that valley. A very common shape is a ravine: a long, narrow canyon that is steep across (the walls) but almost flat along its floor (the direction you actually want to travel). Plain SGD always steps along the gradient, and in a ravine the gradient points mostly toward the nearest wall, not down the gentle floor.

So what happens? The ball overshoots the floor, hits the opposite wall, gets pushed back, overshoots again — it zig-zags from wall to wall, bouncing across the canyon while creeping only a tiny bit forward along it. To stop the bouncing you would have to shrink the learning rate, but then your forward crawl along the floor becomes glacial. You are stuck choosing between 'shaky and overshooting' or 'stable but crawling'. That is the ravine problem.

Two more failure modes pile on. On a plateau — a wide, nearly flat region — the gradient is tiny in every direction, so each step is microscopic and the ball barely moves; training stalls for many iterations. And because each step uses only a noisy mini-batch rather than the whole dataset, the gradient itself wobbles from step to step, adding a jitter on top of everything. Steep zig-zag + flat crawl + noisy jitter is exactly the combination that makes plain SGD slow to train. The fix is to give the ball some memory and some adaptiveness — which is what the rest of this guide builds.

Momentum: giving gradient descent inertia

Here is the first cure. Instead of stepping by whatever gradient we see this instant, we keep a running 'velocity' that accumulates past gradients and step by that. This is exactly what a heavy ball does when it rolls downhill: it does not stop and re-aim at every pebble. It builds up speed in whatever direction it has consistently been pushed, and its inertia carries it straight through small bumps. That accumulated velocity is what we call momentum.

Why does this fix the ravine zig-zag? Think about the two directions separately. Across the canyon, the gradient flips sign every step (it points right, then left, then right...) because the ball keeps crossing the floor. When you add those alternating pushes into a running velocity, they largely cancel out — left and right average to roughly zero. But along the canyon floor, the gradient always points the same way (downhill). Those aligned pushes add up, so the velocity in the useful direction grows and grows. The result: less side-to-side bouncing, more steady forward speed. On a plateau the same accumulation keeps the ball coasting forward even when the local slope is nearly flat.

The momentum update: first blend old velocity with the new gradient, then step by the velocity.

Let us dissect it symbol by symbol. The vector v_t is the velocity — our running memory of recent gradients (it starts at zero before the first step). \nabla L(\theta_t) is the current mini-batch gradient at the current weights \theta_t — the fresh slope under the ball's feet. \beta (beta) is the momentum coefficient, a number between 0 and 1; it acts like friction, deciding how much of the old velocity survives into the next step. \eta (eta) is the learning rate, the overall step size. The first equation says 'new velocity = (a fraction \beta of the old velocity) + (the fresh gradient)'. The second says 'move the weights by \eta times that velocity'. When \beta = 0 the velocity is just today's gradient and we recover plain SGD; turning \beta up adds inertia.

Now the famous intuition for the number. With coefficient \beta, the velocity is an exponentially weighted average that effectively remembers about 1/(1-\beta) of the most recent gradients. The standard choice \beta = 0.9 gives 1/(1-0.9) = 10: the velocity is roughly the average of the last ten gradients. That averaging is the whole trick — average ten noisy, alternating cross-canyon gradients and they nearly cancel; average ten aligned down-the-floor gradients and you get a strong, smooth push. Concretely, suppose the down-floor gradient is a steady +1 each step. With \beta=0.9 the velocity climbs 1, 1.9, 2.71, 3.44, \dots toward a top speed of 1/(1-0.9)=10 — about ten times faster than a single plain step. Meanwhile a cross-canyon gradient of +1, -1, +1, -1 leaves the velocity hovering near zero. Memory speeds up the good direction and damps the bad one at the same time.

# One step of SGD with momentum (beta = 0.9)
# v starts as zeros, same shape as the parameters
def sgd_momentum_step(theta, v, grad, lr=0.01, beta=0.9):
    v = beta * v + grad          # blend old velocity with fresh gradient
    theta = theta - lr * v       # step the weights along the velocity
    return theta, v              # carry v into the next step

Adapting the step per-parameter: RMSProp to Adam

Momentum fixes the direction problem, but there is a second, separate problem: scale. A neural network has millions of weights, and their gradients differ wildly in size. A weight deep in the network might consistently see tiny gradients while one near the output sees huge ones. A single global learning rate has to be a compromise — small enough not to blow up the big-gradient weights, which means it is far too timid for the tiny-gradient ones. Wouldn't it be nice if each parameter got its own step size, automatically tuned to its own typical gradient scale? That is the idea of an adaptive optimizer, and we will build it in two moves.

Move one is RMSProp. For each parameter, keep a running average of its squared gradients. The square throws away the sign and measures only magnitude, so this running average is essentially 'how big are this parameter's gradients, typically?'. Then, when you take a step, divide the gradient by the square root of that average. A parameter that habitually has large gradients gets divided by a large number, so its step shrinks; a parameter with tiny gradients gets divided by a small number, so its step grows. Every parameter ends up taking a step of roughly comparable, sensible size — the global learning rate stops being a one-size-fits-all compromise.

Move two is Adam (Adaptive Moment Estimation): just combine RMSProp's per-parameter scaling with the momentum idea from the last section, plus one technical fix called bias correction. Adam keeps two running averages per parameter: a momentum-like mean of the gradients (so it still has inertia and direction smoothing) and RMSProp's mean of the squared gradients (so it still adapts the scale). Here is the full update.

The Adam update: a momentum mean, a squared-gradient mean, two bias corrections, and a scaled step.

Go line by line. g_t is the current mini-batch gradient at step t. The first line builds m_t, the first moment — a running mean of the gradients, just like momentum; it carries direction. The second line builds v_t, the second moment — a running mean of the squared gradients g_t^2, which estimates each parameter's gradient scale (its variance/magnitude), just like RMSProp. The decay rates are \beta_1 \approx 0.9 (how long the gradient memory lasts, ~10 steps) and \beta_2 \approx 0.999 (a much longer memory for the scale, ~1000 steps, so the scale estimate is smooth and stable). The factors (1-\beta_1) and (1-\beta_2) make each line a proper weighted average that stays on the same scale as the raw gradients.

Now the third line, bias correction, which trips people up. Both m_t and v_t start at zero before training. On the very first step, m_1 = (1-\beta_1) g_1 = 0.1\,g_1 — only a tenth of the true gradient, because the average is still mostly made of the zero we initialized with. The estimates are biased toward zero early on. The fix is to divide by 1-\beta_1^{\,t} and 1-\beta_2^{\,t}, where the exponent t is the step count. At t=1, 1-\beta_1^1 = 1-0.9 = 0.1, so \hat{m}_1 = m_1/0.1 = g_1 — exactly restored to full size. As t grows, \beta_1^{\,t} shrinks toward zero, the divisor approaches 1, and the correction quietly switches itself off. So bias correction is a temporary boost that matters only in the first dozens of steps, preventing a falsely tiny start.

Finally the step: \theta_t = \theta_{t-1} - \eta\, \hat{m}_t / (\sqrt{\hat{v}_t}+\epsilon). Read the fraction as 'move in the momentum direction \hat{m}_t, but shrink the move for parameters whose gradients are large (big \hat{v}_t, big square root in the denominator) and grow it for parameters whose gradients are small'. \eta is still the overall learning rate knob. \epsilon (epsilon, typically 10^{-8}) is a tiny constant added only so we never divide by zero when a parameter's gradients have been essentially nil — pure numerical safety. Tiny worked example: if \hat{m}=0.2 and \hat{v}=0.04, then \sqrt{0.04}=0.2, so the step is \eta \cdot 0.2/0.2 = \eta. The per-parameter scaling has normalized a gradient of magnitude 0.2 into a clean unit-sized step, exactly what we wanted. This combination — momentum for direction, RMSProp for scale, bias correction for a fair start — is why Adam just works out of the box on so many problems.

# One step of Adam. m and v start as zeros; t is the step counter (1, 2, 3, ...)
def adam_step(theta, m, v, grad, t, lr=3e-4, b1=0.9, b2=0.999, eps=1e-8):
    m = b1 * m + (1 - b1) * grad          # first moment: momentum-like mean
    v = b2 * v + (1 - b2) * grad**2       # second moment: mean of squared grads
    m_hat = m / (1 - b1**t)               # bias-correct (matters only early)
    v_hat = v / (1 - b2**t)
    theta = theta - lr * m_hat / (v_hat**0.5 + eps)   # per-parameter step
    return theta, m, v

Which optimizer when? SGD+momentum vs Adam

You now have two strong choices: SGD with momentum and Adam. Which should you reach for? The honest answer is 'it depends', but there is a clear, opinionated heuristic, so let me give it to you straight rather than waffle.

Reach for Adam (or AdamW) by default, especially when you want something working quickly or are unsure how to set the learning rate. Because Adam adapts the step per parameter, it is far more forgiving about the exact learning rate you pick — a value within a wide range usually trains fine. It converges fast in the early epochs, and it is the standard, expected choice for transformers and Vision Transformers (ViT), where plain SGD often struggles to train at all. If you are prototyping, fine-tuning, or working with attention-based models, Adam is the safe, productive starting point.

Reach for SGD with momentum when you are training a convolutional network to its very best final accuracy and you are willing to tune carefully. The classic ImageNet recipes (ResNets and friends) are SGD+momentum with \beta = 0.9, a well-chosen schedule, and weight decay. The folklore — backed by a lot of experience — is that for CNNs, SGD+momentum often generalizes a touch better, reaching a slightly higher test accuracy than Adam, even though Adam may reach a lower training loss faster. The price is sensitivity: SGD+momentum needs you to find a good learning rate and schedule, or it underperforms.

One more name to plant now: AdamW. It is the modern default variant of Adam, and the difference is in how it handles weight decay — a regularization technique that gently pulls weights toward zero to prevent overfitting. Plain Adam folds weight decay into the gradient, where the per-parameter scaling distorts it; AdamW 'decouples' it and applies it as a clean, separate shrink. You do not need to understand weight decay yet — guide 4 in this track is devoted to generalization and explains it properly. For now, just know that 'AdamW' is the name you will see in real configs, and it is the better-behaved cousin of Adam.

Learning-rate schedules: warmup, decay, and cosine

Even with a great optimizer, holding the learning rate at one fixed value for the whole run is suboptimal. Early in training the weights are far from any good solution, so you want big steps to cover ground fast. Late in training you are close to the bottom of the valley, and big steps will just bounce you around the minimum without settling; now you want small steps to home in precisely. The cure is a learning rate schedule: a plan that changes the learning rate over the course of training. The analogy: you sprint toward a door across a dark room, then slow to a careful shuffle as you sense you are nearly there, so you do not slam into it.

There are three pieces worth knowing. (1) Warmup: for the first few hundred or few thousand steps, linearly ramp the learning rate up from near zero to its full value. At the very start the weights are random and the gradient estimates are erratic; a full-size step on top of that can blow training up. Warmup protects these fragile early weights, and it is essentially mandatory for large batch sizes and for transformers. (2) Step decay: keep the rate constant, then drop it by a fixed factor (say ÷10) at chosen milestones — the classic ResNet schedule cuts the rate at epochs 30, 60, and 90. (3) Cosine annealing: instead of abrupt drops, glide the rate down smoothly along a cosine curve from its maximum to near zero. Cosine is the modern favourite because the smooth decay tends to land in a slightly better minimum.

Cosine annealing: the learning rate glides from eta_max down to eta_min following half a cosine wave.

Symbol by symbol: t is the current step, T is the total number of steps (the full length of the cosine phase), and \eta_{\max}, \eta_{\min} are the high and low bounds of the learning rate. The engine is the \cos(\pi t / T) term. At the start t=0, the angle is 0 and \cos 0 = +1; at the end t=T, the angle is \pi and \cos\pi = -1. So the cosine sweeps smoothly from +1 down to -1, and the (1+\cos) factor sweeps from 2 down to 0, which the \tfrac{1}{2} rescales to a glide from 1 down to 0. Multiply by the range and add the floor: \eta_t glides from \eta_{\max} at the start to \eta_{\min} at the end, fast at first and gently flattening near the end (because cosine is flat at its extremes).

A small worked example makes it tangible. Say \eta_{\max}=0.001, \eta_{\min}=0, and T=100 epochs. At t=0: \cos(0)=1, so \eta = \tfrac{1}{2}(0.001)(1+1) = 0.001 — full rate. At the midpoint t=50: \cos(\pi/2)=0, so \eta = \tfrac{1}{2}(0.001)(1+0) = 0.0005 — exactly half. At t=75: \cos(3\pi/4)\approx-0.707, so \eta \approx \tfrac{1}{2}(0.001)(0.293) \approx 0.00015. At t=100: \cos(\pi)=-1, so \eta = 0. Notice it is at half rate by the halfway point, then decays faster through the back stretch and flattens out toward zero — exactly the 'careful shuffle near the door' behaviour.

Warmup has its own tiny formula. For the first t_{\text{warm}} steps, set \eta_t = \eta_{\max}\cdot t / t_{\text{warm}}. Read it plainly: at step 0 the rate is 0; at step t_{\text{warm}} it has climbed linearly to the full \eta_{\max}; in between it is the straight-line fraction t/t_{\text{warm}} of full. For example with t_{\text{warm}}=1000 and \eta_{\max}=0.001, at step 500 the rate is 0.001\times 500/1000 = 0.0005. This gentle ramp keeps the very first updates small while the random initial weights and noisy early gradients settle down, so training does not detonate before it begins. In practice you stitch the two together: linear warmup for the first t_{\text{warm}} steps, then cosine annealing for the rest.

Putting it together: a stable, fast training run

Let us assemble everything into one concrete recipe. The three decisions — optimizer, base learning rate, and schedule — are not independent; you choose them as a set so they cooperate. Here is a config that would look completely normal for training a modern vision model, with the reasoning for every line.

# A sane, fast, stable starting recipe for a vision model
optimizer    = AdamW(lr=3e-4, betas=(0.9, 0.999), weight_decay=0.05)
batch_size   = 256
epochs       = 100
warmup       = 5          # epochs of linear warmup
schedule     = cosine_annealing(eta_max=3e-4, eta_min=0, total=epochs - warmup)
# Per step: if epoch < warmup: lr = 3e-4 * progress_through_warmup
#           else:              lr = cosine value for the remaining epochs

1. Optimizer = AdamW. It converges fast and tolerates a wide learning-rate range, so you spend less time tuning and more time iterating. (Swap to SGD+momentum only if this is a CNN and you are squeezing out the last bit of accuracy.)
2. Base learning rate = 3e-4. This is the well-worn 'Adam default' that works for a huge range of models; it is high enough to make fast early progress but not so high that AdamW's adaptive steps destabilize.
3. Warmup = 5 epochs, linear from 0 to 3e-4. Protects the random initial weights from a full-size step before the gradients are trustworthy — essential at batch size 256 and above.
4. Schedule = cosine decay to ~0 over the remaining 95 epochs. Big steps early to explore, gently shrinking to a careful crawl so the model settles precisely into a good minimum.
5. Batch size = 256, chosen to fill the GPU. It sets how many examples each gradient is averaged over (less noise per step) and interacts with the learning rate — bigger batches usually want a higher LR and a longer warmup.

Two of these knobs deserve their term links so you can dig deeper later: AdamW is the optimizer, the learning rate schedule is the warmup-plus-cosine plan, and the batch size of 256 controls how much real data backs each step. Together they form a single coherent training strategy rather than three isolated settings.

How do you know if it is working? Use the tool from guide 1: the loss curve. A healthy run with this recipe shows the training loss dropping quickly during and just after warmup, then bending into a steady, smooth decline that gently flattens as cosine pulls the rate toward zero. If the loss spikes or goes to NaN in the first few hundred steps, your warmup is too short or your base rate too high. If the loss plateaus far too early, your rate is probably too low or your schedule decays too fast. Read the curve, adjust one knob, re-run.