PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Quasi-Gradient Methods: Variance Reduction via Jacobian Sketching

Robert M. Gower, Peter Richtárik, Francis Bach

arXiv:1805.02632v1math.OC

TL;DR

Minimizing a large average of smooth functions is costly when each iteration scans all functions, and SGD retains gradient-estimate variance. JacSketch maintains and updates a Jacobian estimate from random measurements to form variance-reduced gradient estimates. Its analysis confirms the conjectured rate for importance-sampled SAGA and improves minibatch SAGA rates, including a first result for importance-sampled minibatches.

  • Problem

    For large finite sums, full-data iterations are costly, and the theory for SAGA’s importance-sampling and minibatch variants was incomplete.

    Source: §1 · p. 4, §1.2 · p. 5, §1.7 · p. 13

  • Method

    JacSketch maintains a Jacobian estimate and projects it onto solutions consistent with cheap random linear measurements to form gradient estimates.

    Source: §1.3 · p. 6

  • Results

    The analysis confirms Schmidt et al.’s conjecture for importance-sampled SAGA and gives improved minibatch SAGA rates, including the first result for importance-sampled minibatch SAGA.

    Source: §1.7 · p. 13

  • Takeaways & Limitations

    JacSketch specializes to SAGA variants, with importance sampling and minibatching covered by the framework.

    Source: §1.7 · pp. 12–13, §4 · p. 26

    For the small-dimensional problems discussed, predicted total-complexity improvements are modest, with τ = 2 or 3 optimal.

    Source: §6.2 · p. 44

Abstract

from arXiv · show

We develop a new family of variance reduced stochastic gradient descent methods for minimizing the average of a very large number of smooth functions. Our method --JacSketch-- is motivated by novel developments in randomized numerical linear algebra, and operates by maintaining a stochastic estimate of a Jacobian matrix composed of the gradients of individual functions. In each iteration, JacSketch efficiently updates the Jacobian matrix by first obtaining a random linear measurement of the true Jacobian through (cheap) sketching, and then projecting the previous estimate onto the solution space of a linear matrix equation whose solutions are consistent with the measurement. The Jacobian estimate is then used to compute a variance-reduced unbiased estimator of the gradient. Our strategy is analogous to the way quasi-Newton methods maintain an estimate of the Hessian, and hence our method can be seen as a stochastic quasi-gradient method. We prove that for smooth and strongly convex functions, JacSketch converges linearly with a meaningful rate dictated by a single convergence theorem which applies to general sketches. We also provide a refined convergence theorem which applies to a smaller class of sketches. This enables us to obtain sharper complexity results for variants of JacSketch with importance sampling. By specializing our general approach to specific sketching strategies, JacSketch reduces to the stochastic average gradient (SAGA) method, and several of its existing and many new minibatch, reduced memory, and importance sampling variants. Our rate for SAGA with importance sampling is the current best-known rate for this method, resolving a conjecture by Schmidt et al (2015). The rates we obtain for minibatch SAGA are also superior to existing rates.

1 Introduction

Large-scale empirical risk minimization makes full-data iterations costly, while SGD’s gradient-estimate variance does not vanish. The paper introduces JacSketch, which learns a Jacobian estimate from random measurements to form variance-reduced gradient estimates.

1 Introduction

  • Large datasets make scanning every function per iteration costly, while SGD’s gradient-estimate variance persists and requires a tuned decreasing stepsize.In empirical risk minimization, each function represents one data point’s loss.Source: §1 · p. 4

1.1 Variance-reduced methods

  • Variance-reduced methods address SGD’s high-variance issue; SAGA uses stored per-example gradients for an unbiased estimate, while SVRG uses full-gradient outer loops.SAG’s estimator is biased, and SARAH is unbiased over a history rather than at its last step.Source: §1.1 · p. 4

1.2 Gaps in our understanding of SAGA

  • The paper identifies unresolved SAGA theory: importance sampling lacked the conjectured rate, existing minibatch rates lagged, and importance-sampled minibatch SAGA was unknown.The authors state that their contributions give positive answers to these questions.Source: §1.2 · p. 5

1.3 Jacobian sketching: a new approach to variance reduction

  • The approach reuses random linear information about the Jacobian gathered in prior iterations, rather than discarding it as ordinary stochastic-gradient sampling does.The paper motivates this as a way to address SGD’s large variance.Source: §1.3 · p. 6
  • JacSketch maintains a Jacobian estimate, updates it by projecting onto solutions consistent with a cheap random measurement, and uses it to form an unbiased gradient estimate.The sketch-and-project update is analogous to quasi-Newton Hessian estimation; under appropriate conditions, the method’s variance vanishes and it converges linearly for smooth, strongly convex functions.Source: §1.3 · pp. 6–7

1.4 SAGA as a special case of JacSketch

JacSketch recovers SAGA and supports importance-sampling and minibatch variants, including a reduced-memory variant. Its analysis establishes the conjectured importance-sampling rate and reports improved minibatch SAGA rates.

1.4 SAGA as a special case of JacSketch

  • Choosing singleton sketches recovers standard SAGA, while choosing singleton probabilities p_i recovers SAGA with arbitrary sampling probabilities.Uniform SAGA uses probability 1/n and θ_Sk = n; arbitrary probabilities use θ_Sk = 1/p_i.Source: §1.4 · p. 8
  • A different sketch yields a reduced-memory SAGA variant while retaining SAGA’s gradient-estimator formula.The paper presents this as a special case of its general approach.Source: §1.4 · p. 8
  • For partition-based sketches, identical Jacobian columns within each partition can be stored as one column, yielding a compressed-memory method analyzed with non-uniform probabilities.The shared-column property holds when the initialization assigns equal columns within each partition.Source: §1.4 · p. 9

1.5 Sketch and project

JacSketch’s Jacobian-learning update connects sketch-and-project methods to variance reduction: it uses SGD both to minimize the loss and to learn the Jacobian used to lower search-direction variance. The paper also frames JacSketch as SGD on a controlled stochastic reformulation and derives general and specialized complexity results for SAGA variants.

1.5 Sketch and project

  • JacSketch applies SGD both to minimizing the loss and to learning the Jacobian, which is then used to lower search-direction variance.For constant Jacobians, sketch-and-project methods provide a connection to randomized iterative methods for linear systems.Source: §1.5 · p. 9

1.6 Controlled stochastic reformulation

  • A controlled reformulation writes the objective as an expectation, adds a Jacobian-dependent zero-mean perturbation to reduce gradient variance, then applies SGD and updates the Jacobian by sketch-and-project.The perturbation leaves the objective unchanged.Source: §1.6 · p. 10

1.7 Summary of complexity results

  • The general theorem covers unbiased sketches and positive-definite weight matrices, with complexity depending on expected-smoothness constants, a stochastic condition number, and sketch residual.Its special cases include gradient descent and uniform SAGA.Source: §1.7 · pp. 10, 12
  • The minibatch results cover virtually all strategies, and the bounds for τ-nice and τ-partition sampling can each be slightly better depending on the smoothness constants.For τ-partition sampling, W = Diag(L_i) gives a better complexity than W = I.Source: §1.7 · p. 12
  • The minibatch-SAGA bounds interpolate between the single-sample and gradient-descent complexities as τ increases, and are precise enough to inform an optimal minibatch size.The previous best bound was not precise enough to inform that choice.Source: §1.7 · p. 12
  • For τ-partition minibatch sketches, the refined theorem gives stronger iteration-complexity guarantees for non-uniform probabilities.For τ-partition sampling, the stated complexity with W = Diag(L_i) is better than with W = I.Source: §1.7 · p. 12
  • The refined theorem proves the conjectured importance-sampling rate for SAGA and establishes a complexity result for minibatch SAGA with importance sampling that is better than uniform probabilities.The latter is reported as the first result for minibatch SAGA with importance sampling.Source: §1.7 · pp. 12–13

1.8 Outline of the paper

The paper develops JacSketch’s motivation, general convergence theory, minibatch specializations, refined analysis, and computational experiments across Sections 2–6. It also provides a notation glossary in Section D.

1.8 Outline of the paper

  • Section 2 presents an alternative motivation for JacSketch based on controlled stochastic optimization reformulations.The paper describes this as a novel technical tool.Source: §1.8 · p. 13
  • Section 3 develops general convergence theory, and Section 4 specializes it to minibatch sketches and computes constants for selected samplings.The main result in Section 3 is Theorem 3.6.Source: §1.8 · p. 13
  • Section 5 develops an alternative theory using a stochastic Lyapunov function, and Section 6 contains computational experiments.The main result in Section 5 is Theorem 5.2.Source: §1.8 · p. 13

1.9 Notation

  • A notation glossary is available in Section D, and log t denotes the natural logarithm of t.Matrices are written in upper-case bold letters.Source: §1.9 · p. 13

2 Controlled Stochastic Reformulations

JacSketch broadens stochastic optimization through sketch-based reformulations, then controls estimator variance by maintaining and updating a Jacobian estimate. Its algorithm combines that estimate with an unbiased gradient estimator, while related variants show how biased estimates connect to SAG.

2.1 Stochastic reformulation using sketching

  • Unbiased sketches yield unbiased estimates of both the loss and its gradient, allowing SGD on equivalent stochastic reformulations whose sketch distribution can shape cost and complexity.The reformulations include gradient descent, non-uniform SGD, and minibatch SGD as special cases.Source: §2.1 · pp. 14–15

2.2 The controlled stochastic reformulation

  • Controlled stochastic reformulations add a control variate whose linear model mimics the sampled function, producing an unbiased gradient estimate for any Jacobian parameter J.The method chooses J to reduce estimator variance; the exact Jacobian minimizes it but is costly to compute.Source: §2.2 · pp. 15–16

2.3 The Jacobian estimate, variance reduction and the sketch residual

  • The sketch residual ρ measures average information loss from sketching and weighting and is a key constant affecting JacSketch’s convergence rate.It is the largest eigenvalue of W^1/2BW^1/2.Source: §2.3 · p. 17
  • JacSketch updates its Jacobian estimate by projecting it toward the true Jacobian subject to consistency with a single sketch, reusing that sketch for the gradient calculation.The update makes the estimate closer to the true Jacobian and reduces variance; sparse weighting can lower update cost.Source: §2.3 · pp. 17–18

2.4 JacSketch Algorithm

  • JacSketch combines the controlled-reformulation gradient estimator with a sketch-based Jacobian update to form its algorithm.The same approach also admits a controlled-reformulation view.Source: §2.4 · pp. 18–19

2.5 A window into biased estimates and SAG

  • Using the same Jacobian update with an uncorrected gradient average recovers biased SAG; correcting the sampling bias instead gives the unbiased SAGA estimator.The paper leaves the broader strategy of trading bias for lower variance for future work.Source: §2.5 · pp. 19–20

3 Convergence Analysis for General Sketches

The section develops a convergence analysis for JacSketch that applies to general sketches under expected-smoothness and nondegeneracy conditions. It establishes linear convergence and identifies conditions governing the rate and the decay of gradient-estimator variance.

3.1 Two expected smoothness constants

  • The analysis applies to general sketch distributions and introduces two expected-smoothness assumptions linking the objective, its Jacobian, and the sketch distribution.The authors describe the Jacobian expected-smoothness condition as new; they provide formulas for L1 for smooth functions and a class of minibatch samplings.Source: §3 · p. 20, §3.1 · pp. 20–21

3.2 Stochastic condition number

  • The stochastic condition number satisfies 0 ≤ κ ≤ 1, and κ > 0 requires a sufficiently rich sketch distribution.For unit-vector sketches with W = I, κ equals the smallest sampling probability.Source: §3.2 · p. 21

3.3 Convergence theorem

  • Under one-point strong convexity, an unbiased sketch, two expected-smoothness assumptions, and κ > 0, JacSketch’s iterates and Jacobian estimates converge linearly.The theorem uses a Lyapunov function combining distances of the iterates and Jacobian estimates from their optima.Source: §3.3 · p. 22

3.4 Projection lemmas and the stochastic condition number κ

  • Projection bounds relate sketching to contraction in the W^-1 norm, with the expected projection governed by κ.The bound in the projection lemmas is (1 − κ)W^-1.Source: §3.4 · pp. 22–23

3.5 Key lemmas

  • As xk and Jk approach their optima, the second-moment bound implies that JacSketch’s gradient-estimator variance approaches zero.This diminishing variance is identified as a key property of the method.Source: §3.5 · p. 24

3.6 Proof of Theorem 3.6

  • The proof combines one-point strong convexity, the gradient second-moment bound, and a stepsize satisfying two inequalities to obtain geometric decay.The resulting bound is E[Ψk] ≤ (1 − αμ)^k Ψ0.Source: §3.6 · pp. 25–26

4 Minibatch Sketches

This section specializes JacSketch to random minibatches, examining how sampling and weighting choices shape SAGA updates, expected smoothness and sketch-residual bounds, and resulting complexity analyses.

4 Minibatch Sketches

  • Minibatch sketches select columns of the identity indexed by a random subset, enabling each iteration to compute gradients only for sampled functions.The subset is drawn according to a fixed probability law.Source: §4 · p. 26

4.1 Samplings

  • A sampling is a random subset-valued mapping; the analysis assumes it is non-vacuous and has c1-uniform support, with 1-uniform support receiving special attention.The latter case has support forming a partition of [n] and permits stronger analysis than the general theorem.Source: §4.1 · pp. 26–27

4.2 Minibatch sketches and projections

  • The section establishes basic properties of the minibatch projection matrix and gives an additional lemma for samplings that are both τ-uniform and have c1-uniform support.The associated stochastic condition number is also defined for minibatch sketches.Source: §4.2 · pp. 27–28

4.3 JacSketch for minibatch sampling = minibatch SAGA

  • JacSketch with minibatch sketches yields minibatch SAGA, with standard, nonuniform, partition-based, τ-nice, and full-gradient cases as specializations.The method samples indices, updates the Jacobian estimate, and uses it to form the gradient estimate.Source: §4.3 · pp. 28–29

4.4 Expected smoothness constants L1 and L2

  • For convex smooth components and non-vacuous samplings with uniform support, Theorem 4.16 gives expected smoothness constants L1 and L2, with specialized formulas for τ-uniform sampling.The section also relates these constants to L, Lmax, and averages of component-set smoothness constants under stated sampling conditions.Source: §4.4 · pp. 30–32

4.5 Estimating the sketch residual ρ

  • Under non-vacuous, τ-uniform sampling with c1-uniform support, the section estimates the sketch residual and develops upper bounds for τ-partition and τ-nice samplings.Theorem 4.19 addresses uniform sampling over a support family with equal pair-incidence counts.Source: §4.5 · pp. 33–34
  • When τ≥2, τ-nice sampling satisfies the assumptions of the sketch-residual theorem, so bound (95) simplifies.The supplied passage does not include the simplified expression.Source: §4.5 · p. 34

4.6 Calculating the iteration complexity for special cases

  • Theorem 3.6 yields minibatch SAGA iteration complexity for τ-uniform samplings with c1-uniform support, specializing to gradient descent and uniform SAGA under their stated choices.The analysis combines formulas for L1, L2, κ, and ρ; gradient descent uses W=I and S=[n], while uniform SAGA uses W=I and singleton sampling.Source: §4.6 · p. 34
  • For τ-nice sampling, the analysis derives complexity bounds for both W=I and W=Diag(Li), whose relative performance depends on the component constants Li.When some Li is much larger than the others, (104) is smaller than (103); when all Li are equal, (104) is larger.Source: §4.6 · p. 35
  • For τ-partition sampling, the paper derives complexities for W=I and W=Diag(Li), and reports that the latter bound is better.The two choices are listed in lines 12 and 13 of Table 1, respectively.Source: §4.6 · p. 35

4.7 Comparison with previous mini-batch SAGA convergence results

  • Unlike, which samples a gradient independently and discards it for the Jacobian update, JacSketch uses the bias-correcting variable θS to avoid this data-hungry strategy.The independent sampling in maintains independence between Jk and gk.Source: §4.7 · pp. 35–36
  • The [17] complexity decreases only modestly as τ grows, while JacSketch recovers the gradient-descent complexity when τ=n.The cited comparison notes that τ=n corresponds to gradient descent and gives the expected O(L/µ log(1/ϵ)) complexity for the analyzed mini-batch variants.Source: §4.7 · p. 36

5 A Refined Analysis with a Stochastic Lyapunov Function

For τ-partition samplings, the refined analysis uses a stochastic Lyapunov function to obtain sharper results for JacSketch. It establishes the conjectured importance-sampling SAGA rate and derives complexities for several special cases.

5.1 Convergence theorem

  • For τ-partition minibatch sketches, Theorem 5.2 uses a stochastic Lyapunov analysis to obtain the conjectured importance-sampling SAGA rate.The analysis is restricted to samplings that choose uniformly among sets in a partition of [n].Source: §5 · p. 36, §5.1 · p. 37

5.2 Gradient estimate contraction

  • Partition-sampling orthogonality makes a cross-term vanish in the gradient-estimate contraction argument.The authors identify this orthogonality as fundamental to controlling convergence of the gradient estimate.Source: §5 · p. 36, §5.2 · p. 38

5.3 Bounding the second moment of gk

  • The analysis bounds the second moment of the gradient estimate gk.The section states the bound in Lemma 5.4.Source: §5.3 · p. 39

5.4 Smoothness and strong convexity of fIC,J

  • Under per-block strong-convexity and smoothness assumptions, the analysis establishes a smoothness lemma for f_{I_C,J}.The text identifies f_{I_C,J} as n p_C-smooth and uses this property in the proof of Theorem 5.2.Source: §5.4 · p. 39

5.5 Proof of Theorem 5.2

  • The proof bounds the stepsize so term II ≤ I, drops the nonpositive term III, and obtains recursion (113).The combined stepsize bound uses σ_S = n/(4τL_S).Source: §5.5 · pp. 40–41

5.6 Calculating the iteration complexity in special cases

  • The refined complexity specializes across gradient descent and SAGA variants, including importance sampling and reduced-memory SAGA.The probability optimization extends directly to arbitrary τ.Source: §5.6 · pp. 41–42
  • The refined rate is always better than both (131) and (134), resolving the importance-sampling SAGA rate conjecture [32].The non-uniform sampling speedup over uniform sampling is significant when nµ ≤ Lmin.Source: §5.6 · p. 42

6 Experiments

The experiments compare optimized importance sampling and minibatch choices for SAGA, then benchmark optimized SAGA variants against other methods. They report advantages for optimized sampling and minibatching, while noting limits on practical batch-size selection.

6.1 New non-uniform sampling using optimal probabilities

  • As n grows, SAGA-opt is fastest; SAGA-Li stalls on larger problems and can perform worse than uniform SAGA.The comparison uses an artificially constructed ridge-regression problem.Source: Figure 1 · p. 43, §6.1 · p. 43

6.2 Optimal mini-batch size

  • Minibatch SAGA’s iteration complexity initially decreases super-linearly with batch size, then decreases sublinearly once L_Gmax dominates; total complexity improves up to that transition.The analysis identifies the transition between C(τ) and L_Gmax as a useful point for choosing a batch size.Source: §6.2 · p. 44
  • The predicted total-complexity gains are modest, with τ = 2 or τ = 3 optimal in the small-dimensional problems studied; these batches use 10% or 20% of the data.The authors conjecture that larger problems may yield larger improvements.Source: §6.2 · p. 44

6.3 Comparative experiments

  • On three LIBSVM logistic-regression problems, SAG-opt and SAGA-opt were fastest in both epochs and time, ahead of AM-prev, SVRG, and gradient descent.The authors also note SAG-opt performs well despite its biased gradient estimates.Source: Figure 3 · p. 45, §6.3 · p. 46

6.2 Optimal mini-batch size

  • Practical batch-size selection remains difficult because computing L_Gmax has a combinatorial dependence on n and τ; bounding it for linear-classifier losses is left for future work.The authors propose concentration bounds as a possible way to obtain such a bound.Source: §6.2 · pp. 44, 46

7 Conclusion

The conclusion presents JacSketch as a Jacobian-sketching framework for variance reduction and reports new SAGA and minibatch results. It also identifies structured sketches, bias–variance trade-offs, and efficient sparse sketches as future directions.

7.1 Summary of key contributions

  • The analysis resolves Schmidt et al.’s [32] conjecture affirmatively: a properly designed importance-sampling SAGA has iteration complexity O((n + L̄/μ) log…).The supplied conclusion passage truncates the final complexity expression.Source: §7.1 · p. 46
  • The uniform minibatch SAGA complexity interpolates between the best-known SAGA and gradient-descent rates and can guide batch-size choice; the paper also analyzes a reduced-memory SAGA variant..Source: §7.1 · p. 46

7.2 Future work

  • Future work could combine sketches with highly structured weight matrices, such as block-diagonal or sparse matrices, to make Jacobian updates efficient..Source: §7.2 · p. 47
  • The authors propose studying the bias–variance trade-off, noting that they know of no empirical evidence that unbiased estimators outperform biased ones..Source: §7.2 · p. 47
  • New sparse sketches, including Johnson–Lindenstrauss and Achlioptas transforms, could yield methods analyzed by Theorem 3.6, but efficient use requires adapting them so S is row-sparse..Source: §7.2 · p. 47

A Proof of Inequality (20)

The appendix states an inequality under a partition-sampling assumption with equal-sized sets.

A Proof of Inequality (20)

B Duality of Sketch-and-Project and Constrain-and-Approximate

The appendix derives the least-norm update in the constrained formulation and uses duality to establish equivalence with the sketch-and-project formulation.

B Duality of Sketch-and-Project and Constrain-and-Approximate

C Proof of Theorem 4.19

The proof uses a circulant matrix representation and a formula for its eigenvalues, concluding that only two distinct eigenvalues occur.

C Proof of Theorem 4.19

D Notation Glossary

The glossary defines notation for the function vector and its Jacobian, along with sketching, sampling, and convergence quantities.

D Notation Glossary

Related papers

Loading 1805.02632v1…