Technical Summary: Safeguarded Stochastic Polyak Step Sizes for Non-smooth Optimization

1. Problem Statement

The paper addresses the challenge of optimizing unconstrained finite-sum problems of the form:
$$\min_{x \in \mathbb{R}^d} f(x) = \frac{1}{n} \sum_{i=1}^n f_i(x)$$
where the component functions $f_i$ are convex, Lipschitz, and non-smooth. The goal is to minimize the empirical risk using stochastic methods without relying on strong assumptions such as:

1. Interpolation: The assumption that there exists a solution $x^*$ where $f_i(x^*) = f_i^*$ for all $i$.
2. Oracle Knowledge: The requirement to know the optimal loss values $f_i(x^*)$ or $f^*$ a priori.
3. Step-size Capping: The reliance on manually tuned upper bounds ($\gamma_b$) that often force the algorithm to behave like a fixed learning rate.

Existing Stochastic Polyak Step (SPS) methods either require these strong assumptions or suffer from instability when gradients vanish (common in deep neural networks), leading to exploding step sizes.

2. Methodology

The authors propose Safeguarded Stochastic Polyak Step Sizes (SPSsafe) for two primary algorithms: the Stochastic Subgradient Method (SSM) and the Iterate Moving Average (IMA), which is algebraically equivalent to Stochastic Heavy Ball momentum.

The Core Mechanism

The proposed step size modifies the standard Polyak rule by introducing a safeguard constant $M > 0$ in the denominator. Instead of capping the step size from above (as in previous $SPS_{max}$ approaches), SPSsafe caps the denominator from below to prevent division by near-zero gradients.

For Stochastic Subgradient Method (SSM):
$$\gamma_t = \frac{f_i(x_t) - \ell_i^*}{\max\{\|g_t^i\|^2, M\}}$$
where $g_t^i \in \partial f_i(x_t)$ is the stochastic subgradient, and $\ell_i^*$ is a known lower bound for $f_i$ (e.g., 0 for non-negative losses).

For Iterate Moving Average (IMA) with Momentum:
$$\eta_t = \frac{[f_i(x_t) - \ell_i^* + \lambda_t \langle g_t^i, x_t - x_{t-1} \rangle]_+}{\max\{\|g_t^i\|^2, M\}}$$
where $\lambda_t \geq 0$ is the momentum parameter and $[\cdot]_+$ denotes the positive part.

Theoretical Connection to Clipping

The paper establishes a theoretical link between SPSsafe and Gradient Clipping. Specifically, Proposition 2.1 demonstrates that SSM with SPSsafe (where $M=c^2$) is algebraically equivalent to a Clipped SSM with an adaptive step size. This provides the first theoretical convergence guarantees for Polyak-style updates combined with gradient clipping in the non-smooth stochastic regime.

3. Key Contributions

1. Novel Step Size Rule (SPSsafe):

   • Introduces a Polyak-type step size that requires only a lower bound $\ell_i^*$ and a safeguard parameter $M$.
   • Eliminates the need for interpolation assumptions and oracle knowledge of $f_i(x^*)$.
   • Prevents the "vanishing gradient" issue where $\|g_t^i\| \to 0$ causes the step size to explode.

2. Rigorous Convergence Guarantees:

   • SSM: Proves an $O(T^{-1/2})$ convergence rate to a neighborhood of the solution for convex, Lipschitz objectives. The neighborhood size depends on the variance $\sigma^2$ and the safeguard $M$.
   • IMA (Momentum): Extends these guarantees to the momentum setting, proving $O(T^{-1/2})$ convergence for both the Cesàro average and the last iterate, without requiring knowledge of $f_i(x^*)$.

3. Theoretical Clarification on Smoothing:

   • The authors argue that previous "smoothed" Polyak methods (e.g., $SPS_{max}$ with a decaying upper bound $\gamma_b$) often degenerate into fixed learning rates in practice because the bound is selected in the majority of iterations. SPSsafe avoids this by keeping the numerator intact and only bounding the denominator, ensuring the step size remains adaptive.

4. Empirical Validation:

   • Demonstrates that SPSsafe achieves competitive performance against adaptive baselines (like Adam) and existing Polyak variants on convex benchmarks (SVM, Phase Retrieval).
   • Shows that in Deep Neural Network (DNN) training (ResNet on CIFAR), SPSsafe maintains larger gradient norms compared to smoothed baselines, preventing the collapse of gradients to near-zero while maintaining generalization performance.

4. Results and Performance

• Convergence Rates: Theoretical analysis confirms an $O(T^{-1/2})$ rate for both SSM and IMA variants. The error bound consists of a standard optimization term and a neighborhood term proportional to $\sigma^2/M$.
• Sensitivity Analysis: Experiments on synthetic non-smooth problems show that increasing the safeguard $M$ reduces the limiting error neighborhood, confirming the theoretical relationship.
• DNN Training:
  • On CIFAR-10/100 with ResNet-20/32, SPSsafe achieves test accuracy comparable to or better than tuned SGD, Adam, and smoothed SPS baselines.
  • Gradient Norms: Unlike smoothed Polyak methods which drive gradient norms to near-zero (potentially causing training stagnation), SPSsafe maintains robust gradient norms throughout training, indicating resilience to vanishing gradients.

5. Significance and Claims

The paper claims to provide one of the first Polyak-type rules with rigorous convergence guarantees for stochastic, convex, non-smooth optimization that:

1. Does not assume interpolation.
2. Does not require knowledge of optimal loss values ($f_i(x^*)$).
3. Does not rely on arbitrary upper bounds that force the algorithm into fixed-step behavior.

The authors emphasize that while their theory is derived for convex Lipschitz objectives, the empirical results on highly non-convex deep learning models suggest the method is robust and effective in practice. They position SPSsafe as a stable, adaptive alternative that bridges the gap between the simplicity of Polyak steps and the robustness required for modern non-smooth optimization tasks.

Limitations and Future Work:
The authors acknowledge that their theoretical guarantees are currently limited to convex Lipschitz settings. They identify extending the analysis to structured non-convex settings (e.g., weakly convex, Polyak-Lojasiewicz) and $(L_0, L_1)$-smooth regimes as necessary future directions to fully bridge the gap between theory and modern non-convex training. They also suggest exploring applications in distributed and federated learning.