Adaptive Single-Pass Stochastic Gradient Descent in Input Sparsity Time

We study sampling algorithms for variance reduction methods for stochastic optimization . Although stochastic gradient descent ( SGD ) is widely used for large scale machine learning , it sometimes experiences slow convergence rates due to the high variance from uniform sampling . In this paper , we introduce an algorithm that approximately samples a gradient from the optimal distribution for a common finite-sum form with n terms , while just making a single pass over the data , using input sparsity time , and Õ ( Td ) space . Our algorithm can be implemented in big data models such as the streaming and distributed models . Moreover , we show that our algorithm can be generalized to approximately sample Hessians and thus provides variance reduction for second-order methods as well . We demonstrate the efficiency of our algorithm on large-scale datasets . 1 INTRODUCTION . There has recently been tremendous progress in variance reduction methods for stochastic gradient descent ( SGD ) methods for the standard convex finite-sum form optimization problem min x∈Rd F ( x ) : = 1 n ∑n i=1 fi ( x ) , where f1 , . . . , fn : Rd → R is a set of convex functions that commonly represent loss functions . Whereas gradient descent ( GD ) performs the update rule xt+1 = xt − ηt∇F ( xt ) on the iterative solution xt at iterations t = 1 , 2 , . . . , SGD ( Robbins & Monro , 1951 ; Nemirovsky & Yudin , 1983 ; Nemirovski et al. , 2009 ) picks it ∈ [ n ] in iteration t with probability pit and performs the update rule xt+1 = xt − ηtnpit∇fit ( xt ) , where ∇fit is the gradient ( or a subgradient ) of fit and ηt is some predetermined learning rate . Effectively , training example it is sampled with probability pit and the model parameters are updated using the selected example . The SGD update rule only requires the computation of a single gradient at each iteration and provides an unbiased estimator to the full gradient , compared to GD , which evaluates n gradients at each iteration and is prohibitively expensive for large n. However , since SGD is often performed with uniform sampling so that the probability pi , t of choosing index i ∈ [ n ] at iteration t is pi , t = 1n at all times , the variance introduced by the randomness of sampling a specific vector function can be a bottleneck for the convergence rate of the iterative process . Thus the subject of variance reduction beyond uniform sampling has been well-studied in recent years ( Roux et al. , 2012 ; Johnson & Zhang , 2013 ; Defazio et al. , 2014 ; Reddi et al. , 2015 ; Zhao & Zhang , 2015 ; Daneshmand et al. , 2016 ; Needell et al. , 2016 ; Stich et al. , 2017 ; Johnson & Guestrin , 2018 ; Katharopoulos & Fleuret , 2018 ; Salehi et al. , 2018 ; Qian et al. , 2019 ) . A common technique to reduce variance is importance sampling , where the probabilities pi , t are chosen so that vector functions with larger gradients are more likely to be sampled . Thus for Var ( v ) : = E [ ‖v‖22 ] − ‖E [ v ] ‖22 , for a random vector v , then pi , t = 1 n for uniform sampling implies σ2t = Var ( 1 npit , t ∇fit ) = 1 n2 ( n n∑ i=1 ‖∇fi ( xt ) ‖2 − n2 · ‖∇F ( xt ) ‖2 ) , whereas importance sampling with pi , t = ‖∇fi ( xt ) ‖∑n j=1‖∇fj ( xt ) ‖ gives σ2t = Var ( 1 npit , t ∇fit ) = 1 n2 ( n∑ i=1 ‖∇fi ( xt ) ‖ ) 2 − n2 · ‖∇F ( xt ) ‖2 , which is at most 1n2 ( n ∑ ‖∇fi ( xt ) ‖2 − n2 · ‖∇F ( xt ) ‖2 ) , by the Root-Mean Square-Arithmetic Mean Inequality , and can be significantly less . Hence the variance at each step is reduced , possibly substantially , e.g. , Example 1.3 and Example 1.4 , by performing importance sampling instead of uniform sampling . In fact , it follows from the Cauchy-Schwarz inequality that the above importance sampling probability distribution is the optimal distribution for variance reduction . However , computing the probability distribution for importance sampling requires computing the gradients in each round , which is too expensive in the first place . Second-Order Methods . Although first-order methods such as SGD are widely used , they do sometimes have issues such as sensitivity to the choice of hyperparameters , stagnation at high training errors , and difficulty in escaping saddle points . By considering second-order information such as curvature , second-order optimization methods are known to be robust to several of these issues , such as ill-conditioning . For example , Newton ’ s method can achieve a locally super-linear convergence rate under certain conditions , independent of the problem . Although naı̈ve second-order methods are generally too slow compared to common first-order methods , stochastic Newton-type methods such as Gauss-Newton have shown to be scalable in the scientific computing community ( Roosta-Khorasani et al. , 2014 ; Roosta-Khorasani & Mahoney , 2016a ; b ; Xu et al. , 2019 ; 2020 ) . Our Contributions . We give a time efficient algorithm that provably approximates the optimal importance sampling using a small space data structure . Remarkably , our data structure can be implemented in big data models such as the streaming model , which just takes a single pass over the data , and the distributed model , which requires just a single round of communication between parties holding each loss function . For ∇F = 1n ∑ ∇fi ( x ) , where each ∇fi = f ( 〈ai , x〉 ) · ai for some polynomial f and vector ai ∈ Rd , let nnz ( A ) be the number of nonzero entries of A : = a1◦ . . .◦an1 . Thus for T iterations , where d T n , GD has runtime Õ ( T · nnz ( A ) ) while our algorithm has runtime T · poly ( d , log n ) + Õ ( nnz ( A ) ) , where we use Õ ( · ) to suppress polylogarithmic terms . Theorem 1.1 Let ∇F = 1n ∑ ∇fi ( x ) , where each ∇fi = f ( 〈ai , x〉 ) · ai for some polynomial f and vector ai ∈ Rd and let nnz ( A ) be the number of nonzero entries of A : = a1 ◦ . . . ◦ an . For d T n , there exists an algorithm that performs T steps of SGD and at each step samples a gradient within a constant factor of the optimal probability distribution . The algorithm requires a single pass over A and uses Õ ( nnz ( A ) ) pre-processing time and Õ ( Td ) space . Theorem 1.1 can be used to immediately obtain improved convergence guarantees for a class of functions whose convergence rate depends on the variance σ2t , such as µ-smooth functions or strongly convex functions . Recall that SGD offers the following convergence guarantees for smooth functions : Theorem 1.2 ( Nemirovski et al. , 2009 ; Meka , 2017 ) Let F be a µ-smooth convex function and xopt = argminF ( x ) . Let σ2 be an upper bound for the variance of the unbiased estimator across all iterations and xk = x1+ ... +xkk . Let each step-size ηt be η ≤ 1 µ . Then for SGD with initial position x0 , E [ F ( xk ) − F ( xopt ) ] ≤ 1 2ηk ‖x0 − xopt‖22 + ησ2 2 , so that k = O ( 1 2 ( σ2 + µ ‖x0 − xopt‖22 ) 2 ) iterations suffices to obtain an -approximate optimal value by setting η = 1√ k . In the convergence guarantees of Theorem 1.2 , we obtain a constant factor approximation to the variance σ = σopt from optimal importance sampling , which can be significantly better than the 1We use the notation a ◦ b to denote the vertical concatenation [ a b ] . variance σ = σuniform from uniform sampling in standard SGD . We first show straightforward examples where uniform sampling an index performs significantly worse than importance sampling . For example , if∇fi ( x ) = 〈ai , x〉 · ai , then for A = a1 ◦ . . . ◦ an : Example 1.3 When the nonzero entries of the input A are concentrated in a small number of vectors ai , uniform sampling will frequently sample gradients that are small and make little progress , whereas importance sampling will rarely do so . In an extreme case , the A can contain exactly one nonzero vector ai and importance sampling will always output the full gradient whereas uniform sampling will only find the nonzero row with probability 1n . Example 1.4 It may be that all rows of A have large magnitude , but x is nearly orthogonal to most of the rows of A and heavily in the direction of row ar . Then 〈ai , x〉 · ai is small in magnitude for most i , but 〈ar , x〉 · ar is large so uniform sampling will often output small gradients while importance sampling will output 〈ar , x〉 · ar with high probability . Thus Example 1.3 shows that naı̈ve SGD with uniform sampling can suffer up to a multiplicative n factor loss in the convergence rate of Theorem 1.2 compared to that of SGD with importance sampling whereas Example 1.4 shows a possible additive n factor loss . Unlike a number of previous variance reduction methods , we do not require distributional assumptions ( Bouchard et al. , 2015 ; Frostig et al. , 2015 ; Gopal , 2016 ; Jothimurugesan et al. , 2018 ) or offline access to the data ( Roux et al. , 2012 ; Johnson & Zhang , 2013 ; Defazio et al. , 2014 ; Reddi et al. , 2015 ; Zhao & Zhang , 2015 ; Daneshmand et al. , 2016 ; Needell et al. , 2016 ; Stich et al. , 2017 ; Johnson & Guestrin , 2018 ; Katharopoulos & Fleuret , 2018 ; Salehi et al. , 2018 ; Qian et al. , 2019 ) . On the other hand , for applications such as neural nets in which the parameters in the loss function can change , we can use a second-order approximation for a number of iterations , then reread the data to build a new second-order approximation when necessary . We complement our main theoretical result with empirical evaluations comparing our algorithm to SGD with uniform sampling for logistic regression on the a9a Adult dataset collected by UCI and retrieved from LibSVM ( Chang & Lin , 2011 ) . Our evaluations demonstrate that for various step-sizes , our algorithm has significantly better performance than uniform sampling across both the number of SGD iterations and surprisingly , wall-clock time . We then show that our same framework can also be reworked to approximate importance sampling for the Hessian , thereby performing variance reduction for second-order optimization methods . ( Xu et al. , 2016 ) reduce the bottleneck of many second-order optimization methods to the task of sampling s rows of A = a1 ◦ . . . ◦ an so that a row ai is sampled with probability ‖f ( 〈ai , x〉 ) ·a > i ai‖2F∑n i=1‖f ( 〈ai , x〉 ) ·a > i ai‖2F , for some fixed function f so that the Hessian H has the form H : = ∇2F = 1n ∑ ∇f ( 〈ai , x〉 ) a > i ai . ( Xu et al. , 2016 ) show that this finite-sum form arises frequently in machine learning problems such as logistic regression with least squares loss . Theorem 1.5 Let∇2F = 1n ∑ ∇fi ( x ) , where each∇fi = f ( 〈ai , x〉 ) · a > i ai for some polynomial f and vector ai ∈ Rd and let nnz ( A ) be the number of nonzero entries of A : = a1 ◦ . . . ◦ an . For d T n , there exists an algorithm that subsamples T Hessians within a constant factor of the optimal probability distribution . The algorithm requires a single pass over A and uses Õ ( nnz ( A ) ) pre-processing time and Õ ( Td ) space .

Summary: This paper develops an efficient streaming algorithm to approximate the optimal importance sampling weights for variance reduction in finite-sum SGD. The optimal weights are proportional to each sample's gradient norm; this work uses AMS-like moment estimation to sketch gradient norms which take the form of a bounded-degree polynomial, in time linear in the input sparsity and polynomial in the dimension d, iteration count T, and the log of the number of the samples n. A second-order analogue is derived for approximating optimal importance weights for sampling the Hessian. Some experiments are shown with more simplistic importance weight estimators (not the proposed algorithm), to demonstrate the advantage over uniform sampling.