# Optimal learning rate schedules in high-dimensional non-convex optimization problems

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 63 references that are in the map.
- First posted on arXiv 9 Feb 2022.

## Builds on (most cited first, 30 of 32)

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Cyclical Learning Rates for Training Neural Networks](https://mlsanity.com/paper/1506.01186/) (2015)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [The Loss Surfaces of Multilayer Networks](https://mlsanity.com/paper/1412.0233/) (2014)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [The Marginal Value of Adaptive Gradient Methods in Machine Learning](https://mlsanity.com/paper/1705.08292/) (2017)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [Statistical physics of inference: Thresholds and algorithms](https://mlsanity.com/paper/1511.02476/) (2015)
- [Improving Generalization Performance by Switching from Adam to SGD](https://mlsanity.com/paper/1712.07628/) (2017)
- [Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup](https://mlsanity.com/paper/1906.08632/) (2019)
- [A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillation](https://mlsanity.com/paper/1810.13243/) (2018)
- [Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks](https://mlsanity.com/paper/1806.06763/) (2018)
- [The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares](https://mlsanity.com/paper/1904.12838/) (2019)
- [Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape](https://mlsanity.com/paper/1907.02911/) (2019)
- [Dynamical mean-field theory for stochastic gradient descent in Gaussian mixture classification](https://mlsanity.com/paper/2006.06098/) (2020)
- [Is SGD a Bayesian sampler? Well, almost](https://mlsanity.com/paper/2006.15191/) (2020)
- [A Loss Curvature Perspective on Training Instability in Deep Learning](https://mlsanity.com/paper/2110.04369/) (2021)
- [The landscape of the spiked tensor model](https://mlsanity.com/paper/1711.05424/) (2017)
- [Deep Learning Theory Review: An Optimal Control and Dynamical Systems Perspective](https://mlsanity.com/paper/1908.10920/) (2019)
- [Who is Afraid of Big Bad Minima? Analysis of Gradient-Flow in a Spiked Matrix-Tensor Model](https://mlsanity.com/paper/1907.08226/) (2019)
- [The effective noise of Stochastic Gradient Descent](https://mlsanity.com/paper/2112.10852/) (2021)
- [Marvels and Pitfalls of the Langevin Algorithm in Noisy High-dimensional Inference](https://mlsanity.com/paper/1812.09066/) (2018)
- [Complex energy landscapes in spiked-tensor and simple glassy models: ruggedness, arrangements of local minima and phase transitions](https://mlsanity.com/paper/1804.02686/) (2018)
- [The Effect of Network Width on Stochastic Gradient Descent and Generalization: an Empirical Study](https://mlsanity.com/paper/1905.03776/) (2019)

## Built on by (most cited first, 4 of 4)

- [Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks](https://mlsanity.com/paper/2205.09653/) (2022)
- [Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks](https://mlsanity.com/paper/2507.02119/) (2025)
- [The RL Perceptron: Generalisation Dynamics of Policy Learning in High Dimensions](https://mlsanity.com/paper/2306.10404/) (2023)
- [A statistical physics framework for optimal learning](https://mlsanity.com/paper/2507.07907/) (2025)

## Abstract (arXiv)

Learning rate schedules are ubiquitously used to speed up and improve optimisation. Many different policies have been introduced on an empirical basis, and theoretical analyses have been developed for convex settings. However, in many realistic problems the loss-landscape is high-dimensional and non convex -- a case for which results are scarce. In this paper we present a first analytical study of the role of learning rate scheduling in this setting, focusing on Langevin optimization with a learning rate decaying as η(t)=t-β. We begin by considering models where the loss is a Gaussian random function on the N-dimensional sphere (Nrightarrow ∞), featuring an extensive number of critical points. We find that to speed up optimization without getting stuck in saddles, one must choose a decay rate β<1, contrary to convex setups where β=1 is generally optimal. We then add to the problem a signal to be recovered. In this setting, the dynamics decompose into two phases: an exploration phase where the dynamics navigates through rough parts of the landscape, followed by a convergence phase where the signal is detected and the dynamics enter a convex basin. In this case, it is optimal to keep a large learning rate during the exploration phase to escape the non-convex region as quickly as possible, then use the convex criterion β=1 to converge rapidly to the solution. Finally, we demonstrate that our conclusions hold in a common regression task involving neural networks.