Cosine Learning Rate Schedule › Cosine Annealing with Warm Restarts

Training on a fixed cosine schedule marches the learning rate all the way to $\eta_{\min}$ and stops. This risks getting trapped in a sharp local minimum. Loshchilov and Hutter's 2016 paper "SGDR: Stochastic Gradient Descent with Warm Restarts" introduced a powerful extension: periodically reset the learning rate back to its maximum value, triggering a new cosine decay cycle.

In SGDR, each restart begins a new cosine annealing cycle with a potentially longer period. The period of the $i$-th cycle is given by multiplying the base cycle length $T_0$ by a multiplicative factor raised to the cycle index:

where:

• $T_i$: the number of steps in the $i$-th cycle (starting from $i = 0$)
• $T_0$: the number of steps in the first (base) cycle
• $T_{\text{mult}}$: a multiplicative factor that lengthens each successive cycle. Common values are 1 (equal-length cycles) or 2 (doubling cycles)
• $i$: the zero-based cycle index

Within each cycle, the learning rate follows standard cosine decay from $\eta_{\max}$ down to $\eta_{\min}$. Letting $t_{\text{cycle}}$ denote the step count within the current cycle (reset to 0 at each restart):

where:

• $t_{\text{cycle}}$: the number of steps elapsed since the last restart (or since training began for the first cycle)
• $T_i$: the total length of the current cycle

When $T_{\text{mult}} = 1$, all cycles have the same length $T_0$, giving a periodic saw-tooth pattern. When $T_{\text{mult}} = 2$, each cycle is twice as long as the previous, so the schedule allocates progressively more training time to each refinement phase.

A sharp local minimum is characterized by a loss surface with steep walls: the loss is low at the minimum but increases rapidly in any direction. Models stuck in such minima tend to generalize poorly because sharp minima are sensitive to small perturbations in parameters. When the learning rate suddenly jumps back up at a restart, the optimizer takes larger steps that can escape the steep walls and find a path into a flatter, wider basin.

To understand this geometrically, imagine the parameter space as a landscape with many valleys. Some valleys are narrow and deep, corresponding to sharp minima where the Hessian of the loss has large eigenvalues. Others are wide and shallow, corresponding to flat minima where the loss changes slowly in all directions. With a monotonically decreasing learning rate, once the optimizer enters a narrow valley it tends to stay there because the small learning rate prevents it from climbing back out. With warm restarts, the sudden surge in learning rate gives the optimizer enough momentum to escape the narrow valley and potentially discover a wider, more generalizable one.

Flat minima are associated with better generalization in practice. The parameter values in a flat minimum remain nearly optimal over a broader range of perturbations, meaning that the model will perform well even under the slight distributional shifts that occur between training and test conditions. Warm restarts act as a form of implicit regularization by repeatedly biasing the optimizer toward broader basins.

There is also a practical benefit for model averaging. The models saved at the end of each cosine cycle, when the learning rate is near $\eta_{\min}$ and parameters have converged locally, represent good checkpoints for ensembling. Averaging several such checkpoints often outperforms using the last checkpoint alone, a technique sometimes called snapshot ensembling. Each cycle's endpoint represents a locally good solution reached from a different exploratory trajectory, and their average often achieves better generalization than any individual checkpoint.

Warm restarts are most beneficial when the loss landscape is known to contain many suboptimal sharp minima, which is often the case in computer vision tasks and some classification problems. For large language model pretraining, practitioners have largely converged on single-cycle cosine decay without restarts. The reason is scale: LLMs are trained for hundreds of billions of tokens, and the cost of each "exploratory" restart is measured in millions of compute dollars. At that scale, the uncertainty about whether a restart will improve or hurt performance makes the single-cycle cosine the safer, more predictable choice.

Restarts are still useful in fine-tuning scenarios, where the model may get stuck in a local minimum quickly due to the small dataset size. They are also useful when you want multiple model checkpoints for ensembling without re-running training from scratch. If your task has a clear evaluation signal and you can afford the extra exploration, the doubling schedule ($T_{\text{mult}} = 2$) often offers a good balance between exploration and refinement.

One of the most practical applications of warm restarts is snapshot ensembling, introduced by Huang et al. in the paper "Snapshot Ensembles: Train 1, Get M for Free" (2017). The main insight is that when the learning rate resets to $\eta_{\max}$ at each restart, the model takes a fresh exploratory path through the loss field. Each time the learning rate decays back toward $\eta_{\min}$, the model converges to a locally good solution. By saving model weights at the end of each cycle, you collect a set of diverse models, all trained on the same data in a single training run.

Averaging the predictions of these snapshot models at inference time produces an ensemble that typically outperforms any individual snapshot. Each cycle's endpoint represents a distinct region of parameter space, reached after exploring from a different starting point (the previous cycle's endpoint, after a large learning rate jump displaced it). This is qualitatively different from averaging checkpoints from the same monotonic training run, where consecutive checkpoints are often very similar because the learning rate has been small for many steps.

The computational cost of snapshot ensembling is low compared to training separate models from scratch. You pay for $M$ cycle-endpoints worth of storage and $M$ forward passes at inference, but the total training compute is the same as a single model with warm restarts. For tasks where inference throughput is not critical, this trade-off is often worthwhile.

In practice, the quality of snapshot ensembles degrades if cycle lengths are very short, because the model has not converged within a single cycle and the snapshots are not sufficiently distinct from each other. Cycles of at least 10-20% of total training steps tend to produce better-quality individual snapshots. The doubling schedule, where later cycles are longer, naturally provides better snapshots from later cycles because each subsequent restart gives the model progressively more time to settle into its local minimum before being saved.

For large language models, snapshot ensembling is rarely used in pretraining because the cost of storing and running multiple large model copies at inference time is prohibitive. However, in settings like fine-tuning a model for a classification or retrieval task, where model size is more manageable and warm restarts are commonly used anyway, snapshot ensembling can provide meaningful quality improvements without any additional training cost.