Better Theory for SGD in the Nonconvex World
This work provides improved theoretical foundations for SGD in nonconvex settings, which is incremental but important for practitioners in machine learning dealing with large-scale optimization problems.
The paper tackles the analysis of Stochastic Gradient Descent (SGD) for nonconvex optimization by introducing a new variant of the expected smoothness assumption, showing it is more general and reasonable than prior work, and achieving optimal rates of O(ε^{-4}) for stationary points and O(ε^{-1}) under Polyak-Łojasiewicz conditions.
Large-scale nonconvex optimization problems are ubiquitous in modern machine learning, and among practitioners interested in solving them, Stochastic Gradient Descent (SGD) reigns supreme. We revisit the analysis of SGD in the nonconvex setting and propose a new variant of the recently introduced expected smoothness assumption which governs the behaviour of the second moment of the stochastic gradient. We show that our assumption is both more general and more reasonable than assumptions made in all prior work. Moreover, our results yield the optimal $\mathcal{O}(\varepsilon^{-4})$ rate for finding a stationary point of nonconvex smooth functions, and recover the optimal $\mathcal{O}(\varepsilon^{-1})$ rate for finding a global solution if the Polyak-Łojasiewicz condition is satisfied. We compare against convergence rates under convexity and prove a theorem on the convergence of SGD under Quadratic Functional Growth and convexity, which might be of independent interest. Moreover, we perform our analysis in a framework which allows for a detailed study of the effects of a wide array of sampling strategies and minibatch sizes for finite-sum optimization problems. We corroborate our theoretical results with experiments on real and synthetic data.