Source-linked AI summary
A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent
Eduard Gorbunov, Filip Hanzely, Peter Richtárik
TL;DR
Existing SGD variants use different stochastic-gradient constructions and assumptions, making their relationships and convergence analyses difficult to unify. This paper proposes a framework covering these variants and proves a linear convergence theorem whose special cases recover the best known rates.
Problem
SGD variants offer different design choices but require dedicated analyses, and no comprehensive theory explains their relationships or guides the development of new variants.
Method
The paper introduces a unified theoretical framework covering vanilla, variance-reduced, randomized-coordinate, and quantized SGD methods as special cases.
Results
The framework's general theorem establishes a linear rate for a wide range of proximal SGD methods, with best known rates in all known special cases.
Takeaways & Limitations
The framework connects previously separate SGD variants under one analysis and provides pre-proved convergence guarantees for their special cases.
Takeaways & Limitations
Comparable guarantees for the nonconvex case remain a major open problem.
Abstract
from arXiv · showhide
In this paper we introduce a unified analysis of a large family of variants of proximal stochastic gradient descent ({\tt SGD}) which so far have required different intuitions, convergence analyses, have different applications, and which have been developed separately in various communities. We show that our framework includes methods with and without the following tricks, and their combinations: variance reduction, importance sampling, mini-batch sampling, quantization, and coordinate sub-sampling. As a by-product, we obtain the first unified theory of {\tt SGD} and randomized coordinate descent ({\tt RCD}) methods, the first unified theory of variance reduced and non-variance-reduced {\tt SGD} methods, and the first unified theory of quantized and non-quantized methods. A key to our approach is a parametric assumption on the iterates and stochastic gradients. In a single theorem we establish a linear convergence result under this assumption and strong-quasi convexity of the loss function. Whenever we recover an existing method as a special case, our theorem gives the best known complexity result. Our approach can be used to motivate the development of new useful methods, and offers pre-proved convergence guarantees. To illustrate the strength of our approach, we develop five new variants of {\tt SGD}, and through numerical experiments demonstrate some of their properties.
1 Introduction
The paper studies proximal stochastic optimization when exact gradients are expensive but unbiased stochastic estimates are efficient, including finite-sum and distributed settings.
- The objective combines a smooth convex function f with a proximable proper closed convex regularizer R.
- When exact gradients are unavailable or costly, unbiased stochastic gradients can be computed efficiently from sampled data.
- In finite-sum problems, f averages n smooth functions, and uniformly sampled component gradients are unbiased estimators of the full gradient.
- The full gradient is typically about n times more expensive to compute than one component gradient.
- Distributed formulations assign component losses to machines, with each machine potentially storing a finite sum of training examples.
2 The Many Faces of Stochastic Gradient Descent
SGD variants differ in how they construct stochastic gradients, using sampling, variance reduction, quantization, distributed aggregation, or coordinate updates. These choices create both design flexibility and separate analytical challenges.
- Proximal SGD updates the iterate with a stochastic gradient step followed by the proximal operator of R.
- There are infinitely many ways to construct unbiased stochastic gradients, targeting convergence speed, iteration cost, parallelizability, and generalization.
- Vanilla SGD: Importance sampling and mini-batching can improve linear rates through larger stepsizes and modify the variance parameter σ2, but generally do not eliminate gradient noise.
- Variance reduced SGD: Variance-reduced methods can learn stochastic gradients at the optimum and eliminate the adverse effect of σ2 > 0 in finite-sum problems.
- Distributed SGD, quantization and variance reduction: Quantized distributed methods apply randomized transformations to gradient updates, preserving unbiasedness while reducing communication bits, and DIANA and VR-DIANA provide analyzed examples.
- Randomized coordinate descent: RCD methods use coordinate-specific stepsizes, historically struggled with non-separable regularizers, and are naturally variance-reduced in the discussed setting.
3 Contributions
The paper addresses the lack of a comprehensive theory connecting diverse SGD variants by proposing a unified framework, sharp special-case rates, generalizations, new methods, and experiments.
- Existing SGD variants form largely disconnected communities with different intuitions, applications, and convergence analyses.
- The framework unifies vanilla, variance-reduced, randomized coordinate descent, and quantized SGD methods within one analysis.
- The framework also generalizes known methods to proximal and arbitrary-sampling settings, including L-SVRG in the proximal case.
- Theorem 4.1 yields the best known rates for all known special cases analyzed in the framework.
- Verifying the framework's stochastic-gradient assumption supplies complexity estimates for an array of new SGD variants.
- Experiments examine practical properties of selected new and generalized methods against appropriately selected existing methods.
4 Main Result
The analysis uses a parametric assumption on stochastic gradients and iterates together with strong quasi-convexity. Its main theorem establishes linear convergence, with the limiting behavior determined by additive noise terms.
- The framework first specifies a stochastic-gradient assumption, assumptions on f, and then a unified convergence theorem.
- Key assumption: Assumption 4.1 requires unbiased stochastic gradients and two parametric recursions governing iterates and stochastic gradients.
- Key assumption: The coupled inequalities in Assumption 4.1 are intended to capture vanilla, variance-reduced, arbitrary-sampling, quantized, and coordinate-descent variants.
- Function assumption: The paper assumes f is µ-strongly quasi-convex, a generalization of µ-strong convexity.
- Main theorem: Theorem 4.1 establishes a linear rate for proximal SGD under Assumptions 4.1 and 4.2, using a Lyapunov function.
- Main theorem: Vanilla methods with D1 > 0 or D2 > 0 converge only to an oscillation region, whereas variance-reduced methods with D1 = D2 = 0 can reach x∗.
5 The Classic, The Recent and The Brand New
This section shows how the framework encompasses existing and new SGD variants, including arbitrary sampling, quantization, variance reduction, and coordinate descent. It also summarizes the framework’s parameterization, convergence guarantees, and five newly developed methods.
- The framework includes vanilla, variance-reduced, arbitrary-sampling, quantized, and randomized-coordinate-descent methods.
- Known methods recovered as special cases retain their best known convergence rates.
- Assumption 4.1 parameters characterize the methods’ convergence behavior, with smoothness affecting A and quantization variance affecting B.
- The framework develops five new SGD variants and supplies complexity estimates whenever Assumption 4.1 is verified.
- SGD-SR is analyzed under arbitrary sampling, while Q-SGD-SR extends SGD-SR with quantization and generalizes QSGD.
- SGD-star bridges vanilla and variance-reduced SGD, whereas N-SAGA and N-SEGA use noisy estimators and therefore converge only to oscillation regions.
6 Experiments
The experiments compare SGD-MB with independent SGD on LIBSVM under alternative sampling choices. They report nearly identical iteration complexity, while SGD-MB has cheaper iterations.
- SGD-MB uses with-replacement mini-batches formed by τ independent draws and averages the corresponding stochastic gradients.
- The experiment reports SGD-MB and independent SGD with uniform or importance-sampling probabilities and different replacement settings.
- SGD-MB and independent SGD have almost identical iteration complexity, while SGD-MB has cheaper iterations.
7 Limitations and Extensions
The framework supports many SGD variants and extensions, but its current scope excludes some important settings and leaves several refinements open. It also recovers existing methods and introduces practical sampling and variance-reduction variants within its assumptions.
- Limitations: The theory does not yet extend to the nonconvex case, although the authors believe weakly convex extensions are possible.
- Limitations: A unified treatment of biased gradient estimators remains open, including recovery of SAG and rates for zero-order optimization.
- Limitations: The framework does not recover the best known rates for randomized coordinate descent with importance sampling.
- Extensions: Extending the assumption to iteration-dependent parameters could enable analyses of methods such as SGD with decreasing stepsizes.
- Extensions: A unified analysis of stochastic methods with acceleration and momentum is identified as a promising future direction.
- SGD-MB: With-replacement sampling gives SGD-MB independent importance sampling in O(τ log(n)) time, versus O(n) without replacement.
- Recovered methods: The framework recovers convergence results for SGD-star, SAGA, and SEGA as direct consequences of its general theorem.
A.11 DIANA
This section shows how the unified framework recovers DIANA, quantized SGD, variance-reduced DIANA, and JacSketch, including convergence behavior under exact and stochastic gradients.
- DIANA: DIANA uses quantization of gradient differences to reduce communication between distributed workers and the server.The method quantizes differences rather than actual gradients, which supports convergence to the optimum when full gradients are evaluated.
- DIANA: With full deterministic gradients, DIANA converges linearly to the exact optimum.
- Other recovered methods: The framework also recovers quantized SGD-SR, JacSketch, and τ-L-SVRG as special cases under their respective assumptions.JacSketch maintains and updates a Jacobian estimate from random sketches, while τ-L-SVRG interpolates between vanilla SGD and L-SVRG.
- Variance-reduced DIANA: With stochastic gradients, DIANA converges linearly only to a solution neighborhood, whereas VR-DIANA converges linearly to the exact solution.VR-DIANA applies variance reduction within each machine to eliminate the stochastic-gradient noise effect.
B.1 SGD-MB: remaining experiments and exact problem setup.
The SGD-MB experiments evaluate logistic regression on LIBSVM datasets and compare its sampling implementation with independent SGD.
- Problem setup: The experiments use logistic regression with Tikhonov regularization on eight LIBSVM datasets.The datasets are a1a, a9a, w1a, w8a, gisette, madelon, phishing, and mushrooms, with data rescaled before evaluation.
- Results: SGD-MB and independent SGD have essentially identical iteration complexity in practice.The comparison includes uniform and importance sampling labels, with replacement used for SGD-MB and without replacement for independent SGD.
B.2 Experiments on SGD-star
The SGD-star experiments compare SGD-star, SGD, and SAGA on generated least-squares problems, testing the convergence relationship predicted by the theory.
- Results: SGD-star is always faster than SAGA in the reported experiments, although only by constant factors.The experiments use a least-squares problem with independently generated standard-normal matrix and response entries.
- Results: When d ≥ n, SGD and SGD-shift have identical performance because ∇f_i(x*) = 0 for every i.Under this condition, the stochastic-gradient shift is zero, making the two algorithms identical.
B.3 Experiments on N-SEGA
The N-SEGA experiment studies noisy partial derivatives on constrained least-squares problems and examines how oracle noise affects convergence.
- Experimental setup: N-SEGA is evaluated on unit-ball-constrained least-squares problems with Gaussian noise added to partial-derivative estimates.The noise magnitude σ^2 is varied across problem instances.
- Results: N-SEGA exhibits linear convergence to a specific objective neighborhood, as predicted by the theory.The experiment is intended to support and provide intuition for the theoretical results rather than serve as a practical evaluation.
- Results: Larger strong convexity μ generally reduces the effect of the oracle noise.The reported examples identify problems 1 and 3, without scaling, as having larger μ and a smaller noise effect.
C Proofs for Section 4
The proof develops a single-iteration Lyapunov recurrence under Assumptions 4.1 and 4.2, then rearranges and unrolls it to support the main theorem.
- Lemma C.1 establishes a key single-iteration recurrence for all k ≥ 0 under Assumptions 4.1 and 4.2.
- The proof estimates the distance term using r_k = x_k − x* and conditional expectation given x_k.
- The full expectation of the next Lyapunov iterate V^{k+1} is bounded using the preceding distance estimate.
- After rearrangement, a positive term implied by condition (13) can be omitted from the inequality.
- Unrolling the recurrence and using Lyapunov-function notation completes the transition from the single-step bound to the broader proof argument.