Optimisers: SGD, Momentum and Adam

Getting down the mountain

While training PaLM, a 540-billion-parameter language model, Google engineers saw the loss suddenly spike about 20 times. Their fix was not clever maths: they rolled back to a checkpoint about 100 steps earlier and skipped a few hundred batches of data.

The spikes happened despite clean data and a careful setup (Chowdhery et al., PaLM, section 5.1). They are a reminder that training a network is a numerical process with real failure modes. A model is only as good as the procedure that moves its weights, and the learning rate and the update rule decide whether that procedure finds a good solution, crawls, or blows up.

This lesson builds the main optimisers from plain gradient descent upwards: why the learning rate has a hard ceiling, why mini-batches add noise, how momentum and Adam fix the problems of plain SGD, and why almost every large model uses a warmup and a decaying schedule.

What is an optimiser?

An optimiser is the rule that updates a model’s parameters to reduce its loss. Almost all of them start from the gradient, the vector of slopes that says how the loss changes as each weight moves, and take a step against it. The simplest version is gradient descent:

θ are the parameters, η (eta) the learning rate, ∇L the gradient of the loss.

Why is there more than one?

Real loss surfaces are not round bowls. They have long narrow valleys, flat plateaus and directions with wildly different curvature, and real training only ever sees noisy estimates of the gradient from a small batch of data. Each optimiser in this lesson is a fix for one of those problems.

Key takeaways

• An optimiser turns gradients into parameter updates; gradient descent steps against the gradient by η.
• Badly chosen learning rates can make training crawl or diverge, even at the largest scales.
• Momentum, adaptive scaling and schedules each fix a specific weakness of plain gradient descent.

The learning rate and curvature

“Use a smaller learning rate” is the most common fix in deep learning. A one-dimensional bowl shows exactly why, and exactly where the line between slow and broken sits.

How does the step size interact with the surface?

Take the loss f(x) = (a/2)x², a parabola whose curvature is a. Its gradient is ax, so one step of gradient descent gives x ← x − ηax = (1 − ηa)x. Every step multiplies the distance to the minimum by the same factor r = 1 − ηa.

• 0 < r < 1 (η below 1/a): smooth, steady approach.
• −1 < r < 0 (η between 1/a and 2/a): overshoots each time, zig-zagging inwards.
• r < −1 (η above 2/a): overshoots further each time. Divergence.

The edge is always at η = 2/a. The fastest rate, η = 1/a, jumps to the minimum in one step. A learning rate is therefore never “big” or “small” on its own; it is only big or small relative to the curvature it meets.

A network has millions of directions, each with its own curvature. The steepest direction sets the ceiling on η, while the flattest direction sets how slowly you make progress. The ratio between them, the condition number, is what makes training hard: a surface that is 25 times steeper one way than the other forces a step size that crawls along the gentle direction. That long narrow valley is exactly the shape the next sections fight.

Key takeaways

• On a bowl of curvature a, each step multiplies the error by 1 − ηa; above η = 2/a it diverges.
• The steepest direction caps the learning rate; the flattest one sets the speed.
• Ill-conditioned surfaces (very different curvatures) make plain gradient descent slow.

Mini-batches and noise

Computing the exact gradient means a pass over the whole training set for every single step. With billions of examples, nobody does that. They estimate the gradient from a small random batch instead.

What is stochastic gradient descent?

Stochastic gradient descent (SGD) replaces the full gradient with the average gradient over a random mini-batch of examples, typically tens to thousands. The estimate is unbiased: on average it points the right way. But each individual step is noisy, and the smaller the batch, the noisier it is. The idea goes back to Robbins and Monro, who also showed that a step size shrinking at the right rate makes the noise average out (Robbins & Monro 1951).

With the full batch the path is a clean curve into the minimum. With batches of 1 or 4 it heads the right way quickly but never settles, and the final loss stays above the best possible value. Raise the learning rate and the jitter grows; lower it and the jitter shrinks, but progress slows. That trade-off motivates the schedules later in this lesson.

In practice the batch size is mostly set by hardware: as large as fits efficiently on the accelerators. When the batch grows, the learning rate usually grows with it. Goyal and colleagues trained ResNet-50 on ImageNet in one hour on 256 GPUs with batches of 8,192 by scaling the learning rate linearly with the batch size and warming it up gradually (Goyal et al.). There are limits: very large batches have been found to converge to sharper minima that generalise worse (Keskar et al.).

Key takeaways

• SGD estimates the gradient from a random mini-batch: much cheaper steps, with noise.
• Smaller batches mean noisier steps; the noise keeps a constant-rate run bouncing near the minimum.
• Bigger batches usually allow bigger learning rates, up to a point.

Momentum, RMSProp and Adam

In a narrow valley the gradient points mostly across the valley, not along it. Plain gradient descent bounces from wall to wall. Two ideas fix this: remember where you were heading, and scale each direction by how steep it usually is.

What is momentum?

Momentum keeps a running velocity, a decaying sum of past gradients, and steps along that instead of the latest gradient (Polyak 1964). Across the valley the gradient flips sign each step, so those parts cancel out. Along the valley it points the same way every step, so those parts add up, like a ball gathering speed downhill.

β is typically 0.9, so the velocity is roughly an average of the last ten gradients.

Sutskever and colleagues showed that well-tuned momentum was a key ingredient in training deep networks with plain first-order methods (Sutskever et al. 2013).

What do RMSProp and Adam add?

RMSProp keeps a running average of each parameter’s squared gradient and divides the step by its square root (Tieleman & Hinton). Steep directions, with big gradients, get smaller steps; flat directions get bigger ones. Every parameter effectively gets its own learning rate.

Adam combines the two: a momentum-style average of the gradient (m) and an RMSProp-style average of its square (v) (Kingma & Ba). Both start at zero, which would make the first steps too small, so Adam divides them by a correction factor that fades out over the first few hundred steps.

Default β₁ = 0.9, β₂ = 0.999, ε = 10⁻⁸; t is the step number and g the gradient. Squares, roots and division are per parameter.

On the elongated bowl, SGD zig-zags across the steep direction. With the default rate it still arrives in about 70 steps, but only because that rate sits just under its limit of 2/25 = 0.08: nudge it to 0.1 and it diverges, so it cannot go any faster. Momentum swings wide and then cuts through. Adam heads almost diagonally, because it rescales both coordinates to similar step sizes. RMSProp moves fast but keeps jittering near the minimum, because without momentum or decay its steps stay roughly η in size. On the Rosenbrock valley the difference is starker: SGD reaches the curved valley floor and then crawls along it, still far from the minimum after 600 steps, while momentum and Adam follow the floor all the way in.

Key takeaways

• Momentum averages past gradients: oscillations cancel, consistent directions accelerate.
• RMSProp and Adam divide by a running gradient size, giving each parameter its own step size.
• Adam’s steps barely depend on the gradient’s scale, which makes its defaults portable.

Schedules and warmup

The best learning rate at the start of training is rarely the best one at the end. Large runs change it on a schedule: up gently at first, then down.

Why decay the learning rate?

Early on you are far from any minimum, so big steps make fast progress and their noise does not matter. Late in training, with mini-batch noise, a big step keeps knocking the parameters around the minimum. Shrinking η shrinks the bounce. Common shapes are step decay (divide by 10 at fixed points) and cosine decay, which glides smoothly from the peak to near zero (Loshchilov & Hutter, SGDR).

Why warm up?

Warmup starts with a tiny learning rate and raises it linearly over the first few hundred or thousand steps. At the start, weights are random, Adam’s running averages hold little information, and the surface can be at its steepest, so the full rate would overshoot. The original transformer used 4,000 warmup steps (Vaswani et al.), and Goyal et al. relied on gradual warmup to make their large-batch learning rates stable.

On the noisy problem the constant schedule ends tens of times further from the minimum than the decaying ones, even though all four share the same peak. On the steep problem, every schedule that begins at full strength diverges within its first few steps once the peak passes 0.22, while warmup survives every peak up to 0.6, the top of the slider. Real networks are not quartics, but the same logic applies: early training is where curvature and noise are least predictable.

Key takeaways

• Decay (step or cosine) lets late steps settle instead of bouncing on mini-batch noise.
• Warmup protects the first steps, when weights, statistics and curvature are least settled.
• Warmup followed by cosine decay is the standard recipe for training transformers.

Choosing in practice

Most projects do not invent an optimiser. They pick a well-tested recipe and tune one or two numbers, starting with the learning rate.

AdamW and weight decay

Adding an L2 penalty to the loss puts a term λθ into the gradient. With SGD that shrinks every weight evenly. Inside Adam, that term is divided by √v̂ along with everything else, so weights with large gradients are hardly regularised at all. AdamW instead subtracts η λ θ directly from the weights, outside the adaptive scaling (Loshchilov & Hutter). The Overfitting and Regularisation lesson covers why the penalty matters in the first place.

Key takeaways

• Start from a proven recipe; tune the peak learning rate first.
• AdamW decouples weight decay from Adam’s per-parameter scaling.
• Divergence calls for smaller steps, longer warmup or clipping before anything exotic.

Check your understanding

Six situations you will meet when training real models.

Your training loss is fine for 200 steps, then shoots up to NaN. Nothing about the data changed. What is the first thing to try?

References

The papers behind each optimiser and schedule in this lesson.

References

1. [1]

   PaLM: scaling language modeling with Pathways(opens in a new tab)

   Chowdhery, Narang, Devlin et al., 2022

   Section 5.1 describes loss spikes during training of a 540B-parameter model and the restart-and-skip workaround.

2. [2]

   A stochastic approximation method(opens in a new tab)

   Robbins & Monro, 1951

   Annals of Mathematical Statistics. The origin of stochastic gradient methods and of decaying step sizes.

3. [3]

   Accurate, large minibatch SGD: training ImageNet in 1 hour(opens in a new tab)

   Goyal, Dollár, Girshick et al., 2017

   Batch size 8,192 on 256 GPUs using the linear scaling rule for the learning rate and a gradual warmup.

4. [4]

   On large-batch training for deep learning: generalization gap and sharp minima(opens in a new tab)

   Keskar, Mudigere, Nocedal, Smelyanskiy & Tang, 2017

   ICLR. Very large batches tended to converge to sharper minima that generalised worse.

5. [5]

   Some methods of speeding up the convergence of iteration methods(opens in a new tab)

   Polyak, 1964

   USSR Computational Mathematics and Mathematical Physics. The heavy-ball (momentum) method.

6. [6]

   On the importance of initialization and momentum in deep learning(opens in a new tab)

   Sutskever, Martens, Dahl & Hinton, 2013

   ICML. Well-tuned momentum and initialisation let plain first-order methods train deep and recurrent networks.

7. [7]

   Lecture 6.5: RMSProp, divide the gradient by a running average of its recent magnitude(opens in a new tab)

   Tieleman & Hinton, 2012

   Coursera Neural Networks for Machine Learning. RMSProp was introduced in these lecture slides and never formally published.

8. [8]

   Adam: a method for stochastic optimization(opens in a new tab)

   Kingma & Ba, 2015

   ICLR. Combines momentum with per-parameter RMSProp-style scaling and corrects the bias of zero-initialised averages.

9. [9]

   The marginal value of adaptive gradient methods in machine learning(opens in a new tab)

   Wilson, Roelofs, Stern, Srebro & Recht, 2017

   NeurIPS. On several benchmarks, well-tuned SGD with momentum generalised better than adaptive methods.

10. [10]

    SGDR: stochastic gradient descent with warm restarts(opens in a new tab)

    Loshchilov & Hutter, 2017

    ICLR. Introduced the cosine annealing schedule now used in most large training runs.

11. [11]

    Attention is all you need(opens in a new tab)

    Vaswani, Shazeer, Parmar et al., 2017

    NeurIPS. The original transformer was trained with Adam, 4,000 warmup steps and inverse square-root decay.

12. [12]

    Decoupled weight decay regularization(opens in a new tab)

    Loshchilov & Hutter, 2019

    ICLR. AdamW: apply weight decay directly to the weights instead of adding an L2 term to the gradient.