Cosine annealing shrinks the learning rate following the smooth curve of a cosine function โ starting with slow decay, accelerating through the middle, and gently flattening as it approaches its minimum. It's one of the most popular schedules in modern deep learning, especially for training from scratch.
Formula
\(T\) is the total number of steps (or epochs) in the schedule (or one cycle of it), \(\eta_{\max}\) and \(\eta_{\min}\) are the starting and ending learning rates. At \(t=0\): \(\cos(0)=1\), giving \(\eta_0=\eta_{\max}\). At \(t=T\): \(\cos(\pi)=-1\), giving \(\eta_T=\eta_{\min}\).
Graph
Slow decay at both the start and end, with the fastest decay through the middle โ the smooth "S" shape of half a cosine wave.
Why the Cosine Shape Is Popular
The slow start gives the optimizer time to keep making meaningful progress at a still-high learning rate before decay meaningfully kicks in. The slow finish provides an extended period of very gentle fine-tuning near the end, letting the model settle carefully into a good minimum. This shape has empirically been found to often outperform simpler schedules (like linear or step decay) across a range of architectures.
Warm Restarts โ A Common Extension
"Cosine Annealing with Warm Restarts" (SGDR) periodically resets \(\eta_t\) back up to \(\eta_{\max}\) partway through training, then anneals down again โ running multiple cosine cycles within one training run. Each restart can help the optimizer escape a local minimum it may have settled into, exploring a nearby region of the loss surface again before annealing down to a potentially better one.
Code
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=0.001)
for epoch in range(30):
# ... training loop for this epoch ...
scheduler.step()
print(scheduler.get_last_lr())
Common Mistakes
- Setting \(T_{\max}\) inconsistently with the actual number of training epochs โ if training stops well before \(t=T\), the learning rate never reaches its intended minimum, and if it runs longer, the schedule (without restarts) has no further decay to offer.
Interview Relevance
Q: "Why might cosine annealing be preferred over step decay for training a model from scratch?" Its smooth, non-linear shape spends more time at both high learning rates (fast early progress) and very low learning rates (careful late-stage fine-tuning) compared to a schedule with abrupt, evenly-spaced drops โ this shape has been empirically shown to often produce better final performance across many architectures, and it avoids the small loss-curve discontinuities that step decay's sudden drops can cause.
Practice Question
With \(\eta_{\max}=0.1\), \(\eta_{\min}=0\), \(T=20\), what is \(\eta_t\) at \(t=10\) (the halfway point)? (Hint: \(\cos(\pi/2)=0\).)