# Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 55 of its 82 references that are in the map.
- First posted on arXiv 6 Jun 2024.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Just Interpolate: Kernel "Ridgeless" Regression Can Generalize](https://mlsanity.com/paper/1808.00387/) (2018)
- [AdaGrad stepsizes: Sharp convergence over nonconvex landscapes](https://mlsanity.com/paper/1806.01811/) (2018)
- [A Simple Convergence Proof of Adam and Adagrad](https://mlsanity.com/paper/2003.02395/) (2020)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [On the Convergence of Stochastic Gradient Descent with Adaptive Stepsizes](https://mlsanity.com/paper/1805.08114/) (2018)
- [Stochastic Gradient Push for Distributed Deep Learning](https://mlsanity.com/paper/1811.10792/) (2018)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Fast and Faster Convergence of SGD for Over-Parameterized Models and an Accelerated Perceptron](https://mlsanity.com/paper/1810.07288/) (2018)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm](https://mlsanity.com/paper/1310.5715/) (2013)
- [Painless Stochastic Gradient: Interpolation, Line-Search, and Convergence Rates](https://mlsanity.com/paper/1905.09997/) (2019)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [SGD and Hogwild! Convergence Without the Bounded Gradients Assumption](https://mlsanity.com/paper/1802.03801/) (2018)
- [SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation](https://mlsanity.com/paper/2006.10311/) (2020)
- [Global convergence of the Heavy-ball method for convex optimization](https://mlsanity.com/paper/1412.7457/) (2014)
- [Stochastic (Approximate) Proximal Point Methods: Convergence, Optimality, and Adaptivity](https://mlsanity.com/paper/1810.05633/) (2018)
- [A Unified Analysis of Stochastic Momentum Methods for Deep Learning](https://mlsanity.com/paper/1808.10396/) (2018)
- [Revisiting the Polyak step size](https://mlsanity.com/paper/1905.00313/) (2019)
- [Fast Convergence of Stochastic Gradient Descent under a Strong Growth Condition](https://mlsanity.com/paper/1308.6370/) (2013)
- [Momentum and Stochastic Momentum for Stochastic Gradient, Newton, Proximal Point and Subspace Descent Methods](https://mlsanity.com/paper/1712.09677/) (2017)
- [On the insufficiency of existing momentum schemes for Stochastic Optimization](https://mlsanity.com/paper/1803.05591/) (2018)
- [Quasi-hyperbolic momentum and Adam for deep learning](https://mlsanity.com/paper/1810.06801/) (2018)
- [Training Neural Networks for and by Interpolation](https://mlsanity.com/paper/1906.05661/) (2019)
- [Understanding the Role of Momentum in Stochastic Gradient Methods](https://mlsanity.com/paper/1910.13962/) (2019)

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

- [An Adaptive Stochastic Gradient Method with Non-negative Gauss-Newton Stepsizes](https://mlsanity.com/paper/2407.04358/) (2024)
- [New Perspectives on the Polyak Stepsize: Surrogate Functions and Negative Results](https://mlsanity.com/paper/2505.20219/) (2025)
- [Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation](https://mlsanity.com/paper/2504.01898/) (2025)
- [Taking the Road Less Scheduled with Adaptive Polyak Steps](https://mlsanity.com/paper/2511.07767/) (2025)
- [Improved Convergence in Parameter-Agnostic Error Feedback through Momentum](https://mlsanity.com/paper/2511.14501/) (2025)
- [Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients](https://mlsanity.com/paper/2512.02342/) (2025)
- [MARINA-P: Superior Performance in Non-smooth Federated Optimization with Adaptive Stepsizes](https://mlsanity.com/paper/2412.17082/) (2024)
- [First-ish Order Methods: Hessian-aware Scalings of Gradient Descent](https://mlsanity.com/paper/2502.03701/) (2025)
- [Adaptive Memory Momentum via a Model-Based Framework for Deep Learning Optimization](https://mlsanity.com/paper/2510.04988/) (2025)
- [Fix the Loss, Not the Radius: Rethinking the Adversarial Perturbation of Sharpness-Aware Minimization](https://mlsanity.com/paper/2605.10183/) (2026)
- [Adaptive Sharpness-Aware Minimization with a Polyak-type Step size: A Theory-Grounded Scheduler](https://mlsanity.com/paper/2606.01827/) (2026)
- [Polyak-Type Extragradient Methods for Monotone Root-Finding Problems](https://mlsanity.com/paper/2609.26581/) (2026)
- [Stochastic Heavy Ball with Polyak Step Size and Armijo Line Search: A General Convergence Analysis](https://mlsanity.com/paper/2609.36668/) (2026)

## Abstract (arXiv)

Stochastic gradient descent with momentum, also known as Stochastic Heavy Ball method (SHB), is one of the most popular algorithms for solving large-scale stochastic optimization problems in various machine learning tasks. In practical scenarios, tuning the step-size and momentum parameters of the method is a prohibitively expensive and time-consuming process. In this work, inspired by the recent advantages of stochastic Polyak step-size in the performance of stochastic gradient descent (SGD), we propose and explore new Polyak-type variants suitable for the update rule of the SHB method. In particular, using the Iterate Moving Average (IMA) viewpoint of SHB, we propose and analyze three novel step-size selections: MomSPSmax, MomDecSPS, and MomAdaSPS. For MomSPSmax, we provide convergence guarantees for SHB to a neighborhood of the solution for convex and smooth problems (without assuming interpolation). If interpolation is also satisfied, then using MomSPSmax, SHB converges to the true solution at a fast rate matching the deterministic HB. The other two variants, MomDecSPS and MomAdaSPS, are the first adaptive step-size for SHB that guarantee convergence to the exact minimizer - without a priori knowledge of the problem parameters and without assuming interpolation. Our convergence analysis of SHB is tight and obtains the convergence guarantees of stochastic Polyak step-size for SGD as a special case. We supplement our analysis with experiments validating our theory and demonstrating the effectiveness and robustness of our algorithms.