Optimal Linear Decay Learning Rate Schedules and Further Refinements

TL;DR

This paper bridges the gap between theory and practice in learning rate schedules, deriving problem-adaptive schedules with warm-up and annealing, validated across diverse deep learning tasks, outperforming standard schedules.

Contribution

It provides a refined theoretical analysis of learning rate schedules, introduces a systematic method for adaptive schedule refinement, and offers the most comprehensive evaluation to date.

Findings

Linear decay schedule is optimal in worst-case scenarios.

Adaptive schedules with warm-up and annealing improve training.

Linear decay outperforms cosine annealing and other default schedules.

Abstract

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 main technical contribution is a refined analysis of learning rate schedules for a wide class of optimization algorithms (including SGD). When considering only worst-case analysis, our theory predicts that the optimal choice is the linear decay schedule where the step-size is set proportional 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 outperforms all commonly used default schedules including cosine annealing. Our adaptive schedule refinement method gives further improvements.