Learning-Rate Schedules, Warmup and Clipping

Warm up, then cool down

The learning rate (Chapter 2's step size) is not held fixed. A typical schedule has two parts:

• Warmup. Start near zero and rise linearly to the peak over the first few thousand steps. Early on, the weights are random and Adam's running averages of the gradient have barely started, so a full-size step can push the model somewhere it never recovers from. Goyal and colleagues used a gradual warmup to train with very large batches without losing accuracy Established.
• Decay. After the peak, lower the rate smoothly so the model can settle into a good region. The cosine schedule follows half a cosine curve from the peak down to a minimum Established:

$$\eta_t=\eta_{\min}+\tfrac12(\eta_{\max}-\eta_{\min})\left(1+\cos\frac{\pi t}{T}\right)$$

Tiny example. Halfway through the decay, $\cos(\pi/2)=0$, so the rate is exactly midway between peak and floor. With Llama 3 405B's numbers that is $8\times10^{-7}+\tfrac12(8\times10^{-5}-8\times10^{-7})\approx 4.04\times10^{-5}$.

A linear warmup to 8 × 10⁻⁵ over the first 8,000 steps (dashed line), then a cosine decay to 8 × 10⁻⁷ by step 1.2 million. The warmup is under 1% of the run, too short to see on the left.

Llama 3's peak learning rates fell as models grew: 3 × 10⁻⁴ for 8B, 1.5 × 10⁻⁴ for 70B and 8 × 10⁻⁵ for 405B Established. Larger models are more sensitive to large steps.

AdamW and weight decay

Pretraining usually uses AdamW, which applies weight decay (a pull of every weight towards zero, a form of regularization) separately from Adam's adaptive step. In Llama 3's scaling-law runs, the weight decay at each step was set to 0.1 times the learning rate at that step Established.

Gradient clipping

If the gradient's overall size (its norm) exceeds a threshold, scale it down to the threshold. One unusual batch then can't throw the weights far. OPT's team found that lowering the clipping threshold from 1.0 to 0.3 early in training helped stability Established.