# Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 53 references that are in the map.
- First posted on arXiv 2 Dec 2025.

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

- [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)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Attention Mechanisms in Computer Vision: A Survey](https://mlsanity.com/paper/2111.07624/) (2021)
- [Just Interpolate: Kernel "Ridgeless" Regression Can Generalize](https://mlsanity.com/paper/1808.00387/) (2018)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [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 for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation](https://mlsanity.com/paper/2006.10311/) (2020)
- [Revisiting the Polyak step size](https://mlsanity.com/paper/1905.00313/) (2019)
- [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)
- [Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution](https://mlsanity.com/paper/2205.04583/) (2022)
- [Barzilai-Borwein Step Size for Stochastic Gradient Descent](https://mlsanity.com/paper/1605.04131/) (2016)
- [Subgradient methods for sharp weakly convex functions](https://mlsanity.com/paper/1803.02461/) (2018)
- [Methods for Convex (L0,L1)-Smooth Optimization: Clipping, Acceleration, and Adaptivity](https://mlsanity.com/paper/2409.14989/) (2024)
- [MoMo: Momentum Models for Adaptive Learning Rates](https://mlsanity.com/paper/2305.07583/) (2023)
- [Adaptive SGD with Polyak stepsize and Line-search: Robust Convergence and Variance Reduction](https://mlsanity.com/paper/2308.06058/) (2023)
- [Stochastic Mirror Descent: Convergence Analysis and Adaptive Variants via the Mirror Stochastic Polyak Stepsize](https://mlsanity.com/paper/2110.15412/) (2021)
- [Why Gradients Rapidly Increase Near the End of Training](https://mlsanity.com/paper/2506.02285/) (2025)
- [Cutting Some Slack for SGD with Adaptive Polyak Stepsizes](https://mlsanity.com/paper/2202.12328/) (2022)
- [AI-SARAH: Adaptive and Implicit Stochastic Recursive Gradient Methods](https://mlsanity.com/paper/2102.09700/) (2021)
- [A Stochastic Proximal Polyak Step Size](https://mlsanity.com/paper/2301.04935/) (2023)
- [Generalized Polyak Step Size for First Order Optimization with Momentum](https://mlsanity.com/paper/2305.12939/) (2023)
- [Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance](https://mlsanity.com/paper/2406.04142/) (2024)
- [An Adaptive Stochastic Gradient Method with Non-negative Gauss-Newton Stepsizes](https://mlsanity.com/paper/2407.04358/) (2024)
- [The Power of Factorial Powers: New Parameter settings for (Stochastic) Optimization](https://mlsanity.com/paper/2006.01244/) (2020)
- [Function Value Learning: Adaptive Learning Rates Based on the Polyak Stepsize and Function Splitting in ERM](https://mlsanity.com/paper/2307.14528/) (2023)

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

- [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)

## Abstract (arXiv)

The stochastic Polyak step size (SPS) has proven to be a promising choice for stochastic gradient descent (SGD), delivering competitive performance relative to state-of-the-art methods on smooth convex and non-convex optimization problems, including deep neural network training. However, extensions of this approach to non-smooth settings remain in their early stages, often relying on interpolation assumptions or requiring knowledge of the optimal solution. In this work, we propose a novel SPS variant, Safeguarded SPS (SPSsafe), for the stochastic subgradient method, and provide rigorous convergence guarantees for non-smooth convex optimization with no need for strong assumptions. We further incorporate momentum into the update rule, yielding equally tight theoretical results. Comprehensive experiments on convex benchmarks and deep neural networks corroborate our theory: the proposed step size achieves competitive performance to existing adaptive baselines and exhibits stable behavior across a wide range of problem settings. Finally, in the context of deep neural network training, the gradient norms under our step size do not collapse to (near) zero, indicating robustness to vanishing gradients.