Lower Bounds and Proximally Anchored SGD for Non-Convex Minimization Under Unbounded Variance

TL;DR

This paper establishes lower bounds for stochastic non-convex optimization under unbounded variance conditions and introduces PASTA, an algorithm that achieves optimal complexity matching these bounds across various non-convex regimes.

Contribution

It provides the first information-theoretic lower bounds under BG-0 variance growth and proposes PASTA, a unified algorithm that attains these bounds in diverse non-convex settings.

Findings

Lower bounds of $oldsymbol{ ext{Ω}(oldsymbol{ extepsilon}^{-6})}$ and $oldsymbol{ ext{Ω}(oldsymbol{ extepsilon}^{-4})}$ oracle queries established.

PASTA algorithm achieves minimax optimal complexities matching the lower bounds.

Results hold for unbounded domains and stochastic gradients in multiple non-convex regimes.

Abstract

Analysis of Stochastic Gradient Descent (SGD) and its variants typically relies on the assumption of uniformly bounded variance, a condition that frequently fails in practical non-convex settings, such as neural network training, as well as in several elementary optimization settings. While several relaxations are explored in the literature, the Blum-Gladyshev (BG-0) condition, which permits the variance to grow quadratically with distance has recently been shown to be the weakest condition. However, the study of the oracle complexity of stochastic first-order non-convex optimization under BG-0 has remained underexplored. In this paper, we address this gap and establish information-theoretic lower bounds, proving that finding an $\epsilon$-stationary point requires $\Omega(\epsilon^{-6})$ stochastic BG-0 oracle queries for smooth functions and $\Omega(\epsilon^{-4})$ queries under mean-square smoothness. These limits demonstrate an unavoidable degradation from classical bounded-variance complexities, i.e., $\Omega(\epsilon^{-4})$ and $\Omega(\epsilon^{-3})$ for smooth and mean-square smooth cases, respectively. To match these lower bounds, we consider Proximally Anchored STochastic Approximation (PASTA), a unified algorithmic framework that couples Halpern anchoring with Tikhonov regularization to dynamically mitigate the extra variance explosion term permitted by the BG-0 oracle. We prove that PASTA achieves minimax optimal complexities across numerous non-convex regimes, including standard smooth, mean-square smooth, weakly convex, star-convex, and Polyak-Lojasiewicz functions, entirely under an unbounded domain and unbounded stochastic gradients.