Scheduling the Learning Rate

Until now the learning rate $\eta$ has been one fixed number. Yet the ideal step size changes during training. Early on we want large steps to cover ground quickly. Near the end we want small steps so the parameters can settle instead of bouncing around the minimum. A schedule changes $\eta$ as training goes on.

Decay schedules

The simplest idea is to start with a rate $\eta_0$ and shrink it. Three common forms, where $k$ is the step number:

With $\eta_0 = 0.1$, halving every 20 steps, $\lambda = 0.03$ for the exponential and $\lambda = 0.05$ for inverse-time, the rates at a few steps are:

The decay speed is a hyperparameter like any other. Too fast and training stalls before it has finished learning. Too slow and the noise near the end is not reduced.

Line search: choosing the rate each step

Instead of a fixed plan, we can search for a good step size at every update. Take the gradient direction $g = \nabla L(\theta)$ and try a large step. If it does not reduce the loss enough, shrink it and try again. A standard test, the backtracking rule, accepts $\alpha$ when

for a small constant $c$ such as $10^{-4}$: the loss must drop by a fixed fraction of what the slope predicts. Start with $\alpha = 1$ and halve until the test passes.

On the narrow valley, starting from $(1, 1)$, this rule picks different rates at different steps:

The loss started at 5.5 and falls at every step without us choosing a rate by hand. The price is extra loss evaluations in every trial. For a large network each trial costs a forward pass over the data, and with mini-batches the loss is itself noisy, so line search is rarely used for training deep networks. Fixed schedules are cheaper.

Warm-up

Sometimes the start is the dangerous part. At the very beginning the weights are random, gradient estimates can be wild, and methods that keep running averages (like those in the coming lessons) have had little data to average. A large rate right away can throw training off.

Warm-up starts with a tiny rate and raises it linearly to the target over the first $W$ steps:

Cosine annealing

After warm-up, we can lower the rate smoothly along a half cosine wave, from $\eta_{\max}$ down to $\eta_{\min}$ over $T$ steps:

It falls slowly at first, fastest in the middle, and slowly again at the end, which lets the parameters settle gently. Some schemes restart the cosine cycle several times.

With $\eta_{\max} = 0.1$, $\eta_{\min} = 0$, $W = 10$ and $T = 100$, the rate at steps 0, 4, 9, 10, 32, 55 and 99 is

It ramps up, peaks at step 10, and then declines to zero by the end. A warm-up followed by cosine decay is a very common recipe for training large networks.

In code

import math

def step_decay(k, eta0=0.1, drop=0.5, every=20):
    return eta0 * drop ** (k // every)

def warmup_cosine(t, eta_max=0.1, eta_min=0.0, warmup=10, total=100):
    if t < warmup:
        return eta_max * (t + 1) / warmup
    progress = (t - warmup) / (total - warmup)
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * progress))

print([round(warmup_cosine(t), 3) for t in (0, 4, 9, 10, 32, 55, 99)])

Inside a training loop, just set eta = warmup_cosine(step) before each update.