Stochastic Gradient Descent › Learning Rate Schedules

Training with a constant learning rate has a fundamental tension. Early in training, you want a large learning rate to explore parameter space and make rapid progress. Late in training, you want a small learning rate to converge precisely to a minimum without oscillating around it. A learning rate schedule resolves this tension by changing $\eta$ over the course of training.

This section covers the most important schedules. All of them modify the base learning rate $\eta_0$ over training steps or epochs.

Step decay reduces the learning rate by a fixed factor every $k$ epochs. If $\gamma$ is the decay factor (typically 0.1 or 0.5) and $e$ is the current epoch:

$$\eta_e = \eta_0 \cdot \gamma^{\lfloor e / k \rfloor}$$

where:

• $\eta_e$: the learning rate at epoch $e$
• $\eta_0$: the initial learning rate
• $\gamma$: the decay factor per drop, a value in $(0, 1)$ (typically 0.1 or 0.5)
• $k$: the number of epochs between drops
• $\lfloor e / k \rfloor$: the number of drops so far (the floor of $e/k$)

Step decay is simple and predictable. The loss usually decreases smoothly between drops and then jumps to a lower value immediately after each drop, as the optimizer can now make more precise adjustments in the current neighborhood.

The main drawback of step decay is that the sharp transitions can be undesirable. Hyperparameter tuning requires choosing both $\gamma$ and $k$, and the performance can be sensitive to these choices.

Instead of discrete drops, exponential decay reduces the learning rate continuously at every step:

$$\eta_t = \eta_0 \cdot e^{-\lambda t}$$

where:

• $t$: the current training step
• $\lambda > 0$: the decay rate, controlling how quickly the learning rate shrinks
• $e^{-\lambda t}$: the continuous decay factor, which approaches zero as $t$ increases

Exponential decay produces a smooth learning rate curve. However, if $\lambda$ is too large, the learning rate decays too quickly and training stalls before convergence.

Cosine annealing, introduced by Loshchilov and Hutter (2017) in their "SGDR" paper, is one of the most popular learning rate schedules for modern NLP. It smoothly decreases the learning rate from $\eta_{\max}$ to $\eta_{\min}$ following a cosine curve:

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

where:

• $\eta_t$: the learning rate at step $t$
• $\eta_{\min}$: the minimum learning rate (often 0 or a very small value like 1e-6)
• $\eta_{\max}$: the maximum learning rate (the "base" learning rate)
• $T$: the total number of training steps
• $\cos(\pi t / T)$: the cosine term, which starts at 1 (giving $\eta_{\max}$) and ends at $-1$ (giving $\eta_{\min}$) over the interval $[0, T]$

At $t = 0$, the cosine term is 1, so $\eta_0 = \eta_{\max}$. At $t = T$, the cosine term is $-1$, so $\eta_T = \eta_{\min}$.

Cosine annealing has several attractive properties: it starts large, decays slowly at first, then accelerates, and finally decays slowly again near the end. This matches intuitions about optimization: you want large steps early on, more cautious steps in the middle, and very fine-tuning at the end.

The "SGDR" variant adds warm restarts: at the end of a cosine cycle, the learning rate jumps back to $\eta_{\max}$ and the cycle repeats, often with a longer period. This allows the optimizer to escape any local minimum it has settled into and explore other regions of parameter space.

For large language models and transformers, training often begins with a warmup phase where the learning rate is linearly increased from 0 (or a tiny value) to the target learning rate over a number of steps $T_{\text{warmup}}$:

$$\eta_t = \eta_{\max} \cdot \frac{t}{T_{\text{warmup}}} \quad \text{for } t \leq T_{\text{warmup}}$$

where:

• $T_{\text{warmup}}$: the number of warmup steps (typically a few hundred to a few thousand)
• $\eta_{\max}$: the target learning rate after warmup completes
• $t / T_{\text{warmup}}$: the linear ramp factor, going from 0 to 1

After warmup, the learning rate typically follows a decay schedule (cosine annealing is most common).

Why is warmup necessary? Early in training, the parameters are far from their optimal values. Gradients can be very large in magnitude, and starting with a large learning rate can cause catastrophic parameter updates that destabilize the entire network. By starting small and ramping up, warmup gives the network time to initialize into a reasonable part of parameter space before taking large steps.

In the original transformer (Vaswani et al., 2017), the following combined warmup-decay schedule was used:

$$\eta_t = d_{\text{model}}^{-0.5} \cdot \min\left(t^{-0.5},\; t \cdot T_{\text{warmup}}^{-1.5}\right)$$

where:

• $d_{\text{model}}$: the model dimension (e.g., 512 in the original transformer), which scales the overall learning rate to be smaller for larger models
• $t^{-0.5}$: the decay term active after warmup (polynomial decay)
• $t \cdot T_{\text{warmup}}^{-1.5}$: the warmup term active during the warmup phase
• The $\min$ selects the warmup term early in training and the decay term after

This formula automatically transitions from warmup to decay without requiring separate configuration of the phase boundary.

For most modern transformer training, the practical recommendation is: linear warmup for 1-5% of total training steps, followed by cosine annealing to $\eta_{\min} \approx 0$.