Step-Size Stability in Stochastic Optimization: A Theoretical Perspective

Title: Step-Size Stability in Stochastic Optimization: A Theoretical Perspective

Problem Statement
Training machine learning models requires optimization algorithms that are both efficient and robust. While the learning rate (step size) is widely recognized as the most critical hyperparameter to tune, existing optimization theory primarily focuses on convergence rates for specific step sizes, often neglecting the sensitivity of methods to step-size magnitude. Empirical observations suggest that adaptive methods like the Stochastic Polyak Step Size (SPS) and Non-negative Gauss-Newton (NGN) steps are more robust to large step sizes than standard Stochastic Gradient Descent (SGD), but a theoretical framework explaining why and how this stability occurs has been lacking. This paper aims to fill that gap by providing a theoretical analysis of step-size sensitivity.

Methodology and Theoretical Framework
The authors analyze a general family of model-based stochastic optimization methods. The core update rule is defined as:
$$x_{t+1} = \arg \min_{y \in \mathbb{R}^d} f_{x_t}(y, s_t) + \frac{1}{2\alpha_t} \|y - x_t\|^2$$
where $f_{x_t}(\cdot, s)$ is a model of the objective function $f(\cdot, s)$ around the current iterate $x_t$, and $\alpha_t$ is the step size.

The central contribution of the methodology is the identification of a key quantity called the stability index, denoted as $\delta_t$:
$$\delta_t := f(x_t, s_t) - f_{x_t}(x_{t+1}, s_t) - \frac{1}{2\alpha_t} \|x_{t+1} - x_t\|^2$$

The authors establish that the scaling behavior of the expected stability index, $\mathbb{E}[\delta_t]$, with respect to the step size $\alpha$ determines the method's robustness.

• Convex Case: They derive suboptimality bounds for the average and last iterates. These bounds take the form $\Omega_T = \frac{D^2}{2\sum \alpha_t} + \frac{\sum \alpha_t \mathbb{E}[\delta_t]}{\sum \alpha_t}$. The first term represents bias (decreasing with larger steps), while the second term represents variance/stability (increasing with larger steps). The behavior of the second term dictates the method's performance as $\alpha$ grows.
• Weakly Convex Case: The analysis is extended to $\rho$-weakly convex functions, where the stability index similarly governs the convergence rate of the Moreau envelope gradient norm.

Key Contributions

1. Identification of the Stability Index: The paper identifies $\delta_t$ as the critical quantity describing how performance degrades as the step size becomes too large.
2. Theoretical Proof of Robustness: The authors derive explicit expressions for $\delta_t$ for four specific methods:
   • SGD: Uses a linear model. $\delta_t^{SGD} = \frac{\alpha_t}{2} \|g_t\|^2$. This scales linearly with $\alpha_t$, leading to instability at large step sizes.
   • SPS: Uses a truncated linear model. $\delta_t^{SPS} \leq \min\{\alpha_t \|g_t\|^2, f(x_t, s_t) - C_t\}$. The growth is capped, making it less sensitive to large $\alpha$.
   • NGN: Uses a square-root model for non-negative losses. $\delta_t^{NGN} = \frac{\gamma_t}{2} \|g_t\|^2$, where the effective step size $\gamma_t$ saturates as $\alpha_t \to \infty$. This results in sub-linear scaling of the stability index.
   • SPP (Stochastic Proximal Point): Uses the exact model. The stability index is bounded by the sub-optimality and interpolation constants, also exhibiting sub-linear scaling in the convex case.
3. Quantification of Gains: The analysis reveals that the stability advantage of SPS and NGN over SGD is determined by the interpolation constant (how well the model fits the data) and the lower-bound estimation error (for SPS).
4. Non-Convex Empirical Validation: While the theoretical bounds are derived for convex and weakly convex settings, the authors demonstrate experimentally that the theoretical bound qualitatively mirrors actual performance even in non-convex deep learning tasks.

Experimental Results
The authors validate their theory through experiments on both convex and non-convex problems:

• Non-Convex (Deep Learning): Training a ResNet20 on CIFAR10. The theoretical bound $\Omega_T$ closely predicts the range of base learning rates ($\alpha$) that yield good training loss. SGD degrades rapidly for $\alpha \geq 1$, whereas SPS and NGN remain stable up to $\alpha \approx 10$. The bound accurately reflects the improved stability of adaptive methods and the effect of warmup schedules on SGD.
• Convex (Regression/Classification): Experiments on linear regression and logistic regression confirm that the theoretical bound matches empirical performance. Variations in problem settings (e.g., adding noise, changing batch sizes, or using incorrect lower bounds for SPS) align with theoretical predictions regarding the interpolation constant and lower-bound estimation error.
• Stability Index Evolution: Tracking $\delta_t$ over time shows that for SGD, the index explodes with large $\alpha$, while for SPS and NGN, it remains bounded or grows much slower, confirming the theoretical inequality $\delta_t^{adaptive} \leq \delta_t^{SGD}$.

Significance and Claims
The paper claims to provide the first theoretical evidence that adaptive step-size methods (SPS, NGN, SPP) are provably more robust to large step sizes than SGD in convex and weakly convex settings. This robustness is quantified by the stability index $\delta_t$, which scales sub-linearly for adaptive methods versus linearly for SGD.

The authors emphasize that their work bridges the gap between empirical observations of learning-rate sensitivity and theoretical convergence analysis. They note that while their theory is restricted to model-based methods and does not currently cover momentum or preconditioning (e.g., Adam), the framework successfully explains the superior stability of SPS and NGN. The paper concludes that the theoretical bound serves as a reliable qualitative predictor for the range of effective learning rates, even in non-convex scenarios where the strict theoretical assumptions do not hold.

Limitations
The authors acknowledge several limitations:

• The analysis is restricted to model-based methods and does not cover momentum, weight decay, or preconditioning techniques prevalent in practice.
• The theoretical bounds are derived for convex and weakly convex problems; the applicability to general non-convex problems is currently supported only by empirical observation, not proof.
• The paper focuses on explaining observed phenomena rather than providing immediate practical heuristics for exploiting the stability index in training loops.