Source-linked AI summary

SGD: General Analysis and Improved Rates

Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, Peter Richtarik

arXiv:1901.09401v4cs.LGmath.OCstat.ML

TL;DR

SGD lacked a unified convergence analysis covering arbitrary sampling without relying on uniformly bounded stochastic gradients. The paper uses expected smoothness to analyze arbitrary sampling schemes, deriving convergence, sampling-dependent stepsizes, mini-batch complexity, and stepsize-switching rules. It shows that the optimal mini-batch size grows with gradient noise at the minimum and equals one when the noise is zero.

  • Problem

    Existing SGD analyses did not provide a unified treatment of virtually all sampling schemes and often relied on bounded-gradient or growth assumptions.

  • Method

    The paper analyzes SGD through expected smoothness, a joint property of the objective and sampling distribution, under arbitrary sampling schemes.

  • Results

    SGD converges linearly up to the additive constant 2γσ2/µ, with the bound depending on gradient noise σ2 and stepsize γ.

  • Takeaways & Limitations

    The framework yields sampling-dependent stepsizes and complexity choices, including mini-batch rules and a closed-form switch from constant to decreasing stepsizes.

  • Takeaways & Limitations

    Optimizing iteration complexity over sampling probabilities leads to an n-dimensional linearly constrained nonsmooth convex problem that may be harder than the original problem.

Abstract

from arXiv · show

We propose a general yet simple theorem describing the convergence of SGD under the arbitrary sampling paradigm. Our theorem describes the convergence of an infinite array of variants of SGD, each of which is associated with a specific probability law governing the data selection rule used to form mini-batches. This is the first time such an analysis is performed, and most of our variants of SGD were never explicitly considered in the literature before. Our analysis relies on the recently introduced notion of expected smoothness and does not rely on a uniform bound on the variance of the stochastic gradients. By specializing our theorem to different mini-batching strategies, such as sampling with replacement and independent sampling, we derive exact expressions for the stepsize as a function of the mini-batch size. With this we can also determine the mini-batch size that optimizes the total complexity, and show explicitly that as the variance of the stochastic gradient evaluated at the minimum grows, so does the optimal mini-batch size. For zero variance, the optimal mini-batch size is one. Moreover, we prove insightful stepsize-switching rules which describe when one should switch from a constant to a decreasing stepsize regime.

1. Introduction

The paper develops a unified analysis of SGD under arbitrary sampling, covering broad sampling and mini-batch variants without requiring convex component functions or a direct growth-condition assumption. It derives sampling-dependent learning-rate and complexity results, including optimal mini-batch behavior and stepsize switching.

  • General analysis: Theorem 3.1 analyzes SGD for virtually all sampling schemes through expected smoothness, including settings where neither f nor the fi are convex.The analysis provides linear convergence under strong quasi-convexity and uses a joint property of the objective and sampling distribution.
  • Assumptions: The analysis uses expected smoothness instead of directly assuming a growth condition or uniformly bounded stochastic-gradient variance.This assumption is presented as weaker than the direct growth-condition approach used in prior analyses.
  • Mini-batching: Exact stepsize formulas show that the optimal learning rate increases and iteration complexity decreases with mini-batch size for independent sampling and sampling with replacement.The paper proves the increase only up to a certain optimal mini-batch size.
  • Mini-batching: For mini-batch size n, the analysis recovers the gradient-descent rate L/µ log(1/ϵ), while zero-noise over-parameterized models have optimal mini-batch size 1.The latter result extends earlier findings to independent sampling and sampling with replacement.
  • Learning schedules: The paper provides a closed-form rule for switching SGD from a constant to a decreasing stepsize.This rule is stated as Theorem 3.2.
  • Experiments: The theoretical results are corroborated by extensive experimental testing.

2. Expected Smoothness and Gradient Noise

The paper formulates expected smoothness as a joint condition on the objective and sampling distribution, then relates it to gradient-noise and weak-growth conditions. This framework supports convergence analysis beyond convex objectives and avoids uniformly bounded stochastic-gradient assumptions.

  • Expected smoothness: Expected smoothness combines properties of the sampling distribution D with smoothness properties of f.It is formalized by an inequality defining an expected smoothness constant L=L(f,D).
  • Expected smoothness: Expected smoothness can hold when both the component functions fi and the aggregate function f are non-convex.The paper gives a smooth non-convex function with a global minimum as an example.
  • Gradient noise: Finite gradient noise is treated as a weak assumption primarily concerning the sampling distribution D rather than f.For nonnegative sampling vectors with finite relevant expectations, σ is finite; over-parameterized models have σ=0.
  • Gradient noise: The analysis does not assume uniformly bounded stochastic gradients, an assumption that may fail even for strongly convex objectives.Instead, expected smoothness supplies a direct bound on the expected stochastic-gradient norm.
  • Weak growth: When σ=0, expected smoothness implies the weak growth condition with ρ=2L.This bound can be tighter than prior bounds using ρ=2Lmax because Lmax≥L.
  • Gradient noise: Uniform gradient bounds along iterates can create a circular argument because bounded iterates are implicitly assumed while proving convergence into a bounded region.This concern is raised for an alternative assumption used in prior analyses.

3. Convergence Analysis

The analysis establishes convergence guarantees for SGD under arbitrary sampling using expected smoothness and gradient noise, including linear convergence without convexity and a decreasing-stepsize regime.

  • Main convergence theorem: SGD converges linearly up to the additive constant 2γσ2/µ, which depends on gradient noise and the stepsize.Smaller stepsizes improve solution accuracy but slow the convergence rate.
  • Main convergence theorem: The framework does not require f_i or f to be convex, while allowing control of both σ2 and L through the sampling distribution D.The paper computes these parameters for several sampling schemes.
  • Decreasing stepsizes: Theorem 3.2 gives a decreasing-stepsize result after a threshold involving K := L/µ and iterations k ≥ 4⌈K⌉.This result controls the additive constant through a carefully chosen stepsize schedule.
  • Sampling framework: A proper sampling generates a valid sampling vector v, and each sampling law induces a corresponding stochastic reformulation and SGD variant.The analysis specializes expected-smoothness and gradient-noise bounds to independent, partition, single-element, and τ-nice samplings.
  • Sampling-specific bounds: Under convex smooth component functions, the paper derives closed-form expected-smoothness and gradient-noise expressions or bounds for several sampling strategies.Theorem 3.6 uses the component smoothness matrices and sampling pair probabilities to bound L.
  • Sampling-specific bounds: The paper proves that the approximation in (24) is a valid upper bound on L, extending earlier experimental evidence of its tightness.This is presented as the first proof of that bound’s validity.

4. Optimal Mini-Batch Size

The paper studies how mini-batch size affects iteration and total complexity, showing that larger batches improve iteration complexity but need not reduce total gradient evaluations.

  • Gradient descent: The full-batch case recovers gradient descent with σ = 0, Lmax = L, and iteration complexity k ≥ 2L/µ.The authors identify this as the first generic SGD convergence theorem to recover the gradient-descent rate as a special case.
  • Independent sampling: Independent sampling yields a new mini-batch complexity result that permits optimizing both mini-batch size and sampling probabilities.For fixed uniform probabilities, the resulting stepsize is given explicitly as a function of τ.
  • Total complexity: The optimal mini-batch size is τ* = 1 when smoothness dominates noise at τ = 1; otherwise, the optimum follows the balancing condition between the two terms.The condition changes with noise and target precision.
  • Stepsize scaling: The stepsize increases with mini-batch size because the relevant smoothness and noise terms decrease as τ increases.The stated monotonicity applies to the stepsize expression for independent and τ-nice sampling.
  • Total complexity: Total complexity is defined as iterations multiplied by τ gradient evaluations per iteration and is optimized by balancing increasing and decreasing terms.The analysis reduces the objective to a maximum of τ × l(τ) and τ × r(τ).
  • Zero-noise regime: With zero gradient noise, the iteration complexity improves as τ increases to n, but total complexity is optimized at τ = 1.This conclusion is reported for the independent, τ-nice, and related zero-noise cases.

5. Importance Sampling

The paper proposes importance-sampling schemes for single elements and minibatches, using sampling probabilities chosen to control complexity parameters. These schemes can substantially improve iteration or total complexity over uniform sampling, especially for minibatches.

  • Importance-sampling schemes: Importance sampling is proposed for single-element sampling and independent sampling with expected minibatch size τ.The minibatch construction is introduced as a new importance-sampling approach for minibatch SGD.
  • Importance-sampling schemes: Optimizing iteration complexity directly over sampling probabilities is a linearly constrained nonsmooth convex problem, so the paper separately minimizes Lmax and σ2.The resulting probabilities are then combined to construct a sampling scheme.
  • Single-element sampling: Partially biased sampling recovers the scheme of Needell et al. (2016) and yields a total complexity bounded using Lmax and the average squared norms ∑i∈[n]∥hi∥2.The probabilities combine Lmax-minimizing probabilities with uniform sampling.
  • Single-element sampling: Compared with uniform sampling, partially biased sampling is at most two times larger in iteration complexity but can be n/2 smaller when Lmax = nL.The comparison is stated for the extreme case where the maximum smoothness parameter is n times the average parameter.
  • Minibatch sampling: For independent minibatch sampling, the proposed importance-sampling variant eliminates dependence on Lmax and improves as the minibatch size τ increases.The result is stated for the construction using partially biased sampling with independent sampling.

6. Experiments

Experiments evaluate SGD on strongly convex regression problems and PCA, comparing stepsize regimes and sampling strategies across synthetic and real data. The results support the theoretical benefits of decreasing stepsizes, minibatching, importance sampling, and theoretically selected batch sizes.

  • Experimental setup: Experiments cover ridge regression, regularized logistic regression, and PCA using synthetic and real datasets, with comparisons across sampling distributions.PCA is treated as having a strongly convex objective with non-convex component functions.
  • Experimental setup: SGD is evaluated by relative error from x∗, starting from a standard-Gaussian point and stopping at error 10^-3 or a maximum epoch count.The horizontal axis in the experiments is the number of epochs.
  • Stepsize regimes: The decreasing stepsize regime is vastly superior at reaching higher precision than the constant-stepsize variant in ridge and logistic regression experiments.A vertical line marks 4⌈L/µ⌉, the theorem-predicted point for switching from constant to decreasing stepsizes.
  • Minibatch strategies: Minibatching with the theoretically optimal τ∗ gives the best convergence, while importance-sampling variants are comparable to or better than corresponding uniform variants.The comparisons include single-element, τ-independent, and τ-nice sampling strategies.
  • PCA: For PCA with non-convex component functions, decreasing stepsizes reach higher accuracy and the theoretically optimal τ = τ∗ yields a vast improvement over single-element SGD.All four minibatch variants outperform single-element SGD, but τ = n/5 produces only a slight improvement.

A. Elementary Results

This appendix section collects elementary smoothness, counting, and complexity results used repeatedly in the paper’s proofs. It explicitly presents these results for completeness and clarity rather than as novel contributions.

  • Smoothness results: The appendix introduces an elementary smoothness proposition for an Lφ-smooth function with a minimizer x∗.Its proof uses Lipschitz continuity of the gradient and the fact that ∇φ(x∗) = 0.
  • Role of the appendix: The section states that these elementary results are included for completeness and clarity, without claiming novelty.They are used repeatedly in the paper’s proofs.
  • Auxiliary lemmas: A double-counting lemma is stated for scalar quantities indexed by elements and collections of subsets of [n].The lemma is listed among the elementary results used in later proofs.
  • Auxiliary lemmas: A complexity-bounds lemma gives conditions involving E, ρ, and c under which an integer iteration count satisfies a prescribed bound.The surrounding proof derives the result by taking logarithms and rearranging the inequality.

B. Proof of Lemma 2.4

The proof establishes the decreasing-stepsize complexity bound by choosing a stepsize sequence that eventually falls below a smoothness threshold. It then treats early and later iterations separately and selects a balancing index.

  • Stepsize construction: The proof uses the decreasing sequence γk = (2k+1)/((k+1)^2µ) and chooses k∗ such that γk∗ ≤ 1/(2L).This choice ensures the later-iteration smoothness condition holds.
  • Stepsize construction: Because γk decreases with k, the threshold condition remains valid for every k ≥ k∗.The proof consequently applies the relevant bound with γk in place of a constant stepsize.
  • Later-iteration bound: For later iterations, the proof multiplies by (k + 1)^2, sums from t = k∗ to k, and uses telescopic cancellation.These algebraic steps produce the bound for the decreasing-stepsize phase.
  • Two-regime analysis: Choosing the minimizing index gives k∗ = 4⌈K⌉, which is then substituted into the resulting bound.The proof uses this choice to complete the complexity analysis.

D. Proof of Theorem 3.6

This proof establishes bounds on the expected smoothness quantities under fixed-size sampling, using smoothness and convexity properties of the component functions.

  • Each component function's smoothness and convexity support the inequalities used to bound the relevant quantities.
  • For fixed-size sampling, the proof specializes the bounds by taking expectations and maximizing over component indices.
  • The proof derives lower and upper bounds on the expected smoothness constant L = Lmax for fixed-size sampling.

G. Proof of Proposition 3.8

This proof derives expected-smoothness bounds for independent and τ-nice sampling, then uses them to characterize importance-sampling probabilities and iteration complexity.

  • For independent sampling, Pij = pipj for i ≠ j, giving c2 = 1 and the corresponding expected-smoothness bound.
  • For τ-nice sampling, the pairwise probabilities yield c2 = n(τ−1) / [τ(n−1)] and the corresponding bound.
  • The probabilities minimizing Lmax are identified, while the resulting importance-sampling iteration complexity is at most two times the minimum over probabilities.
  • Interpolated probabilities preserve valid sampling probabilities and provide an upper bound on the iteration complexity.

J.2. Independent sampling

This section analyzes partially biased independent sampling, including probability construction and bounds showing favorable smoothness and variance scaling with minibatch size.

  • The construction clips probabilities at one for indices whose preliminary values exceed one, then adjusts the remaining probabilities to meet the minibatch constraint.
  • The resulting upper bounds are nonincreasing as sampling probabilities increase, supporting the stated partially biased construction.
  • Partially biased probabilities combine a smoothness-based distribution with uniform sampling when the component information needed for optimal probabilities is unavailable.
  • The sampling keeps Lmax close to L while achieving nearly linear speedup in σ2 as τ increases.
  • Compared with τ-nice and τ-uniform independent sampling, its iteration complexity is at most two times larger and can be about 2τ/n smaller in an extreme case.

K.1. From fixed to decreasing stepsizes: analysis of the switching time

The experiments examine switching from constant to decreasing stepsizes and compare minibatch strategies and stepsize behavior across synthetic and real data.

  • The switching-time experiment compares theoretical, early, late, and optimized regime-switch choices on synthetic ridge-regression data.
  • When x0 is close to x∗, the switching moment has little effect except for a very early switch.
  • When x0 is far from x∗, preliminary one-dimensional optimization reduces the error during the early iterations.
  • For τ-nice sampling, the stepsize increases with minibatch size τ across three synthetic regularized ridge-regression datasets.
Loading 1901.09401v4…