Better Theory for SGD in the Nonconvex World
Abstract: 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 \emph{expected smoothness} assumption which governs the behavior 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}(\epsilon^{-4})$ rate for finding a stationary point of nonconvex smooth functions, and recover the optimal $\mathcal{O}(\epsilon^{-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.
Certifications: Survey Certification
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: 1. Added some notes on a comparison one of the reviewers asked for.
2. Corrected a few typos.
Assigned Action Editor: ~Raman_Arora1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 456
Loading