Notes

Cosine Annealing Schedule

Training a neural network is not just about choosing a learning rate; it is about choosing how that rate changes as the model improves. A cosine annealing schedule starts with relatively large updates and gradually shrinks them along a smooth cosine-shaped curve, letting training move quickly at first and settle carefully near the end.

How the schedule works

At training step or epoch t, the learning rate is commonly set by:

lr(t) = lr_min + 0.5 × (lr_max - lr_min) ×
        (1 + cos(π × t / T_max))

lr_max is the initial learning rate, lr_min is the final floor, and T_max is the number of steps or epochs in the decay cycle. The cosine begins flat, falls more quickly through the middle, then flattens again near the minimum. That gentle ending is useful: parameter updates become small enough to refine a good solution instead of repeatedly stepping past it.

Why the shape helps

A fixed high learning rate keeps the optimizer bouncing around a low-loss region; a fixed tiny one makes early training painfully slow. Cosine annealing provides both phases without abrupt drops:

  • Large early updates explore the loss landscape and make rapid progress.
  • Smaller late updates reduce instability and support fine adjustment.
  • The smooth curve avoids the sudden training disruption caused by step-decay schedules.
Use in real training

For example, an AdamW-trained transformer might warm up for a short period, then use cosine annealing until the final training step. In PyTorch, torch.optim.lr_scheduler.CosineAnnealingLR implements the single-decay version. CosineAnnealingWarmRestarts instead repeatedly resets the rate upward; these restarts can help the optimizer leave a mediocre region, but they also make the learning-rate pattern less quiet near the end. Set T_max to match the intended training duration: ending the cycle far too early can leave learning rates too small while useful progress remains.

Cosine annealing is a learning-rate schedule that smoothly decreases the learning rate from a maximum to a minimum along a half-cosine curve over a fixed training interval. Variants can periodically restart the rate at a higher value. It enables large early updates and progressively finer late-stage optimization, helping training converge stably without abrupt learning-rate drops.

Imagine learning to throw a ball into a basket. At first, you make big changes to your throw because you are far from the goal. Later, once you are close, you make smaller, gentler adjustments so you do not overshoot.

A cosine annealing schedule gives an AI the same kind of pacing while it learns. It starts with relatively bold adjustments, then gradually reduces their size along a smooth, curved path shaped like part of a cosine wave. This helps the network learn quickly early on, then settle into more careful fine-tuning later. The result can be steadier training and a better final solution.