Optimal Linear Decay Learning Rate Schedules and Further Refinements

TL;DR

The paper addresses the gap between theory and practice in learning-rate scheduling by deriving last-iterate convergence guarantees and proposing problem-adaptive, data-dependent refinements. It shows that linear decay is theoretically optimal for SGD in the last-iterate sense and provides a practical refinement procedure based on observed gradient norms, including median smoothing, to tailor schedules to specific tasks. The approach extends to per-coordinate optimizers like Adam and demonstrates strong empirical performance across convex problems, deep learning tasks, and large-language-model benchmarks, often outperforming cosine and other standard schedules. Limitations include degenerate refinements on some tasks and the need for prior runs to estimate gradient norms, but overall the method offers actionable guidance for designing LR schedules in practice.

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.

Figures (7)

• Figure 1: Example gradient norm sequences (top row) and the resulting refined schedules given by Algorithm \ref{['alg:refinement']} (bottom row). Black dashed line at $y=1$ shown for reference. Percentage of runtime on the x-axis.
• Figure 2: Training curves on CIFAR-10 for a sweep over inverse-sqrt and inverse-linear hyper-parameters. A linear-decay schedule baseline is shown in orange. All combinations are out-performed by the linear-decay schedule.
• Figure 3: Gradient Norm sequences and the resulting Refined schedules, generated using an initial linear decay schedule with warmup for the initial run. Log scale views are inset for scale reference.
• Figure 4: Normalized learning rate schedules learned on vanilla Transformer-based LLM trained on C4 dataset. Left: schedules obtained when setting $w_t\propto \|g_t\|^{-2}$. Right: schedules when $w_t\propto 1/\|g_t\|_1$. We find that $\ell_1$ norm is much more consistent across model sizes and baseline schedules. Thus, we used $w_t\propto 1/\|g_t\|_1$ weighting from linear baselines to obtained Refined schedule mentioned in Table \ref{['tab:llm']}. The x axis is training steps.
• Figure 5: Limitations of refinement: if model over-fits, our method produces degenerate schedules.

Theorems & Definitions (24)

• Theorem 1
• Corollary 1
• Theorem 2
• Lemma 3
• proof
• Lemma 4
• proof
• Remark 5
• Lemma 5
• proof