Optimal Linear Decay Learning Rate Schedules and Further Refinements
This work addresses the problem of suboptimal learning rate schedules for practitioners in machine learning, offering a systematic and adaptive approach that improves training efficiency and performance.
The paper tackles the gap between theoretical and practical learning rate schedules by deriving optimal linear decay schedules and adaptive refinements, showing that linear decay outperforms common defaults like cosine annealing across 10 deep learning problems, LLMs, and logistic regression tasks, with adaptive methods providing further gains.
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.