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

TL;DR

This paper introduces a new safeguarded stochastic Polyak step size for non-smooth optimization that guarantees convergence without strong assumptions and performs robustly in deep learning tasks.

Contribution

It proposes SPS$_{safe}$, a novel step size method with theoretical guarantees for non-smooth convex optimization, and demonstrates its effectiveness in deep neural network training.

Findings

Achieves competitive performance on deep neural networks.

Ensures convergence without requiring interpolation assumptions.

Maintains stable gradient norms, indicating robustness to vanishing gradients.

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. On non-smooth convex benchmarks, our experiments are consistent with the theoretical predictions on how the safeguard affects the convergence neighborhood. On deep neural networks the proposed step size achieves competitive performance to existing adaptive baselines and exhibits stable behavior across a wide range of problem settings. Moreover, in these experiments, the gradient norms under our step size do not collapse to (near) zero, indicating robustness to vanishing gradients.