Better Theory for SGD in the Nonconvex World

Ahmed Khaled; Peter Richtárik

Source-linked AI summary

Better Theory for SGD in the Nonconvex World

Ahmed Khaled, Peter Richtárik

arXiv:2002.03329v3Feb 9, 2020math.OCcs.LGstat.ML

TL;DR

The paper asks how SGD's stochastic-gradient second moment should be modeled for practical nonconvex optimization. It proposes expected smoothness, proves its generality and applicability to subsampling and compression, and obtains optimal rates for stationarity and global-solution problems under the stated conditions.

Problem
Existing assumptions on the stochastic-gradient second moment do not satisfactorily model practical nonconvex SGD settings such as subsampling and compression.
Method
The paper proposes expected smoothness (ES) for nonconvex problems and analyzes SGD across perturbation, subsampling, compression, sampling strategies, and minibatch sizes.
Results
O(ε^-4) is obtained for smooth nonconvex stationarity and O(ε^-1) under the PL condition, while ES is shown to be the weakest reviewed assumption.
Takeaways & Limitations
The framework provides an analysis of SGD that applies to practical subsampling and compression settings, including their combination, and yields sampling and minibatch recommendations.
Takeaways & Limitations
The general bounds can be too pessimistic, and specific sampling distributions are needed for tighter bounds.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies SGD for nonconvex optimization, where practical stochastic gradients arise from subsampling and compression but existing second-moment assumptions can be unrealistic. It proposes expected smoothness (ES) and derives optimal convergence rates while analyzing sampling and minibatch choices.

Problem: SGD is studied for finding ε-stationary points when global optimality is generally intractable for nonconvex objectives.The stochastic gradient is unbiased, and stationarity relaxes the requirement of global optimality.
Limitations of prior models: Existing second-moment assumptions may be too strong or unrealistic and fail even for subsampling from a sum of two functions.The paper motivates more accurate modeling of stochastic gradients in nonconvex optimization.
Sources of stochasticity: Practical finite-sum SGD uses minibatching, importance sampling, arbitrary sampling, gradient compression, or combinations of subsampling and compression.These settings produce stochastic gradients from sampled component gradients or randomized compressed gradients.
Expected smoothness: Expected smoothness (ES) is proposed to model the second moment of stochastic gradients for nonconvex problems.The assumption is introduced as a nonconvex extension inspired by expected-smoothness developments in convex analyses.
Expected smoothness: ES is the weakest and most general among the reviewed existing assumptions, including bounded variance, strong-growth, relaxed-growth, and gradient-confusion conditions.The assumption hierarchy is formalized in Theorem 1.
Convergence and sampling: O(ε^-4) is recovered for general smooth nonconvex problems, while O(ε^-1) is recovered under the PL condition, with importance-sampling probabilities and optimal minibatch formulas also derived.The analysis is designed to cover practical subsampling and compression settings, including their combination.

4 ES in the Nonconvex World

The paper extends expected smoothness to nonconvex SGD by replacing convex minimizer-based formulations with a lower-bound-based condition and an added third term. The resulting framework models practical stochasticity, including subsampling and compression, while retaining useful convergence guarantees.

The nonconvex formulation removes reliance on a global minimizer and instead uses a global lower bound on f, with an additional term needed for convergence guarantees.
Expected smoothness is presented as the weakest condition among prior assumptions for modelling stochastic-gradient second moments in nonconvex problems.
Subsampling: The stochastic reformulation supports user-defined sampling distributions, including broad subsampling strategies for finite-sum optimization.
Subsampling: The framework automatically satisfies expected smoothness under mild and natural assumptions on finite-sum components and sampling vectors.
Subsampling: Generic sampling bounds can be pessimistic, whereas specific practical distributions yield substantially tighter constants.
Compression: The framework covers stochastic gradients combining subsampling with unbiased gradient compression, a combination previously not covered for nonconvex SGD.

5 SGD in the Nonconvex World

The paper develops convergence guarantees for SGD under nonconvex, PL, and convex-QFG settings. It obtains optimal accuracy dependence for stationarity and global solution finding, while identifying condition-number differences between nonconvex and convex regimes.

General convergence theory: O(ε^-4) stochastic gradient evaluations attain an ε-stationary point, an optimal dependence without additional second-order or noise assumptions.
General convergence theory: 12δ0Lε^-2 iterations recover the standard nonconvex gradient-descent rate when B = 1 and A = C = 0.
Polyak-Łojasiewicz condition: 1/K convergence finds an optimal global solution under the Polyak-Łojasiewicz condition using the stepsize scheme of Lemma 3.
Polyak-Łojasiewicz condition: O(ε^-1) dependence on accuracy is recovered for strongly convex functions and is reported as optimal.
Convexity and quadratic functional growth: The nonconvex PL result can have multiplicative dependence κfκS, whereas whether the convex additive dependence reflects an inherent difference remains open.
Convexity and quadratic functional growth: Convexity with quadratic functional growth permits O(1/L) stepsizes and yields condition-number dependence κf + κS rather than κfκS.

6 Importance Sampling and Optimal Minibatch Size

The paper applies its complexity bounds to importance sampling and minibatch selection in finite-sum optimization. The analyzed distribution had prior use, while this work provides justification for it in nonconvex SGD.

Importance sampling is analyzed by choosing the sampling distribution to maximize convergence speed for minibatch size τ.
The sampling distribution in (23) is justified for SGD on nonconvex objectives for the first time in this work.
τ ≤ τ* = 1 + ... minimizes the stated expression over minibatch size τ.
Similar importance-sampling and minibatch-size expressions can be obtained for other sampling distributions.

7 Experiments

Experiments test importance sampling on regularized linear regression and compare expected smoothness with relaxed growth on logistic regression. The results support the proposed modeling approach in both settings.

Linear regression: Importance sampling is tested on nonconvex-regularized linear regression with n = 1000, d = 50, minibatches of size τ = 10, and K = 5000 iterations.The experiment samples minibatches with replacement and initializes x = 0.
Linear regression: Importance sampling outperforms uniform sampling when unnormalized Gaussian features produce Li proportional to ∥ai∥2.After normalizing the feature vectors, the two sampling methods are not expected to differ substantially, consistent with Figure 2.
Logistic regression: The logistic-regression experiment uses the a9a dataset with n = 32561, d = 123, τ = 1, and K = 500 iterations.It measures stochastic-gradient and loss quantities during SGD and fits constants for expected smoothness and relaxed growth.
Logistic regression: Expected smoothness achieves a smaller residual error than relaxed growth by using function values to model stochastic gradients.The first fitted coefficient is close to the theoretically estimated Lmax, while coefficient C is difficult to estimate.

8 Basic Facts and Notation

This section introduces probability and linear-algebra facts, smoothness properties, and the one-step gradient-descent mapping used later in the analysis.

Probability: The section collects probability-theory facts involving a random variable and a constant vector.
Linear algebra: It records linear-algebra inequalities for vectors in Rd and uses convexity of the squared norm with Jensen’s inequality.
Smoothness and descent: For an L-smooth function, the section states a smoothness relation and defines the one-step gradient-descent mapping for stepsize γ > 0.

9 Relations Between Assumptions

The paper establishes implication relations among stochastic-gradient assumptions and gives an example where relaxed growth fails while expected smoothness holds.

Counterexample: A simple finite-sum problem shows that relaxed growth need not hold, even though expected smoothness can hold with A = 1/2, B = 0, and C = 2.
Assumption relations: The maximal strong growth condition implies the strong growth condition, which implies relaxed growth.
Assumption relations: Bounded stochastic-gradient variance and gradient confusion for finite-sum problems each imply relaxed growth.
Assumption relations: Relaxed growth implies expected smoothness, while sure smoothness implies expected smoothness under the paper’s expectation structure.
Assumption relations: The theorem also derives relaxed-growth parameters α = n and β = η(n −1) in one implication case.

10 Proofs for Section 4.5 and 4.6

The proofs derive expected-smoothness bounds using smoothness, nonnegativity, expectation, and sampling identities. They cover independent and without-replacement sampling schemes and establish the stated assumption result.

Proof ingredients: The proofs use the definition of stochastic gradients, linearity of expectation, convexity of the squared norm, and lower bounds on component functions.
Sampling schemes: Sampling calculations treat independent sampling with replacement, independent sampling without replacement, and τ-nice sampling without replacement.
Sampling schemes: For independent sampling with replacement, indicator variables encode dice-roll outcomes and yield the required second-moment identities.
Sampling schemes: Without-replacement analyses use pairwise sampling probabilities, including distinct-index and equal-index cases, before substituting into the moment calculation.
Expected-smoothness proof: The final proof decomposes conditional variance, applies Assumption 2 and smoothness bounds, and concludes that Assumption 2 is satisfied.

11 Proofs for General Smooth Objectives

The proofs derive convergence bounds for general smooth objectives by combining smoothness, Assumption 2, weighted inequalities, and stepsize conditions. The resulting bound is obtained after summing the weighted recursion and collecting requirements.

General smooth objectives: Smoothness and Assumption 2 are combined through conditional expectations to bound the one-step progress of SGD.The proof then applies the tower property and rearranges the resulting inequality.
General smooth objectives: A recursively defined weighting sequence is used as a proof technique to weight the inequalities before summation.The weights do not appear in the final bounds.
General smooth objectives: The weighted inequalities are summed over iterations and rearranged to obtain the lemma’s statement.The proof explicitly concludes this step after summation.
General smooth objectives: The final convergence guarantee requires satisfying multiple stepsize constraints and combining the resulting iteration bounds.The proof separately controls terms relative to ε2 before collecting them into a single bound.

12 Proofs Under Assumption 5

Under Assumption 5, the proofs establish convergence by recursively controlling the expected error and telescoping weighted inequalities. The argument concludes by taking the maximum of two derived bounds.

Proof under Assumption 5: A threshold k0 separates the initial and later iterations, with K−k0≥K/2 used to control the latter block.This estimate is substituted into the subsequent bound.
Proof under Assumption 5: The proof starts from a recursion that contracts the error by a factor involving aγ while adding a γ2c error term.The recursion is then expanded over iterations.
Proof under Assumption 5: For later iterations, weighted inequalities are summed and telescoped to derive a bound on the accumulated error.The proof divides by the terminal weight after establishing the telescoping relation.
Proof under Assumption 5: The final result is obtained by taking the maximum of the two bounds derived in the proof.This is the stated closing step.
Proof under Assumption 5: The theorem’s constants are recovered by applying the preceding lemma with specified values of a, b, and c.The proof then uses the definition of κS to match the theorem statement.

13 Proofs Under Assumption 6

Under Assumption 6, the proofs relate projected-iterate distance, gradient norms, and function gaps using smoothness, convexity, and projection properties. A separate consequence converts distance guarantees into function-value guarantees.

Proof under Assumption 6: For the gradient step Tγ(x), smoothness yields a bound combining projected distance, squared gradient norm, and the function gap.The bound applies for γ≤1/L.
Proof under Assumption 6: The proof obtains this inequality by applying smoothness and convexity before substituting Tγ(x)−x=−γ∇f(x).The projection point is set to π(x) in the final step.
Proof under Assumption 6: Assumption 2 bounds the function-decrease term using the function gap and a constant C, producing a coefficient 2ρ.The proof defines ρ=A+BL.
Proof under Assumption 6: The resulting convergence guarantee is first expressed through the distance of iterates from the set of optima.The proof takes unconditional expectations after applying the projection properties.
Proof under Assumption 6: Quadratic Functional Growth supplies consequences that convert the distance-to-optima guarantee into convergence in function values.The passage explicitly identifies this as the next conversion step.

14 Experimental Details

The experiments fit Assumption 2’s parameters from SGD snapshots and compare its modeling residual with relaxed growth assumptions. Logistic-regression experiments track gradients and loss over iterations.

Experimental setup: SGD is run for 500 iterations with minibatch size 1 and uniform sampling on the a9a dataset.The dataset has n=32561 and d=123, with λ=1/2.
Data collection: Snapshots collected during SGD contain the squared full-gradient norm, full-data loss, and average stochastic-gradient norm.Snapshots are collected every five iterates.
Results: In regularized logistic regression, full gradient norms approach zero while stochastic-gradient norms and loss quickly plateau around 0.5.The passage connects the plateau pattern to Lemma 2 under an appropriately chosen stepsize.
Hyperparameter estimation: The individual Lipschitz constants are estimated from the data, while global and variance-related parameters are estimated using prescribed sampling formulas.The stepsize is then selected from the estimated constants.
Data fitting: Assumption 2 parameters A, B, and C are fitted by minimizing squared residual error against average stochastic-gradient norms.The design matrix uses function gaps, squared full-gradient norms, and a constant feature.
Results: The fitted expected-smoothness model achieves a smaller residual error than relaxed growth by using function values to model stochastic gradients.The first coefficient is close to the estimated Lmax, whereas C is difficult to estimate.