為甚麼接近終點時要放慢學習率？

Rolling down a bowl with noisy pushes

Training an AI is like rolling a ball down a bowl. The ball's position is the AI's dials (its weights), the height is how many mistakes it makes, and the bottom means the fewest mistakes. Each training step moves the ball a little downhill:

w ← w − step size × push

The step size is the learning rate and the push is the gradient. But each push comes from just a few examples (a mini-batch), not all the data, so it always points a little off the true downhill direction.

The bowl in the Short: mistake = ½ · (w1² + 4 · w2²), a long oval with its bottom at (0, 0). The ball starts at (−4, 1.5), 4.27 from the bottom. Each push is the true slope (w1, 4 · w2) plus 0.6 × a normal random number in each direction, from numpy's default_rng(2252). All three runs below use the same noise and take 60 steps.

Big steps bounce, small steps crawl

With a big step size, 0.45, the ball is about 0.2 from the bottom by step 6. Then the noise takes over: near the bottom the true slope is tiny, so the push is mostly noise, and a big step turns it into a big jump past the middle. The ball bounces around and never settles. Over the last 10 steps it is 0.50 from the bottom on average.

Big steps: near the bottom by step 6, then every off push throws the ball past the middle. Step 60, still bouncing.

With a small step size, 0.02, the bounces are tiny, but the trip is slow: 3.33 from the bottom after 10 steps, 2.73 after 20, 1.75 after 40 and 1.21 after 60. Same number of steps, still far away.

Small steps: tiny bounces, but after the same 60 steps the ball is still 1.2 away. The big-step cloud stays as a grey ghost.

Lecture 09 puts it in one line: you need a high learning rate when far away and a low learning rate when close.

Start big, then cut the step size in half

A learning-rate schedule uses both. Start with the big step size, 0.45, and halve it every 10 steps: 0.45, 0.225, 0.113, 0.056, 0.028, 0.014.

The first 10 steps match the big-step run exactly. After that, each halving shrinks the bounces: the ball is 0.25 from the bottom at step 20, 0.14 at step 40 and 0.07 at step 60. Over the last 10 steps it averages 0.09, more than five times closer than big steps alone.

Over 2,000 different noise sequences, the last-10-step averages are: big steps 0.49, small steps 1.31, the schedule 0.09. That is about 5.6 times closer than big steps.

The staircase: 0.45, halved at steps 10, 20, 30, 40 and 50. Big steps end 0.5 away, the schedule 0.09 away.

Other schedules

Halving every n steps is called step decay. Lecture 09 lists other ways to shrink the learning rate α:

• Step decay: multiply α by a decay rate (often 0.5) every n steps.
• Exponential decay: α = tᵉᵖᵒᶜʰ · α0, where the decay rate t is below 1, so α shrinks a little every epoch.
• α = α0 / (1 + decay_rate · epoch): with α0 = 0.1 and decay_rate = 1.0, the learning rate goes 0.1, 0.05, 0.033, 0.025 over epochs 0 to 3.

A popular smooth option is the cosine schedule, which lowers the learning rate along half a cosine wave. SGDR (Loshchilov and Hutter) adds warm restarts: the learning rate decays, jumps back up, and decays again.

Why it works

In 1951, Herbert Robbins and Sutton Monro studied how to find an answer when every measurement is noisy, just like a ball pushed by a few examples at a time. They showed that noisy steps still settle if the step sizes aₜ shrink so that:

• Σ aₜ = ∞: added up, the steps are unlimited, so the ball can still travel as far as it needs.
• Σ aₜ² < ∞: the steps shrink fast enough for the noise they carry to average out.

For example, aₜ = a0 / (1 + k · t) meets both. Halving every few steps is the practical version used in training. Strictly, halving forever adds up to a finite total and misses the first condition, but over a fixed training run it does the same job: big steps far away, small steps up close.

Robbins and Monro, "A Stochastic Approximation Method" (1951): noisy steps settle when the step sizes shrink.

Try it in PyTorch

In PyTorch a schedule wraps the optimizer. This one halves the learning rate every 10 epochs:

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

Call scheduler.step() once per epoch, after that epoch's optimizer steps. This article's schedule on its own:

import torch as t
o=t.optim.SGD([t.ones(1)],lr=.45)
s=t.optim.lr_scheduler.StepLR(o,10,.5)
for e in range(60):
  if e%10==0:print(e,s.get_last_lr())
  o.step();s.step()

It prints the staircase: 0 [0.45], 10 [0.225], 20 [0.1125], 30 [0.05625], 40 [0.028125], 50 [0.0140625]. Change 10 to 5 and run it again: how small is the step size by step 50 now?

References

• Robbins, H., & Monro, S. (1951). A Stochastic Approximation Method. The Annals of Mathematical Statistics, 22(3), 400–407. https://doi.org/10.1214/aoms/1177729586
• Loshchilov, I., & Hutter, F. (2016). SGDR: Stochastic Gradient Descent with Warm Restarts. https://arxiv.org/abs/1608.03983
• CS231n notes, "Annealing the learning rate": https://cs231n.github.io/neural-networks-3/
• CS878 Deep Learning, SEECS NUST, lecture 09 (slides 5–7).

This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.