When, Why and How Much? Adaptive Learning Rate Scheduling by Refinement
Aaron Defazio Ashok Cutkosky Harsh Mehta Konstantin Mishchenko
Learning rate schedules used in practice bear little resemblance to those
recommended by theory. We close much of this theory/practice gap, and as a
consequence are able to derive new problem-adaptive learning rate schedules.
Our key technical contribution is a refined analysis of learning rate schedules
for a wide class of optimization algorithms (including SGD). In contrast to
most prior works that study the convergence of the average iterate, we study
the last iterate, which is what most people use in practice. When considering
only worst-case analysis, our theory predicts that the best choice is the
linear decay schedule: a popular choice in practice that sets the stepsize
proportionally to $1 - t/T$, where $t$ is the current iteration and $T$ is the
total number of steps. To go beyond this worst-case analysis, we use the
observed gradient norms to derive schedules refined for any particular task.
These refined schedules exhibit learning rate warm-up and rapid learning rate
annealing near the end of training. Ours is the first systematic approach to
automatically yield both of these properties. We perform the most comprehensive
evaluation of learning rate schedules to date, evaluating across 10 diverse
deep learning problems, a series of LLMs, and a suite of logistic regression
problems. We validate that overall, the linear-decay schedule matches or
outperforms all commonly used default schedules including cosine annealing, and
that our schedule refinement method gives further improvements.