The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training

TL;DR

The paper links practical learning-rate scheduling for large-scale training to a suboptimality bound from non-smooth convex optimization, explaining why cosine and warmup-stable-decay (wsd) perform similarly and why cooldown helps. It derives a bound for the wsd schedule that reduces logarithmic terms and demonstrates how this theory can guide schedule design, including continued-training horizon extension and LR transfer across schedules. The authors validate the theory with theoretical simulations and real-model experiments (124M and 210M Llama-style models), showing tangible improvements when using theory-informed schedules. Overall, the work suggests that convex optimization theory can provide actionable guidance for tuning LR schedules in deep learning, even in non-convex settings, by focusing on gradient norms and cooldown dynamics.

Abstract

We show that learning-rate schedules for large model training behave surprisingly similar to a performance bound from non-smooth convex optimization theory. We provide a bound for the constant schedule with linear cooldown; in particular, the practical benefit of cooldown is reflected in the bound due to the absence of logarithmic terms. Further, we show that this surprisingly close match between optimization theory and practice can be exploited for learning-rate tuning: we achieve noticeable improvements for training 124M and 210M Llama-type models by (i) extending the schedule for continued training with optimal learning-rate, and (ii) transferring the optimal learning-rate across schedules.

Figures (28)

• Figure 1: Strikingly similar: Validation loss for a 210M Llama model trained with AdamW(left) and the theoretical suboptimality bound \ref{['eqn:bound-schedule-convex']} from convex optimization (right). Both plots show wsd and cosine schedule with different training lengths $T$, and with base learning-rate of cosine being twice as large as for wsd.
• Figure 2: Schedule (left) and theoretical bound (right) for cosine and wsd, and various $T$, with base learning-rate $\gamma^\star$.
• Figure 3: Optimal base learning-rate decays with inverse square-root of training horizon $T$(right). For cosine, it is roughly twice as large as for wsd (as $0.92/0.47 \approx 2$). The dashed curve in the right-hand side plot is obtained with a least-squares fit.
• Figure 4: (Left) Optimal base learning-rate increases with cooldown fraction. (Right) For fixed $\gamma$, the optimal cooldown fraction can be smaller than $1$. The analogous curves for real experiments with similar parabola shapes are in \ref{['fig:reanalysis-cooldown-length']}.
• Figure 5: Schedule (left) and theoretical convergence (right) for varying cooldown fraction. With optimal base learning-rate $\gamma^\star$, starting the cooldown at $T_0=1$ is optimal. \ref{['fig:reanalysis-cooldown-length']} shows the analogous plot for real experiments with the same behavior.

Theorems & Definitions (19)

• Theorem 3.1: cf. Thm. 10 from Defazio2023a
• Remark 3.2
• Corollary 3.3
• Theorem 3.4
• Theorem 1.1
• Lemma 4.1: Lemma 5 from Defazio2023a
• Lemma 4.2
• proof
• Lemma 4.3
• Lemma 5.1