Learning rate

Finding a good rate: the sweep

You can't calculate the best rate ahead of time for a real model. So you test. Train briefly with several rates and compare the loss after a fixed number of steps. This is a learning-rate sweep.

Space the rates by factors of ten, with a stop in between: 0.001, 0.003, 0.01, 0.03, and so on. The best rate could be anywhere across a huge range, and this spacing covers it in a few tries.

The code below builds a small made-up dataset: 500 rows, three features and a column of ones for the intercept. Then it trains a linear model for 100 steps at each rate. The final_loss function gives up and returns infinity as soon as the weights blow up. The chart caps the loss at 20 so the failures stay on screen. Its x-axis uses a log scale, where each tick is ten times the one before.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(0)
X = np.column_stack([rng.normal(size=(500, 3)), np.ones(500)])
y = X @ np.array([2.0, -1.0, 0.5, 3.0]) + rng.normal(0, 0.5, 500)

def final_loss(lr, steps=100):
    w = np.zeros(4)
    for _ in range(steps):
        w -= lr * 2 / len(y) * X.T @ (X @ w - y)
        if not np.all(np.isfinite(w)) or np.abs(w).max() > 1e6:
            return np.inf
    return np.mean((X @ w - y) ** 2)

rates = [0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1.0, 1.3]
losses = [final_loss(r) for r in rates]
for r, l in zip(rates, losses):
    print(f"lr {r:<6} loss {l:.3f}" if np.isfinite(l) else f"lr {r:<6} diverged")
plt.plot(rates, [min(l, 20) for l in losses], marker="o")
plt.xscale("log")
plt.xlabel("learning rate")
plt.ylabel("loss after 100 steps (capped at 20)")

Read the printout from top to bottom:

• 0.001 and 0.003 are too small. After 100 steps the loss is still 8.764 and 4.583.
• 0.01 is close, at 0.670.
• 0.03, 0.1 and 0.3 all reach 0.265, the bottom of this valley.
• 1.0 and 1.3 diverge: the loss runs off to infinity instead of settling.

On the log-scale chart, the curve is roughly U-shaped. The left side is too slow. The right side blows up. A wide good region sits in between.

Pick the largest rate that trains smoothly, or a little below it. Larger rates make faster progress. A little room below the cliff keeps you safe if the data or model changes slightly. Some deep learning libraries automate the same idea with a tool called a "learning-rate finder".

Schedules: change the rate during training

Think of parking a car. You drive in at a decent speed, then creep the last few feet. One constant speed would be either too slow early or too fast at the end.

A constant learning rate has the same problem. Large early steps make fast progress. Near the end you want small steps that settle precisely. A schedule changes the rate as training goes on. Three are common:

• Step decay: cut the rate by a fixed factor every so many epochs. For example, start at 0.1 and divide by 10 every 30 epochs: 0.1, then 0.01, then 0.001. (An epoch, from the last lesson, is one pass over all the training data.)
• Warmup: start tiny and ramp up in a straight line over the first few hundred or few thousand steps. Early gradients in big networks can be large and erratic. Warmup avoids an early blow-up while the model finds its feet.
• Cosine decay: after warmup, glide down along a smooth curve to near zero by the end.

print(0.1 / 10)

Cosine decay has a formula. First, the new symbols:

• $t$ is the step number, and $\eta_t$ is the learning rate at step $t$.
• $\eta_{\max}$ is the peak rate, the one you reach at the end of warmup.
• $t_{\text{warm}}$ is the number of warmup steps.
• $T$ is the total number of steps in training.
• $\cos$ is the cosine function from trigonometry, and $\pi$ (pi) is about 3.14.

In plain words: the fraction inside measures progress through the decay phase. It runs from 0 at the end of warmup to 1 at the last step. The cosine part turns that progress into a smooth slide from 1 down to 0. Multiply by the peak rate, and the rate glides from the peak to zero. It falls gently at first, fastest in the middle, and gently again at the end.

You don't need to know how cosine works. For this formula, its input runs from 0 to π. Over that range, cosine slides smoothly from 1 down to −1. Adding 1 and halving turns that into a slide from 1 down to 0.

Try it with small numbers: a peak rate of 0.1, 10 warmup steps and 110 steps in total. So the decay phase runs from step 10 to step 110.

• Step 10: progress is 0. Cosine of 0 is 1, so the rate is 0.1 × ½ × (1 + 1) = 0.1, the peak.
• Step 60: progress is 50 / 100 = 0.5. Cosine of π/2 is 0, so the rate is 0.1 × ½ × 1 = 0.05.
• Step 110: progress is 1. Cosine of π is −1, so the rate is 0.1 × ½ × 0 = 0.

import math

peak, warm, total = 0.1, 10, 110
for t in [10, 35, 60, 85, 110]:
    progress = (t - warm) / (total - warm)
    print(t, round(peak * 0.5 * (1 + math.cos(math.pi * progress)), 4))

It prints 0.1, 0.0854, 0.05, 0.0146 and 0.0. The rate halves by the midpoint and reaches zero at the end.

Warmup followed by cosine decay is the standard recipe for training and fine-tuning transformers, the kind of model behind LLMs.

lr0 = 0.1
epoch = 65
print(round(lr0 * 0.1 ** (epoch // 30), 6))

Scaling and the learning rate

Back to the shower. Now picture one tap for hot and one for cold. The hot tap is very touchy: a tiny turn changes everything. The cold tap is very stiff. Now say one setting decides how hard you turn both taps. You're stuck. A setting gentle enough to keep the hot tap in check barely moves the cold one.

Unscaled features cause the same trap. Say one feature is order value, in hundreds of dollars. Another is the share of the order that was discounted, between 0 and 1. The loss is very steep along the order-value weight and very flat along the discount weight. The loss surface becomes a long, narrow valley. You saw this as the zig-zag 2D bowl in Phase 2.

A rate small enough to avoid blowing up in the steep direction barely moves the flat one. Training is slow and zig-zags. No single rate suits both weights.

Standardizing the features fixes this. You did it in the data-prep lesson: rescale each column so its average is 0 and its spread is 1. Every feature then has a similar size. The valley becomes rounder, so one learning rate suits every weight. Optimizers such as Adam, which you'll meet in Phase 4, also help. They give each weight its own step size.