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 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}({\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.

arXiv comment33 pages, 3 figures, 4 theorems, and 4 propositions. V3 updates: added several references on error conditions (Tseng, Solodov, Bottou and Tsitsiklis, Grimmer), added a full proof of Corollary 1, cleaned up several proofs, and made minor adjustments to text for clarity

References in corpus (13)

• QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding
• Optimal Distributed Online Prediction using Mini-Batches
• Making Gradient Descent Optimal for Strongly Convex Stochastic Optimization
• Stochastic Gradient Descent for Non-smooth Optimization: Convergence Results and Optimal Averaging Schemes
• Optimization for deep learning: theory and algorithms
• Stochastic Distributed Learning with Gradient Quantization and Variance Reduction
• The Error-Feedback Framework: Better Rates for SGD with Delayed Gradients and Compressed Communication
• Unified Optimal Analysis of the (Stochastic) Gradient Method
• Distributed learning with compressed gradients
• The Impact of Neural Network Overparameterization on Gradient Confusion and Stochastic Gradient Descent
• Sharp Analysis for Nonconvex SGD Escaping from Saddle Points
• Tight Dimension Independent Lower Bound on the Expected Convergence Rate for Diminishing Step Sizes in SGD
• The Complexity of Finding Stationary Points with Stochastic Gradient Descent

Cited by in corpus (27)

• Random Reshuffling: Simple Analysis with Vast Improvements
• On the Sample Complexity of Quantum Boltzmann Machine Learning
• A Unified Analysis of Stochastic Gradient Methods for Nonconvex Federated Optimization
• PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization
• A Second look at Exponential and Cosine Step Sizes: Simplicity, Adaptivity, and Performance
• Stochastic Optimization with Heavy-Tailed Noise via Accelerated Gradient Clipping
• ZeroSARAH: Efficient Nonconvex Finite-Sum Optimization with Zero Full Gradient Computation
• A general sample complexity analysis of vanilla policy gradient
• Global Convergence and Stability of Stochastic Gradient Descent
• Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ball
• SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation
• Stochastic Gradient Descent on Nonconvex Functions with General Noise Models
• Convergence of stochastic gradient descent schemes for Lojasiewicz-landscapes
• Proximal and Federated Random Reshuffling
• UAdam: Unified Adam-Type Algorithmic Framework for Non-Convex Stochastic Optimization
• Permutation Compressors for Provably Faster Distributed Nonconvex Optimization
• Stochastic Gradient Descent-Ascent and Consensus Optimization for Smooth Games: Convergence Analysis under Expected Co-coercivity
• Stochastic Extragradient: General Analysis and Improved Rates
• Spatio-Temporal Neural Network for Fitting and Forecasting COVID-19
• Accelerated Almost-Sure Convergence Rates for Nonconvex Stochastic Gradient Descent using Stochastic Learning Rates
• FLIX: A Simple and Communication-Efficient Alternative to Local Methods in Federated Learning
• Online Stochastic Gradient Methods Under Sub-Weibull Noise and the Polyak-Łojasiewicz Condition
• Online Stochastic Optimization for Unknown Linear Systems: Data-Driven Synthesis and Controller Analysis
• A Stochastic Objective-Function-Free Adaptive Regularization Method with Optimal Complexity
• Random-reshuffled SARAH does not need a full gradient computations
• Det-CGD: Compressed Gradient Descent with Matrix Stepsizes for Non-Convex Optimization
• Stochastic Polyak Stepsize with a Moving Target