Learning Rate Annealing Improves Tuning Robustness in Stochastic Optimization

The learning rate in stochastic gradient methods is a critical hyperparameter that is notoriously costly to tune via standard grid search, especially for training modern large-scale models with billions of parameters. We identify a theoretical advantage of learning rate annealing schemes that decay the learning rate to zero at a polynomial rate, such as the widely-used cosine schedule, by demonstrating their increased robustness to initial parameter misspecification due to a coarse grid search. We present an analysis in a stochastic convex optimization setup demonstrating that the convergence rate of stochastic gradient descent with annealed schedules depends *sublinearly* on the multiplicative misspecification factor $\rho$ (i.e., the grid resolution), achieving a rate of $O(\rho^{1/(2p+1)}/\sqrt{T})$ where $p$ is the degree of polynomial decay and $T$ is the number of steps. This is in contrast to the $O(\rho/\sqrt{T})$ rate obtained under the inverse-square-root and fixed stepsize schedules, which depend linearly on $\rho$. Experiments confirm the increased robustness compared to tuning with a fixed stepsize, that has significant implications for the computational overhead of hyperparameter search in practical training scenarios.

In modern machine-learning training, one often needs to choose a learning rate, which controls how much the model learns from each training update step. This knob is often set by trying many values, which can be wasteful or impossible when each training run is expensive. We study whether common schedules that gradually lower the learning rate, especially cosine annealing and polynomial decay, make this choice less fragile. Our analysis shows that they do: if the starting learning rate is too large because the search was coarse, these schedules can still bring the update size into a safer range later in training. A fixed learning rate, by contrast, can suffer in direct proportion to how badly the value was chosen. We prove this effect in a standard mathematical model of learning from noisy updates. This suggests that widely used training schedules can reduce tuning costs and make model training more reliable when compute is limited. We also perform experiments on both a synthetic prediction task and image classification, which provide empirical support for our theory.