The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure

By Rong Ge, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli. 2019

188 citations recorded by the source corpus.

There is a stark disparity between the step size schedules used in practical large scale machine learning and those that are considered optimal by the theory of stochastic approximation. In theory, most results utilize polynomially decaying learning rate schedules, while, in practice, the "Step Decay" schedule is among the most popular schedules, where the learning rate is cut every constant number of epochs (i.e. this is a geometrically decaying schedule). This work examines the step-decay schedule for the stochastic optimization problem of streaming least squares regression (both in the non-strongly convex and strongly convex case), where we show that a sharp theoretical characterization of an optimal learning rate schedule is far more nuanced than suggested by previous work. We focus specifically on the rate that is achievable when using the final iterate of stochastic gradient descent, as is commonly done in practice. Our main result provably shows that a properly tuned geometrically decaying learning rate schedule provides an exponential improvement (in terms of the condition number) over any polynomially decaying learning rate schedule. We also provide experimental support for wider applicability of these results, including for training modern deep neural networks.

Subfields: Artificial Intelligence, Computational Mechanics, Statistics and Probability

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

In Statistics and Probability

Citations in the corpus

Understanding the Role of Training Regimes in Continual Learning (2020, background): “[18] study the effectiveness of exponentially decaying learning rate and show that in its final iterate, it can achieve a near-optimal error in stochastic optimization problems.”
Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks (2018, background): “[Ge et al., 2019] studied the learning rate schedule for quadratic functions.”
On the Generalization Benefit of Noise in Stochastic Gradient Descent (2020, unknown): “Although the critical learning rate will perform poorly for high curvature directions of the loss, we can introduce learning rate decay to minimize the loss along these directions later in training (Ge et al., 2019).”