training-methodsoptimizationdeep-learning

Learning Rate Schedules: A Deep Mathematical Dive

Beyond cosine annealing — understanding why different learning rate schedules work from an optimization theory perspective.

ML for Everyone TeamMay 5, 202611 min read

The Most Important Hyperparameter

The learning rate is arguably the single most impactful hyperparameter in deep learning. Too high and training diverges; too low and it takes forever or gets stuck in suboptimal minima.

Constant Learning Rate: When It Fails

With a constant learning rate \eta, SGD converges to a neighborhood of the minimum with radius proportional to \eta:

\mathbb{E}[f(x_T) - f^*] \leq \frac{\|x_0 - x^*\|^2}{2\eta T} + \frac{\eta L \sigma^2}{2}

The first term decreases with T (more steps), but the second term is a fixed floor set by the learning rate and gradient noise.

Warmup: Building Momentum

Learning rate warmup serves a critical mathematical purpose: in the early stages of training, the model's statistics (batch norm, attention patterns) are far from their converged values. Large gradients on random features can cause irreversible damage to the loss landscape exploration.

Warmup lets the model "find its footing" before taking large steps.

Cosine Annealing

\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{t\pi}{T}\right)\right)

Why cosine? It provides:

Slow initial decay (preserve exploration)

Fast middle decay (rapid convergence)

Slow final decay (fine-tune near the minimum)

The WSD Schedule (Warmup-Stable-Decay)

Recent LLM training has converged on WSD:

Warmup: Linear increase over ~2000 steps

Stable: Constant maximum learning rate for most of training

Decay: Cosine or linear decay in the final 10-20%

This works because large models need sustained high learning rates to explore the vast parameter space before settling.

Practical Wisdom

The optimal schedule depends on your loss landscape. Highly non-convex landscapes (deep networks) benefit from aggressive exploration followed by careful refinement. The mathematics of stochastic optimization gives us the theory; practice requires experimentation within theoretically motivated bounds.