Source-linked AI summary
Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm
Deanna Needell, Nathan Srebro, Rachel Ward
TL;DR
The paper studies how to improve convergence guarantees for smooth, strongly convex stochastic gradient descent and connects SGD with randomized Kaczmarz. It analyzes reweighted sampling and transfers insights between the two algorithms to obtain tighter conditioning guarantees and discuss sampling limitations.
Problem
Prior strongly convex SGD guarantees depend on conditioning, motivating tighter convergence bounds and understanding when weighted sampling can improve them.
Method
The paper connects SGD with randomized Kaczmarz, analyzes reweighted SGD through alternative stochastic representations, and studies partially biased sampling.
Results
The paper improves SGD's conditioning dependence from quadratic to linear under uniform sampling and identifies importance sampling as necessary for dependence on average conditioning.
Takeaways & Limitations
The SGD–Kaczmarz connection enables contributions across both literatures, including tighter SGD guarantees and broader discussion of importance sampling.
Takeaways & Limitations
Optimal weighting in the non-realizable regime depends on the unknown minimizer, while importance sampling may be impractical for online or source-distribution-only data.
Abstract
from arXiv · showhide
We obtain an improved finite-sample guarantee on the linear convergence of stochastic gradient descent for smooth and strongly convex objectives, improving from a quadratic dependence on the conditioning $(L/μ)^2$ (where $L$ is a bound on the smoothness and $μ$ on the strong convexity) to a linear dependence on $L/μ$. Furthermore, we show how reweighting the sampling distribution (i.e. importance sampling) is necessary in order to further improve convergence, and obtain a linear dependence in the average smoothness, dominating previous results. We also discuss importance sampling for SGD more broadly and show how it can improve convergence also in other scenarios. Our results are based on a connection we make between SGD and the randomized Kaczmarz algorithm, which allows us to transfer ideas between the separate bodies of literature studying each of the two methods. In particular, we recast the randomized Kaczmarz algorithm as an instance of SGD, and apply our results to prove its exponential convergence, but to the solution of a weighted least squares problem rather than the original least squares problem. We then present a modified Kaczmarz algorithm with partially biased sampling which does converge to the original least squares solution with the same exponential convergence rate.
1. INTRODUCTION
The paper connects stochastic gradient descent with randomized Kaczmarz, using the relationship to sharpen convergence guarantees and analyze weighted sampling. Importance sampling improves conditioning dependence in strongly convex SGD and extends benefits to other optimization settings.
- The paper connects randomized Kaczmarz for linear systems with SGD for convex optimization, transferring ideas between the two literatures.This connection highlights weighted sampling in SGD and yields tighter linear-convergence guarantees.
- SGD minimizes a convex objective from unbiased stochastic gradient estimates, often obtained by sampling a component function and using its gradient.The objective is represented as an expectation of component functions.
- Earlier strongly convex SGD results achieved exponential convergence under shared minimizers but depended on average squared conditioning and degraded with residual variance.When component minimizers differ, an unavoidable residual term behaves as σ2/k.
- Without reweighting, the paper obtains linear dependence on uniform conditioning but rules out linear dependence on average conditioning.This improves the conditioning dependence over previous results when component smoothnesses are similar.
- Partial biasing toward component smoothnesses achieves linear dependence on average conditioning without amplifying residual dependence.Fully weighting by Lipschitz constants reaches average conditioning but increases dependence on σ2; partial biasing avoids that trade-off.
- Importance sampling also improves smooth non-strongly-convex SGD and removes dependence on component Lipschitz-constant variance for non-smooth objectives.These improvements are unavailable without importance sampling in the stated settings.
- The paper recasts randomized Kaczmarz as SGD and obtains exponential convergence, then uses partially biased sampling to retain the rate while targeting the original least-squares solution.The hybrid strategy exposes a trade-off between convergence rate and residual that is absent from other stated results.
2. SGD FOR STRONGLY CONVEX SMOOTH OPTIMIZATION
This section analyzes SGD for smooth, strongly convex objectives and improves the conditioning dependence of finite-sample convergence guarantees. It shows why source-distribution sampling cannot achieve average conditioning and motivates importance sampling.
- The analysis considers minimizing a smooth convex function represented as an expectation of smooth component functionals over a Euclidean or Hilbert space.The strongly convex case has a unique minimizer x⋆, and convergence is measured by iterate distance to it.
- The assumptions require component gradients to have Lipschitz constants Li and the objective to be µ-strongly convex.The analysis also defines supremum, infimum, and average component Lipschitz constants.
- SGD samples indices independently from a source distribution and updates iterates using unbiased component gradients with a fixed step size.The target is a prescribed expected squared distance from the unique minimum.
- Earlier guarantees had iteration complexity proportional to expected squared conditioning, while this paper reduces that dependence to linear uniform conditioning.The improved recursion replaces one component-smoothness factor with the objective’s strong-convexity parameter µ.
- When component smoothnesses are similar, replacing average squared conditioning with uniform conditioning yields a quadratic improvement in required iterations.With highly variable Li, the supremum can instead exceed the average squared conditioning.
- Uniform source sampling cannot generally achieve linear average conditioning: an example requires at least N+1 samples despite average conditioning near 2.Both uniform conditioning and average squared conditioning match the correct SGD iteration scale in that example.
- Importance sampling obtains linear average conditioning by sampling from a modified distribution rather than the source distribution.This reweighting is introduced specifically to overcome the impossibility for source-distribution sampling.
3. IMPORTANCE SAMPLING
Importance sampling reweights SGD without changing the objective, minimizer, or suboptimality, allowing convergence bounds to depend on average rather than worst-case smoothness. Fully biased weights achieve linear average conditioning in the realizable case, while partially biased sampling preserves this scaling with only a constant residual penalty.
- 3.2. Reweighted SGD: Weighted SGD samples from D(w) and rescales component gradients by 1/w(i), preserving unbiasedness for the original objective.The weighted representation keeps F(x) and its minimizer unchanged while modifying component smoothness quantities.
- 3.3. Strongly Convex Smooth Optimization using Weighted SGD: Weights proportional to component smoothness minimize the weighted worst-case smoothness quantity used in the convergence bound.The paper identifies the fully biased choice w(i) = Li/L as the optimizer of this quantity.
- 3.3. Strongly Convex Smooth Optimization using Weighted SGD: Linear dependence on average conditioning is obtained for realizable strongly convex smooth SGD with fully biased sampling.Here σ² = 0, so the residual term vanishes and the desired linear-convergence guarantee follows.
- 3.4. Partially biased sampling: Partially biased sampling retains linear average-conditioning dependence without amplifying the residual term beyond a constant factor of two.This result dominates the Bach–Moulines guarantee up to that factor and improves quadratic conditioning dependence to linear.
- 3.4. Partially biased sampling: When component smoothness varies widely, importance sampling is necessary for average-conditioning dependence, but it can be difficult with online or source-only sampling access.Direct sampling is more feasible when smoothness constants are known or easily computed and repeated data passes are available.
- 3.5. Implementing Importance Sampling: Rejection sampling accepts source-distribution samples with probability proportional to Li/sup L, yielding linear dependence on sup L in sample complexity.It is most useful when operating on a sample costs more than obtaining it and estimating its smoothness bound.
4. IMPORTANCE SAMPLING FOR SGD IN OTHER SCENARIOS
The paper extends importance sampling beyond strongly convex smooth objectives to smooth non-strongly-convex and non-smooth settings. Reweighting can replace worst-case smoothness dependence with average smoothness and remove dependence on variance in component Lipschitz constants.
- 4.1. Smooth, Not Strongly Convex: Without importance sampling, smooth non-strongly-convex SGD bounds retain linear dependence on sup L, which cannot generally be replaced by average Lipschitz smoothness.The cited prior result identifies the supremum dependence as unavoidable under source-distribution sampling.
- 4.1. Smooth, Not Strongly Convex: For smooth but not strongly convex objectives, importance sampling improves dependence from sup Li to average smoothness E[Li].The paper states that this improvement is unavailable without importance sampling.
- 4.2. Non-Smooth Objectives: For non-smooth objectives, importance sampling eliminates dependence on variance in the components’ Lipschitz constants.The relevant standard SGD complexity depends linearly on G², which weighted SGD replaces with a less variable weighted quantity.
2. Using importance sampling,
Importance sampling improves SGD conditioning bounds, but the best weighting depends on whether convergence or residual terms dominate. Partial or average-based weighting can balance these competing objectives.
- Weighting by Gi removes the variance of the Lipschitz constants from the relevant bound, whereas weighting by Li can leave a suboptimal gradient-dependent term.
- The residual term σ2 depends on the sampling weights, and minimizing it would use weights proportional to ∥∇fi(x⋆)∥2.These optimal weights are unavailable when x⋆ is unknown.
- Linear dependence on average conditioning E[Li/µ] is achievable by weighting components proportionally to Li.This can increase dependence on the residual σ2.
- For generalized linear objectives, weighting by Li and weighting by Gi correspond to different dependencies on the data vectors zi.
5. THE LEAST SQUARES CASE AND THE RANDOMIZED KACZMARZ METHOD
The paper recasts randomized Kaczmarz as weighted SGD for least squares, transferring SGD convergence results to row-sampling variants. This yields finite-iteration guarantees, step-size tradeoffs, and modified methods with improved solution behavior.
- Randomized Kaczmarz iterates are precisely weighted SGD iterates with fully biased weights.
- Applying the improved SGD theorem recovers the Strohmer–Vershynin guarantee and reduces quadratic conditioning dependence to linear dependence.
- Using Step-sizes: Arbitrary step sizes provide finite-iteration guarantees and trade convergence rate against the convergence horizon.
- Uniform Row Selection: Uniform row selection converges exponentially toward a weighted least-squares solution, which can differ from the original least-squares solution.
- Partially Biased Sampling: Partially biased sampling matches fully biased sampling’s convergence rate up to a factor of two while improving dependence on residual error.
6. NUMERICAL EXPERIMENTS
Experiments vary the sampling mixture λ across five matrix and noise settings. The preferred sampling strategy depends on row-norm variation and residual-noise level.
- Weighted sampling significantly outperforms uniform sampling when matrix rows are far from normalized, with performance improving monotonically as λ decreases.
- When rows are close to normalized, different λ values produce similar convergence rates.
- In the high-noise regime, fully weighted sampling is relatively slow, while hybrid sampling outperforms both weighted and uniform selection.
- For the fixed-error iteration count, pure weighted sampling is optimal in Cases 2 and 5, while Cases 3 and 4 prefer intermediate λ values.
- Uniform sampling in Case 1 failed to attain the target error before the iteration count was cut off at 50,000.
7. SUMMARY AND OUTLOOK
The paper’s contributions are improved SGD conditioning bounds, importance-sampling analysis, and the SGD–Kaczmarz connection. Its scope includes fixed-step-size analysis, strong convexity, and static rather than dynamic reweighting.
- The paper makes three contributions: improved conditioning dependence for smooth strongly convex SGD, importance-sampling analysis, and the SGD–Kaczmarz connection.
- The analysis uses fixed step sizes, which are sufficient for optimal iteration complexity when the target accuracy is known.
- The strong-convexity assumption might be weakened to restricted strong convexity, but this is left for future work.
- The importance-sampling discussion is limited to static reweighting, with no rigorous analysis of dynamically updated sampling distributions.
APPENDIX A. PROOFS
The appendix’s main results rely on co-coercivity of the gradient for smooth functions with Lipschitz-continuous gradients. The appendix states this elementary fact and recalls its proof.
- The main results use co-coercivity of the gradient as an elementary fact about smooth functions.
- Co-coercivity applies when a function has a Lipschitz-continuous gradient.
- The appendix includes the lemma and its proof for completeness.
A.1. The Co-coercivity Lemma.
The co-coercivity lemma is stated for smooth functions with gradient Lipschitz constant L and established by applying a minimizer-based inequality to two auxiliary convex functions.
- The lemma concerns a smooth function whose gradient has Lipschitz constant L.
- The proof begins from an inequality involving a minimizer x⋆ of the function.
- Two convex auxiliary functions are defined, each with Lipschitz constant L and respective minimizers x and y.
- Applying the inequality to both auxiliary functions and combining the resulting bounds yields the desired result.
- Theorem 2.1’s notation tracks the random index i selected at iteration k and sets w = wλ.
A.2. Proof of Theorem 2.1.
The proof expands one stochastic-gradient update, bounds its squared distance using Jensen’s inequality and co-coercivity, then takes expectations and recursively applies strong convexity.
- The squared distance after the stochastic update is expanded into inner-product and gradient-difference terms.
- Jensen’s inequality and the co-coercivity lemma bound the update expansion.
- Taking expectation over i uses F(x) = Efi(x), E∇fi(x) = ∇F(x), and σ2 = E∥∇fi(x⋆)∥2.
- The expected-distance bound contains the strong-gradient term, a smoothness-weighted term, and the variance contribution 2γ2σ2.
- Strong convexity of F supplies the next bound, which is recursively applied over the first k iterations.