Cosine Learning Rate Schedule › The Loss Landscape and Why Decay Shape Matters

To develop real intuition for why the cosine schedule works, it helps to think carefully about what the loss landscape looks like during different phases of training and what kind of learning rate behavior each phase calls for.

Deep neural network loss surfaces are high-dimensional and highly non-convex. They contain many local minima, saddle points, and flat regions where gradients are small. Empirical research over the past decade has produced a rough picture of how gradient descent moves through this landscape during training.

In the earliest phase, immediately after random initialization, the model's parameters are far from any interesting region of the loss surface. The loss is high, and gradients point roughly in directions that reduce it. At this point, the exact step size matters less than staying stable and making clear directional progress. This is why warmup is so valuable: large learning rates on random initializations cause the optimizer to take wild, unstable steps that can overshoot into regions of catastrophically high loss. A small initial rate lets the optimizer establish a sensible direction before committing to larger steps.

Once the model is making meaningful progress, it enters a phase of broad exploration. The optimizer is traversing large regions of parameter space, and the loss is decreasing substantially with each step. The learning rate should be high enough to cover ground quickly but not so high that the optimizer overshoots the broad basin it is trying to find. This phase benefits from a learning rate that stays near its peak and only decreases slowly, which is exactly what the cosine schedule provides in the first quarter of training.

As training continues, the model begins to find the broad basin of a good solution. The loss is now decreasing more slowly, and gradients are pointing more consistently downward rather than in a noisy, exploratory fashion. The optimizer needs to make progressively finer adjustments. A learning rate that decreases at an accelerating pace during this phase, matching the cosine curve's steepest descent at the midpoint, aligns well with this transition from exploration to exploitation.

In the final phase, the model is already near a local minimum and needs to settle in. The gradient directions are consistent, but the useful step sizes are tiny. A very small learning rate lets the optimizer make these final micro-adjustments without overshooting the minimum. The cosine schedule's gentle approach to $\eta_{\min}$ in the final quarter of training mirrors this behavior perfectly.

The geometry of the minimum matters for generalization. A sharp minimum is one where the loss surface has steep walls on all sides, meaning the loss increases quickly as you move away from the optimal point. A flat minimum is one where the loss surface is wide and shallow, and the loss increases only slowly in all directions.

Sharp minima are problematic for generalization. When a model trained on one dataset is deployed on a slightly different distribution (test data, slightly different user inputs), its parameters are never exactly at the training minimum. A flat minimum tolerates this imprecision gracefully because the loss changes slowly. A sharp minimum punishes it severely because even a small perturbation takes the model to a point where loss is much higher.

This observation, developed rigorously by Hochreiter and Schmidhuber in the 1990s and revisited extensively in the modern deep learning era, has direct implications for learning rate scheduling. Larger learning rates tend to escape sharp minima because the optimizer takes steps large enough to hop over the steep walls. Smaller learning rates can get trapped in them because each step is too small to escape. The cosine schedule's slow early decay keeps the learning rate in a regime that favors flat minima for longer, before gradually reducing to the fine-grained regime that allows precise convergence.

This is also why warm restarts improve generalization: the sudden increase in learning rate at each restart gives the optimizer a chance to escape any sharp minimum it has found and continue searching for flatter, wider basins before the next round of slow convergence.

In mini-batch stochastic gradient descent, gradients are noisy estimates of the true gradient computed from a small sample of the training data. The noise level in gradients interacts with the learning rate in an important way: a large learning rate magnifies gradient noise, potentially causing the optimizer to take steps that are primarily noise-driven rather than signal-driven. A small learning rate attenuates noise but also attenuates the useful signal.

This interaction creates a natural argument for learning rate decay. Early in training, with the model far from any minimum, signal-to-noise ratios are generally high: gradient directions are consistent across mini-batches because they all point away from the same high-loss initialization. As the model approaches a minimum, gradients become smaller in magnitude and more noisy relative to their mean direction, because there is less clear "which way is down" once you are in a local basin. The optimal learning rate should therefore be smaller when the model is close to a minimum than when it is far away, which is exactly what decay schedules provide.

The cosine schedule does not adapt to gradient noise directly; it decays as a function of training steps regardless of what the gradients are doing. But because training steps and proximity to a minimum are loosely correlated (more steps means more progress toward a minimum), the step-based decay works reasonably well in practice. Adaptive learning rate methods like AdamW already handle per-parameter gradient noise through their moment estimates, but even they benefit from a global learning rate schedule that provides a coarse adjustment to the overall step magnitude.