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

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 (SPS$_{safe}$), 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.

Training a machine learning model means repeatedly nudging it to reduce its mistakes, and a crucial choice is how big each nudge should be: too big and training destabilizes, too small and it crawls. Setting this "step size" by hand is tedious, so researchers built rules that set it automatically. One elegant family, Polyak step sizes, picks the step from how wrong the model is, but these rules struggle on "rough" problems full of sharp corners. Existing versions either assume the model can fit every example perfectly, or rely on a manual cap that quietly turns the rule into a fixed step. These versions can blow up when the gradients becomes tiny. We introduce a "safeguarded" version that, rather than capping the step from above, places a floor under the quantity the rule divides by. This single safeguard prevents the blow-up, and never collapses into a fixed step; we prove it works reliably, including with momentum. The result is a simple learning-rate rule that is robust, nearly tuning-free and matches or beats leading optimizers when training image-recognition networks, with the bonus that gradients never collapse toward zero.