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

2025/01/31 by Fabian Schaipp, Alexander Hägele, Schaipp, Fabian +7 · 2 voices · 12 citations
Computer Science · #Machine Learning and Algorithms #cs.LG #math.OC #stat.ML

paper · pdf · doi:10.48550/arxiv.2501.18965

openalex publication_date 2025/01/31 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/01

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.

Citations

• No More Adam: Learning Rate Scaling at Initialization is All You Need
• Chinchilla Scaling: A replication attempt
• Language Models are Few-Shot Learners
• Attention Is All You Need
• Deep Residual Learning for Image Recognition

Cited by

• A Defense of the Quadratic Model
• Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
• Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
• Step-Size Stability in Stochastic Optimization: A Theoretical Perspective
• Scaling and Transferability of Annealing Strategies in Large Language Model Training
• Deep Progressive Training: scaling up depth capacity of zero/one-layer models
• Optimization Benchmark for Diffusion Models on Dynamical Systems
• Loss-Transformation Invariance in the Damped Newton Method
• Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs
• Scaling with Collapse: Efficient and Predictable Training of LLM Families
• Towards joint scaling laws with optimal batch size schedules
• Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler

Discussions

• Learning rate schedules seem mysterious? Why is the loss going down so fast during cooldown? Turns out that this behaviour can be described with a bound from *convex, nonsmooth* optimization. A short [bsky, 30 points, 2 comments]
• arxiv.org/abs/2501.18965 [bsky, 0 points, 0 comments]

Related

• Optimal learning rate schedules in high-dimensional non-convex optimization problems
• Optimal Linear Decay Learning Rate Schedules and Further Refinements
• The Two Regimes of Deep Network Training
• Towards joint scaling laws with optimal batch size schedules
• Learning an Adaptive Learning Rate Schedule