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

Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, Francis Bach
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:53267-53294, 2025.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-schaipp25a, title = {The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training}, author = {Schaipp, Fabian and H\"{a}gele, Alexander and Taylor, Adrien and Simsekli, Umut and Bach, Francis}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {53267--53294}, year = {2025}, editor = {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, volume = {267}, series = {Proceedings of Machine Learning Research}, month = {13--19 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v267/main/assets/schaipp25a/schaipp25a.pdf}, url = {https://proceedings.mlr.press/v267/schaipp25a.html}, 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.} }
Endnote
%0 Conference Paper %T The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training %A Fabian Schaipp %A Alexander Hägele %A Adrien Taylor %A Umut Simsekli %A Francis Bach %B Proceedings of the 42nd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2025 %E Aarti Singh %E Maryam Fazel %E Daniel Hsu %E Simon Lacoste-Julien %E Felix Berkenkamp %E Tegan Maharaj %E Kiri Wagstaff %E Jerry Zhu %F pmlr-v267-schaipp25a %I PMLR %P 53267--53294 %U https://proceedings.mlr.press/v267/schaipp25a.html %V 267 %X 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.
APA
Schaipp, F., Hägele, A., Taylor, A., Simsekli, U. & Bach, F.. (2025). The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:53267-53294 Available from https://proceedings.mlr.press/v267/schaipp25a.html.

Related Material