Near-Optimal Non-Convex Stochastic Optimization under Generalized Smoothness

TL;DR

This paper introduces a new analysis of a simple variant of the STORM algorithm for generalized smooth non-convex stochastic optimization, achieving near-optimal high-probability and expected convergence guarantees with constant batch size.

Contribution

It provides the first near-optimal high-probability sample complexity for generalized smoothness and improves expected convergence bounds, all with constant batch size requirements.

Findings

Achieves $O( ext{log}(1/(\delta ext{,}\epsilon)) ext{ extasciicircum}3)$ high-probability sample complexity.

Recovers the optimal $O( ext{ extasciicircum}3)$ expected sample complexity.

Requires only a constant batch size, unlike previous methods.

Abstract

The generalized smooth condition, $(L_{0},L_{1})$-smoothness, has triggered people's interest since it is more realistic in many optimization problems shown by both empirical and theoretical evidence. Two recent works established the $O(\epsilon^{-3})$ sample complexity to obtain an $O(\epsilon)$-stationary point. However, both require a large batch size on the order of $\mathrm{ploy}(\epsilon^{-1})$, which is not only computationally burdensome but also unsuitable for streaming applications. Additionally, these existing convergence bounds are established only for the expected rate, which is inadequate as they do not supply a useful performance guarantee on a single run. In this work, we solve the prior two problems simultaneously by revisiting a simple variant of the STORM algorithm. Specifically, under the $(L_{0},L_{1})$-smoothness and affine-type noises, we establish the first near-optimal $O(\log(1/(\delta\epsilon))\epsilon^{-3})$ high-probability sample complexity where $\delta\in(0,1)$ is the failure probability. Besides, for the same algorithm, we also recover the optimal $O(\epsilon^{-3})$ sample complexity for the expected convergence with improved dependence on the problem-dependent parameter. More importantly, our convergence results only require a constant batch size in contrast to the previous works.