# A Stochastic Proximal Polyak Step Size

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 28 references that are in the map.
- First posted on arXiv 12 Jan 2023.

## Builds on (most cited first, 16 of 16)

- [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)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [Three Mechanisms of Weight Decay Regularization](https://mlsanity.com/paper/1810.12281/) (2018)
- [SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation](https://mlsanity.com/paper/2006.10311/) (2020)
- [Stochastic (Approximate) Proximal Point Methods: Convergence, Optimality, and Adaptivity](https://mlsanity.com/paper/1810.05633/) (2018)
- [Error bounds, quadratic growth, and linear convergence of proximal methods](https://mlsanity.com/paper/1602.06661/) (2016)
- [Efficiency of minimizing compositions of convex functions and smooth maps](https://mlsanity.com/paper/1605.00125/) (2016)
- [Revisiting the Polyak step size](https://mlsanity.com/paper/1905.00313/) (2019)
- [Understanding AdamW through Proximal Methods and Scale-Freeness](https://mlsanity.com/paper/2202.00089/) (2022)
- [Training Neural Networks for and by Interpolation](https://mlsanity.com/paper/1906.05661/) (2019)
- [Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution](https://mlsanity.com/paper/2205.04583/) (2022)
- [Stochastic Gradient Descent with Polyak's Learning Rate](https://mlsanity.com/paper/1903.08688/) (2019)

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

- [An Exploration of Non-Euclidean Gradient Descent: Muon and its Many Variants](https://mlsanity.com/paper/2510.09827/) (2025)
- [MoMo: Momentum Models for Adaptive Learning Rates](https://mlsanity.com/paper/2305.07583/) (2023)
- [Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance](https://mlsanity.com/paper/2406.04142/) (2024)
- [Locally Adaptive Federated Learning](https://mlsanity.com/paper/2307.06306/) (2023)
- [Stochastic Gradient Descent with Preconditioned Polyak Step-size](https://mlsanity.com/paper/2310.02093/) (2023)
- [Analysis of an Idealized Stochastic Polyak Method and its Application to Black-Box Model Distillation](https://mlsanity.com/paper/2504.01898/) (2025)
- [Stochastic Weakly Convex Optimization Beyond Lipschitz Continuity](https://mlsanity.com/paper/2401.13971/) (2024)
- [Demystifying the Myths and Legends of Nonconvex Convergence of SGD](https://mlsanity.com/paper/2310.12969/) (2023)
- [On the Convergence of Federated Learning Algorithms without Data Similarity](https://mlsanity.com/paper/2403.02347/) (2024)
- [Taking the Road Less Scheduled with Adaptive Polyak Steps](https://mlsanity.com/paper/2511.07767/) (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)
- [Adaptive Sharpness-Aware Minimization with a Polyak-type Step size: A Theory-Grounded Scheduler](https://mlsanity.com/paper/2606.01827/) (2026)
- [Layer-wise Adaptive Step-Sizes for Stochastic First-Order Methods for Deep Learning](https://mlsanity.com/paper/2305.13664/) (2023)
- [Step-Size Stability in Stochastic Optimization: A Theoretical Perspective](https://mlsanity.com/paper/2602.09842/) (2026)

## Abstract (arXiv)

Recently, the stochastic Polyak step size (SPS) has emerged as a competitive adaptive step size scheme for stochastic gradient descent. Here we develop ProxSPS, a proximal variant of SPS that can handle regularization terms. Developing a proximal variant of SPS is particularly important, since SPS requires a lower bound of the objective function to work well. When the objective function is the sum of a loss and a regularizer, available estimates of a lower bound of the sum can be loose. In contrast, ProxSPS only requires a lower bound for the loss which is often readily available. As a consequence, we show that ProxSPS is easier to tune and more stable in the presence of regularization. Furthermore for image classification tasks, ProxSPS performs as well as AdamW with little to no tuning, and results in a network with smaller weight parameters. We also provide an extensive convergence analysis for ProxSPS that includes the non-smooth, smooth, weakly convex and strongly convex setting.