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

arXiv: https://arxiv.org/abs/2501.18965 · HTML: https://mlsanity.com/paper/2501.18965/

## Key facts

- Cited by 33 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 44 references that are in the map.
- First posted on arXiv 31 Jan 2025.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [Gradient Descent Learns Linear Dynamical Systems](https://mlsanity.com/paper/1609.05191/) (2016)
- [Conservative set valued fields, automatic differentiation, stochastic gradient method and deep learning](https://mlsanity.com/paper/1909.10300/) (2019)
- [Chinchilla Scaling: A replication attempt](https://mlsanity.com/paper/2404.10102/) (2024)
- [Optimal Linear Decay Learning Rate Schedules and Further Refinements](https://mlsanity.com/paper/2310.07831/) (2023)
- [Smooth Strongly Convex Interpolation and Exact Worst-case Performance of First-order Methods](https://mlsanity.com/paper/1502.05666/) (2015)
- [(S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability](https://mlsanity.com/paper/2302.08982/) (2023)
- [A Second look at Exponential and Cosine Step Sizes: Simplicity, Adaptivity, and Performance](https://mlsanity.com/paper/2002.05273/) (2020)
- [Exact convergence rate of the last iterate in subgradient methods](https://mlsanity.com/paper/2307.11134/) (2023)
- [Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling](https://mlsanity.com/paper/2409.15156/) (2024)
- [Aiming towards the minimizers: fast convergence of SGD for overparametrized problems](https://mlsanity.com/paper/2306.02601/) (2023)
- [Training trajectories, mini-batch losses and the curious role of the learning rate](https://mlsanity.com/paper/2301.02312/) (2023)
- [Exact Worst-case Performance of First-order Methods for Composite Convex Optimization](https://mlsanity.com/paper/1512.07516/) (2015)
- [PEPit: computer-assisted worst-case analyses of first-order optimization methods in Python](https://mlsanity.com/paper/2201.04040/) (2022)
- [Loss Landscape Characterization of Neural Networks without Over-Parametrization](https://mlsanity.com/paper/2410.12455/) (2024)
- [No More Adam: Learning Rate Scaling at Initialization is All You Need](https://mlsanity.com/paper/2412.11768/) (2024)
- [Why you don't overfit, and don't need Bayes if you only train for one epoch](https://mlsanity.com/paper/2411.14478/) (2024)

## Built on by (most cited first, 30 of 33)

- [A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules](https://mlsanity.com/paper/2503.12811/) (2025)
- [Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks](https://mlsanity.com/paper/2507.02119/) (2025)
- [A Defense of the Quadratic Model](https://mlsanity.com/paper/2607.21716/) (2026)
- [Replaying pre-training data improves fine-tuning](https://mlsanity.com/paper/2603.04964/) (2026)
- [Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler](https://mlsanity.com/paper/2508.01483/) (2025)
- [Anytime Pretraining: Horizon-Free Learning-Rate Schedules with Weight Averaging](https://mlsanity.com/paper/2602.03702/) (2026)
- [Deriving Hyperparameter Scaling Laws via Modern Optimization Theory](https://mlsanity.com/paper/2603.15958/) (2026)
- [Scaling with Collapse: Efficient and Predictable Training of LLM Families](https://mlsanity.com/paper/2509.25087/) (2025)
- [Optimal Learning Rate Schedules under Functional Scaling Laws: Power Decay and Warmup-Stable-Decay](https://mlsanity.com/paper/2602.06797/) (2026)
- [Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training](https://mlsanity.com/paper/2503.02844/) (2025)
- [Optimization Benchmark for Diffusion Models on Dynamical Systems](https://mlsanity.com/paper/2510.19376/) (2025)
- [On the Role of Batch Size in Stochastic Conditional Gradient Methods](https://mlsanity.com/paper/2603.21191/) (2026)
- [ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models](https://mlsanity.com/paper/2605.19095/) (2026)
- [The Design Space of Tri-Modal Masked Diffusion Models](https://mlsanity.com/paper/2602.21472/) (2026)
- [Universal Dynamics of Warmup Stable Decay: understanding WSD beyond Transformers](https://mlsanity.com/paper/2601.09000/) (2026)
- [AdaLRS: Loss-Guided Adaptive Learning Rate Search for Efficient Foundation Model Pretraining](https://mlsanity.com/paper/2506.13274/) (2025)
- [Convex Dominance in Deep Learning I: A Scaling Law of Loss and Learning Rate](https://mlsanity.com/paper/2602.07145/) (2026)
- [Optimistic Dual Averaging Unifies Modern Optimizers](https://mlsanity.com/paper/2605.11172/) (2026)
- [Anytime Training with Schedule-Free Spectral Optimization](https://mlsanity.com/paper/2605.23061/) (2026)
- [Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors](https://mlsanity.com/paper/2606.25971/) (2026)
- [Predicting Training Re-evaluation Curves Enables Effective Data Curriculums for LLMs](https://mlsanity.com/paper/2509.25380/) (2025)
- [Scaling depth capacity via zero/one-layer model expansion](https://mlsanity.com/paper/2511.04981/) (2025)
- [Scaling and Transferability of Annealing Strategies in Large Language Model Training](https://mlsanity.com/paper/2512.13705/) (2025)
- [Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts](https://mlsanity.com/paper/2604.19835/) (2026)
- [How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size](https://mlsanity.com/paper/2607.01487/) (2026)
- [Towards joint scaling laws with optimal batch size schedules](https://mlsanity.com/paper/2607.27731/) (2026)
- [Step-Size Stability in Stochastic Optimization: A Theoretical Perspective](https://mlsanity.com/paper/2602.09842/) (2026)
- [WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training](https://mlsanity.com/paper/2607.10959/) (2026)
- [Same Loss, Same Noise, Opposite Schedules: Noise Structure and Optimizer Normalization Jointly Determine Whether Learning-Rate Cooldown Helps](https://mlsanity.com/paper/2607.12360/) (2026)
- [Understanding the Subspace Stabilization of the Hessian and Gradient Covariance Matrix](https://mlsanity.com/paper/2609.31983/) (2026)

## Abstract (arXiv)

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.