Stochastic quasi-gradient methods: variance reduction via Jacobian sketching

By Robert Mansel Gower, Peter Richtárik, Francis R. Bach. 2018

115 citations recorded by the source corpus.

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. Our method can also be seen as stochastic gradient descent applied to a controlled stochastic optimization reformulation of the original problem, where the control comes from the Jacobian estimates. 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, featuring a novel proof technique based on a stochastic Lyapunov function. 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 celebrated stochastic average gradient (SAGA) method, and its several 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. (Proceedings of the eighteenth international conference on artificial intelligence and statistics, AISTATS 2015, San Diego, California, 2015). The rates we obtain for minibatch SAGA are also superior to existing rates and are sufficiently tight as to show a decrease in total complexity as the minibatch size increases. Moreover, we obtain the first minibatch SAGA method with importance sampling.

Subfields: Artificial Intelligence, Computational Mechanics, Statistics and Probability, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

In Statistics and Probability

In Optimization & Training

Citations in the corpus

Stochastic Hamiltonian Gradient Methods for Smooth Games (2020, methodology): “For more examples of variance reduced methods in different settings, see Defazio (2016); Kone ˇ cn ´ y et al. (2016); Gower et al. (2018); Sebbouh et al. (2019).”
SGD: General Analysis and Improved Rates (2019, methodology): “Expected smoothness (Gower et al., 2018) is an assumption that combines both the properties of the distribution D and the smoothness properties of function f .”
Stochastic Distributed Learning with Gradient Quantization and Variance Reduction (2019, methodology): “…Defazio et al., 2014; Qu et al., 2015; Shalev-Shwartz, 2016; Qu et al., 2016; Csiba & Richt ´ arik, 2015; Csiba et al., 2015; Nguyen et al., 2017; Gower et al., 2018; Zhou, 2018) admit linear convergence on empirical risk minimization problems, surpassing the rate of vanilla SGD (Bottou, 2010).”
Federated Learning of a Mixture of Global and Local Models (2020, background): “We ﬁst show that our gradient estimator G ( x ) satisﬁes the expected smoothness property (Gower et al., 2018, 2019).”
One Method to Rule Them All: Variance Reduction for Data, Parameters and Many New Methods (2019, methodology): “XR XRE [R]−1 18 JS JacSketch [8] with AS and prox N”
A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent (2019, methodology): “It was developed using the stochastic reformulation approach (whence the “SR”) pioneered in (Richtárik and Takáč, 2017) in a numerical linear algebra setting, and later extended to develop the JacSketch variance-reduction technique for finite-sum optimization (Gower et al., 2018).”
Randomized Iterative Methods for Linear Systems: Momentum, Inexactness and Gossip (2019, background): “Similar stochastic reformulations that have the sketch and project method as special case were also proposed in the more general setting of convex optimization [68] and in the context of variance-reduced methods [70].”
Unified Analysis of Stochastic Gradient Methods for Composite Convex and Smooth Optimization (2020, methodology): “The appendix is organized as follows: • Section A : we present the arbitrary sampling framework for Stochastic Gradient methods introduced in [13], which will be used for the analysis of SGD and L-SVRG .”
99% of Parallel Optimization is Inevitably a Waste of Time (2019, background): “Unfortunately, numerical experiments show that this is not a proof flaw, but rather a real property of these methods (Gower et al., 2018).”
The Error-Feedback Framework: Better Rates for SGD with Delayed Gradients and Compressed Communication (2019, background): “Closely related is the refined notion of expected smoothness, see the discussions in (Gower et al., 2018, 2019).”
Unified Optimal Analysis of the (Stochastic) Gradient Method (2019, background): “, 2016) or expected smoothness conditions (Gower et al., 2018).”
SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation (2020, methodology): “Instead, we rely on the recently introduced expected residual (ER) condition (Gower et al., 2020).”
Nonconvex Variance Reduced Optimization with Arbitrary Sampling (2018, methodology): “For instance, an importance sampling variant of the popular SAGA method, with the “correct” convergence rate, was only designed very recently [10]; and the analysis applies to strongly convex f only.”
Better Theory for SGD in the Nonconvex World (2020, methodology): “…contributions of our work are: • Inspired by recent developments in the analysis of SGD in the (strongly) convex setting (Richt ´ arik & Tak ´ a ˇ c, 2017; Gower et al., 2018; 2019), we propose a new assumption, which we call expected smoothness (ES) , for modelling the second moment of the…”
Dualize, Split, Randomize: Fast Nonsmooth Optimization Algorithms (2020, background): “Assumption 1 is satisfied by several stochastic gradient estimators used in machine learning, including the full gradient, some kinds of coordinate descent [25], variance reduction [20, 24, 27, 31], and also compressed gradients used to reduce the communication cost in distributed optimization [28],”
Stochastic Reformulations of Linear Systems: Algorithms and Convergence Theory (2017, methodology): “eduction. A major open problem in the variancereduction literature for ﬁnite-sum optimization problems is shedding light on the mechanism behind variance reduction. Motivated by our work, Gower et el [24] develop a controlled stochastic reformulation of ﬁnite sum problems, which enables them to desi”
Towards closing the gap between the theory and practice of SVRG (2019, methodology): “new tighter mini-batch complexity bounds was by using the recently introduced concept of expected smoothness [9] alongside a new constant we introduce in this paper called the expected residual constant.”
Don't Jump Through Hoops and Remove Those Loops: SVRG and Katyusha are Better Without the Outer Loop (2019, unknown): “, 2016), RBCN (Doikov & Richtárik, 2018), JacSketch (Gower et al., 2018) and SAGD (Bibi et al.”