Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization: Robust Performance Without Small (Sub)Gradients
Dimitris Oikonomou, Nicolas Loizou
TL;DR
The paper addresses non-smooth stochastic optimization by introducing Safeguarded Stochastic Polyak Step Sizes (SPS_safe) for the Stochastic Subgradient Method and its momentum variant (IMA_SPS_safe). These adaptive rules use a safeguard M in the denominator to prevent instability from vanishing gradients and eliminate the need for interpolation or exact f_i^* values, while preserving an O(1/√T) convergence rate to a neighborhood of the optimum. The authors extend the analysis to momentum, prove convergence for both Cesàro averages and last iterates, and provide extensive numerical validation on convex problems and deep neural networks, showing faster convergence, reduced variance, and robustness to vanishing gradients. The work also connects safeguarded steps to adaptive gradient clipping and demonstrates practical benefits in DNN training, offering a parameter-free, theoretically grounded alternative to traditional Polyak-type updates in non-smooth settings.
Abstract
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 accelerates convergence, reduces variance, and consistently outperforms existing adaptive baselines. Finally, in the context of deep neural network training, our method demonstrates robust performance by addressing the vanishing gradient problem.