Cosine Learning Rate Schedule › Cosine Learning Rate Schedule

The learning rate is one of the most consequential hyperparameters in training neural networks. Set it too high and the model diverges; set it too low and it crawls toward a poor solution. As we saw in the Learning Rate Warmup chapter, starting with a small learning rate and gradually ramping it up helps the optimizer find a stable initial direction. And as covered in Learning Rate Decay, reducing the learning rate over time lets the model settle into the fine-grained structure of the loss landscape. The cosine learning rate schedule combines both ideas into a single, smooth function that has become the dominant choice for training large language models and vision transformers alike.

The central intuition is elegant. Instead of dropping the learning rate in abrupt steps or along a straight line from start to finish, the cosine schedule follows the shape of a cosine curve: it starts high, descends gradually at first, accelerates through the middle of training, and then slows again to arrive gently near zero. This smooth trajectory avoids the shock of sudden step reductions while also being faster than a linear decay in the early and late phases where the model benefits most from larger and then smaller steps respectively.

Why does the cosine curve, of all curves, suit this purpose so well? The answer comes down to the geometry of the loss landscape and the characteristic phases of deep learning optimization. Early in training, the model is far from any meaningful minimum, and a learning rate that stays moderately high lets the optimizer take bigger exploratory steps. In the middle phase, the model is making substantial parameter adjustments as it locates the broad basin of a good solution. Near the end, the model is refining its parameters within that basin, and a very small learning rate lets it settle precisely without oscillating around the minimum. The cosine curve allocates the most aggressive decay to this middle phase and decelerates gently at both ends, which closely mirrors what practitioners discovered empirically through years of trial and error with step decay schedules.

This chapter covers the cosine decay formula in full detail, explains when and why to use cosine annealing with warm restarts, walks through the role each schedule parameter plays, and compares cosine decay against linear and step decay alternatives. We will implement these schedules from scratch and then show how to integrate them with PyTorch's built-in scheduler API. We will also examine how this schedule appears in the published training configurations of GPT, LLaMA, Mistral, and other prominent language models.