[Paper Review] Better Theory for SGD in the Nonconvex World

Ahmed Khaled, Peter Richtárik|arXiv (Cornell University)|Feb 9, 2020

Stochastic Gradient Optimization Techniques46 references60 citations

TL;DR

This paper introduces a general expected smoothness (ES) condition to model stochastic gradients for nonconvex SGD, proving optimal convergence rates and flexibility across subsampling and compression schemes.

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

Motivation & Objective

• Motivate accurate modeling of stochastic gradient second moments in nonconvex SGD.
• Introduce the expected smoothness (ES) assumption as the weakest general condition for SGD analysis.
• Show that ES holds under practical subsampling and compression schemes.
• Derive optimal convergence rates: O(ε^-4) for general nonconvex and O(ε^-1) under the Polyak-Lojasiewicz condition.
• Provide guidance on minibatch sizing and importance sampling.
• Compare ES-based results with existing convex/nonconvex convergence results.

Proposed method

• Define stochastic gradient g(x) as an unbiased estimator of ∇f(x) under a finite-sum framework.
• Propose the ES inequality: E[||g(x)||^2] ≤ 2A(f(x)-f_inf) + B||∇f(x)||^2 + C.
• Show ES is the weakest among known assumptions and automatically holds for many subsampling/compression schemes.
• Develop convergence analysis for SGD under ES, obtaining ε-stationary point with O(ε^-4) gradient evaluations.
• Extend analysis to the Polyak-Lojasiewicz (PL) condition to obtain O(ε^-1) rates and discuss convergence under QFG and convexity.
• Examine sampling schemes (with/without replacement, τ-nice) and compression operators (ω-compressors) to instantiate ES.

Experimental results

Research questions

• RQ1 What minimal condition on the second moment of stochastic gradients ensures reliable nonconvex SGD convergence?
• RQ2 Is there a general, weak assumption that covers practical stochasticity sources like subsampling and compression?
• RQ3 What are the convergence rates for SGD under ES, and do they match known optimal rates under PL or convex cases?
• RQ4 How do sampling strategies and minibatch sizes influence ES parameters and resulting complexity?
• RQ5 Can ES recover or improve upon existing results for nonconvex optimization, including full gradient descent and PL-based convergence?

Key findings

• Introduces ES as the weakest known condition governing the second moment of stochastic gradients for nonconvex problems.
• Demonstrates ES automatically holds under common practical settings such as subsampling and compression.
• Achieves the optimal O(ε^-4) convergence rate to an ε-stationary point for general smooth nonconvex functions.
• Under the Polyak-Lojasiewicz condition, recovers the O(ε^-1) rate for finding a global solution.
• Provides a framework to derive optimal minibatch sizes and importance sampling probabilities within ES.
• Relates ES to and subsumes several existing gradient-noise models (BV, M-SG, SG, RG, GC) while remaining more general.