Learning Rate Schedules and Warmup

A learning rate schedule means varying η during training instead of holding it fixed. A constant value is rarely best across all phases: at the start you want to move carefully, then advance quickly, and near the end fine-tune with small steps. A schedule (a function that says how large η is at each step) captures that idea, and warmup adds an initial ramp that keeps the network from destabilising in its first iterations. The same explanation is available in Spanish.

Key takeaways

• A learning rate schedule replaces the fixed value of η with a function $\eta_t = \eta_0 \cdot \operatorname{decay}(t)$ that changes step by step.
• The three most common decays are step (multiply by 0.1 every so many epochs), exponential and cosine.
• Warmup raises η from near zero over the first steps (4,000 in the original Transformer) to stabilise the start.
• The cosine annealing of Loshchilov and Hutter (2016) lowers η along half a cosine wave down to a floor near 10% of the peak.
• Transformers combine linear warmup and decay, a recipe that almost every large language model uses today.

Why vary the learning rate

The learning rate sets the size of each step when updating the weights with the rule $\mathbf{W} \leftarrow \mathbf{W}-\eta\cdot\nabla L$. Too high, and the loss oscillates or diverges; too low, and training takes forever. The catch is that the ideal value is not the same at every moment.

Early in training the weights are freshly initialised and the gradients are large and noisy, so a cautious step is wise. In the middle, once the descent direction is clear, a larger step speeds up progress. Near the minimum, a big step would overshoot the valley, so it pays to shrink it and settle. A fixed η forces a compromise between these three phases; a schedule treats them separately and often cuts the number of epochs needed by 10% to 30% versus a well-chosen constant rate.

Step, exponential and cosine decay

Step decay is the simplest: every so many epochs you multiply η by a factor γ smaller than 1. A classic recipe in computer vision starts at 0.1 and multiplies by 0.1 every 30 epochs, so by epoch 60 the rate is a hundred times smaller. In closed form:

$$\eta_t = \eta_0\cdot\gamma^{\lfloor t/\text{step_size}\rfloor}$$

Exponential decay does the same thing continuously, without steps. The rate falls smoothly as $\eta_t = \eta_0\cdot e^{-\lambda t}$, where λ controls the decay speed. It is handy when you prefer a curve with no abrupt jumps.

Cosine annealing, proposed by Loshchilov and Hutter in 2016, lowers η along the first half-wave of a cosine, slow at first, fast in the middle and slow again at the end:

$$\etat = \eta{\min}+0.5\,(\eta0-\eta{\min})\,(1+\cos(\pi t/T))$$

Here $T$ is the total cycle length and $\eta_{\min}$ the floor (often zero or 10% of the peak). Cosine is today the favourite for training large networks because its shape spends more time at useful high rates and then fine-tunes gently. You can see the restart variant in the PyTorch^[1] documentation, which ships all three decays as ready-to-use schedulers.

Warmup

Warmup solves a start-up problem. In the first steps, with random weights and optimiser statistics not yet settled (Adam’s moments take several iterations to warm up), a high η can push the weights into a bad region the network never recovers from. Warmup starts η at a very small value and raises it linearly to the peak over the first steps, typically between 500 and 10,000 depending on model size.

BERT, for instance, uses 10,000 warmup steps before it begins to decay, and many language models spend around 1% of total steps on warmup. The usual combination is linear warmup followed by cosine decay: first up, then down. This pattern cushions the first updates and avoids the loss spikes that appear when training with very large batches.

Schedules in transformers

The recipe that popularised warmup comes from «Attention Is All You Need»^[2] by Vaswani and co-authors (2017). Their schedule combines both ideas in a single formula:

$$\etat = d{\text{model}}^{-0.5}\cdot\min!\left(t^{-0.5},\; t\cdot\text{warmup_steps}^{-1.5}\right)$$

With $d_{\text{model}}=512$ and $\text{warmup_steps}=4000$, the rate rises linearly for the first 4,000 steps and then falls proportionally to the inverse square root of the step. As the authors explain, «this corresponds to increasing the learning rate linearly for the first warmup_steps training steps, and decreasing it thereafter proportionally to the inverse square root of the step number».

The factor $d_{\text{model}}^{-0.5}$ couples the rate to the model width: the larger the network, the smaller the peak η. This idea is still alive in current models, which almost always train with linear warmup plus cosine. For the full context of where η sits within training, the mathematics behind neural networks roadmap places optimisation next to linear algebra and calculus.

Configuration examples

In practice a schedule is defined with two or three parameters. These are the most common:

step:    eta_0 = 0.1,  gamma = 0.1,  step_size = 30 epochs
cosine:  eta_0 = 3e-4, eta_min = 3e-5, T = total steps
warmup:  warmup_steps = 4000, then cosine down to eta_min

A simple guide: use step decay when reproducing classic vision recipes, cosine when training a large model from scratch, and always add warmup if you work with transformers or batches above a few thousand examples. The general idea of adjusting η is also summarised on Wikipedia^[3], useful as a quick reference for the terms.

Frequently asked questions

What is the difference between warmup and decay?

Warmup is the initial phase in which η rises from near zero to the peak over the first steps; it stabilises the start. Decay is what happens afterwards: η falls progressively (step, exponential or cosine) to fine-tune convergence. They are almost always used together, one first and then the other.

Which schedule should I pick to start?

For most modern models, linear warmup over 1% of the steps followed by cosine decay down to a small floor is a safe choice. If you reproduce classic image-classification work, step decay with γ = 0.1 every 30 epochs still works well.

Does warmup replace choosing a good base learning rate?

No. Warmup and decay manage how η changes over time, but the peak value is still a hyperparameter you must tune. A badly chosen peak diverges even with warmup; the ramp only smooths the start, it does not fix a peak η that is too high.

Conclusion

A learning rate schedule turns η from a fixed number into a function of time: warmup to start calmly, a peak to advance and a decay (step, exponential or cosine) to settle. It is one of the choices that most influences whether a training run converges well or stalls. The natural next step is to revisit the base learning rate and fit it into the rest of the optimisation route.

Sources

1. PyTorch
2. «Attention Is All You Need»
3. Wikipedia
4. SGDR: Stochastic Gradient Descent with Warm Restarts (Loshchilov and Hutter, 2016)

Route: Advanced Neural Network Optimization