Optimisers: momentum and Adam

Training deep networksabout 10 minutes

Plain stochastic gradient descent has two weaknesses you've already seen: it zig-zags across narrow valleys, and one learning rate must serve parameters whose gradients differ by orders of magnitude. Modern optimisers fix both.

Momentum

Give the parameters a velocity that accumulates past gradients, like a heavy ball rolling downhill (Polyak, 1964). Components that keep pointing the same way build up speed; components that flip sign each step (the zig-zag) cancel out.

v \leftarrow β v + \nabla L (θ), θ \leftarrow θ - η v, β \approx 0.9

Adaptive learning rates: RMSProp and Adam

RMSProp (Hinton, 2012) divides each parameter's step by a running average of its recent squared gradients, so parameters with consistently large gradients take smaller steps and rarely-updated ones take bigger steps. Adam (Kingma & Ba, 2014) combines this with momentum, plus a bias correction because both averages start at zero:

m \hat{m} θ \leftarrow β_{1} m + (1 - β_{1}) g = m / (1 - β_{1}^{t}) \leftarrow θ - η \hat{m} / (\hat{v} + ϵ) v \hat{v} \leftarrow β_{2} v + (1 - β_{2}) g^{2} = v / (1 - β_{2}^{t})

Typical defaults are $η = 1 0^{- 3}$ , $β_{1} = 0.9$ , $β_{2} = 0.999$ , $ϵ = 1 0^{- 8}$ . Adam is forgiving and the most common choice; AdamW (Loshchilov & Hutter, 2017), which applies weight decay separately from the adaptive step, is the standard for transformers. Well-tuned SGD with momentum still sometimes generalises better on image models.

Three optimisers race down the same narrow valley from the same start. Plain gradient descent zig-zags; momentum and Adam make progress along the valley.

Learning-rate schedules

The best learning rate changes during training: large early to make fast progress, small late to settle into a minimum. Common schedules are step decay, cosine decay to near zero, and a short warm-up at the start (ramping up from a tiny rate), which stabilises the early steps of large models, especially transformers with Adam.

The learning rate is still the single most important hyperparameter. If you can only tune one thing, tune it on a logarithmic grid (1e-4, 3e-4, 1e-3, 3e-3, …).

try itedit the code and run it, nothing can break

python sleeps until you press run

Momentum averages gradients over time; Adam also scales each parameter's step by its gradient history. Pair either with a learning-rate schedule.

← Initialisation and vanishing gradients