Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence

Abstract

We propose a stochastic variant of the classical Polyak step-size (Polyak, 1987) commonly used in the subgradient method. Although computing the Polyak step-size requires knowledge of the optimal function values, this information is readily available for typical modern machine learning applications. Consequently, the proposed stochastic Polyak step-size (SPS) is an attractive choice for setting the learning rate for stochastic gradient descent (SGD). We provide theoretical convergence guarantees for SGD equipped with SPS in different settings, including strongly convex, convex and non-convex functions. Furthermore, our analysis results in novel convergence guarantees for SGD with a constant step-size. We show that SPS is particularly effective when training over-parameterized models capable of interpolating the training data. In this setting, we prove that SPS enables SGD to converge to the true solution at a fast rate without requiring the knowledge of any problem-dependent constants or additional computational overhead. We experimentally validate our theoretical results via extensive experiments on synthetic and real datasets. We demonstrate the strong performance of SGD with SPS compared to state-of-the-art optimization methods when training over-parameterized models.

arXiv commentProceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS) 2021

References in corpus (16)

• HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent
• Making Gradient Descent Optimal for Strongly Convex Stochastic Optimization
• Stochastic Gradient Descent for Non-smooth Optimization: Convergence Results and Optimal Averaging Schemes
• Guaranteed Matrix Completion via Non-convex Factorization
• Adaptive Gradient Methods with Dynamic Bound of Learning Rate
• AdaGrad stepsizes: Sharp convergence over nonconvex landscapes
• Barzilai-Borwein Step Size for Stochastic Gradient Descent
• The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning
• SGD and Hogwild! Convergence Without the Bounded Gradients Assumption
• Stop Wasting My Gradients: Practical SVRG
• The importance of better models in stochastic optimization
• L4: Practical loss-based stepsize adaptation for deep learning
• Revisiting the Polyak step size
• Stochastic Gradient Descent with Polyak's Learning Rate
• Training Neural Networks for and by Interpolation
• Randomized Iterative Methods for Linear Systems: Momentum, Inexactness and Gossip

Cited by in corpus (14)

• A Unified Theory of Decentralized SGD with Changing Topology and Local Updates
• Weight-Sharing Neural Architecture Search: A Battle to Shrink the Optimization Gap
• Counting Cows: Tracking Illegal Cattle Ranching From High-Resolution Satellite Imagery
• Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ball
• Training Neural Networks for and by Interpolation
• Multi-modal anticipation of stochastic trajectories in a dynamic environment with Conditional Variational Autoencoders
• SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation
• On the Convergence of Stochastic Gradient Descent with Bandwidth-based Step Size
• On the Convergence of Step Decay Step-Size for Stochastic Optimization
• AI-SARAH: Adaptive and Implicit Stochastic Recursive Gradient Methods
• Doubly Adaptive Scaled Algorithm for Machine Learning Using Second-Order Information
• About some works of Boris Polyak on convergence of gradient methods and their development
• A Probabilistically Motivated Learning Rate Adaptation for Stochastic Optimization
• Stochastic Polyak Stepsize with a Moving Target