Stochastic Reweighted Gradient Descent

Importance sampling is a promising strategy for improving the convergence rate of stochastic gradient methods . It is typically used to precondition the optimization problem , but it can also be used to reduce the variance of the gradient estimator . Unfortunately , this latter point of view has yet to lead to practical methods that provably improve the asymptotic error of stochastic gradient methods . In this work , we propose stochastic reweighted gradient ( SRG ) , a variance-reduced stochastic gradient method based solely on importance sampling that can improve on the asymptotic error of stochastic gradient descent ( SGD ) in the strongly convex and smooth case . We show that SRG can be extended to combine the benefits of both importance-sampling-based preconditioning and variance reduction . When compared to SGD , the resulting algorithm can simultaneously reduce the condition number and the asymptotic error , both by up to a factor equal to the number of component functions . We demonstrate improved convergence in practice on ` 2-regularized logistic regression problems . 1 INTRODUCTION . Unconstrained optimization of finite-sum objectives is a core algorithmic problem in machine learning . The prototypical way of solving such problems is by viewing them through the lens of stochastic optimization , where the source of stochasticity resides in the choice of the index in the sum . Stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) remains the standard algorithm for this class of problems . A natural way to improve on SGD is by considering importance sampling schemes . This idea is not new and dates back to ( Needell et al. , 2014 ) who uses importance sampling as a preconditioning technique . They propose sampling the indices with probabilities proportional to the smoothness constants of the corresponding component functions , and show that this sampling scheme provably reduces the condition number of the problem . In another line of work , variance-reduced methods were found to achieve linear convergence in the strongly convex and smooth case ( Roux et al. , 2012 ; Schmidt et al. , 2017 ) . Many of these methods rely on control variates to reduce the variance of the gradient estimator used by SGD ( Johnson & Zhang , 2013 ; Defazio et al. , 2014 ) . While very successful , the applicability of these methods is limited by the large memory overhead that they introduce , or the periodic full-gradient recomputation that they require ( Defazio & Bottou , 2019 ) . Despite strong progress in this research area , an importance-sampling-based analogue to these algorithms , which is potentially free from these drawbacks , has yet to emerge . In this work , we propose such an analogue . We introduce stochastic reweighted gradient ( SRG ) , an importance-sampling-based variance-reduced optimization algorithm . Similar to SGD , SRG requires a single gradient oracle call per iteration , and only requires O ( n ) additional memory , and O ( log n ) additional floating point operations per iteration . We analyze the convergence rate of SRG in the strongly-convex and smooth case , and show that it can provably improve the asymptotic error of SGD . Finally , we show how our importance sampling strategy can be combined with smoothnessbased importance sampling , and prove that the resulting algorithm simultaneously performs variance reduction and preconditioning . We demonstrate improved convergence in practice on ` 2-regularized logistic regression problems . 2 PRELIMINARIES . We consider the finite-sum optimization problem : min x∈Rd { F ( x ) : = 1 n n∑ i=1 fi ( x ) } ( 1 ) where F is µ-strongly convex for µ > 0 , and for all i ∈ [ n ] , fi is convex and Li-smooth for Li > 0 . Note that by strong-convexity , F has a unique minimizer x∗ ∈ Rd . We define the maximum Lmax : = maxi∈ [ n ] Li and average L : = ∑n j=1 Li/n smoothness constants . Similarly , we define the maximum κmax : = Lmax/µ and average κ : = L/µ condition numbers . The classical way of solving ( 1 ) is by viewing it as a stochastic optimization problem where the randomness comes from the choice of the index i ∈ [ n ] . Starting from some x0 ∈ Rd , and for an iteration number k ∈ N , stochastic gradient descent ( SGD ) performs the following update : xk+1 = xk − αk∇fik ( xk ) for a step size αk > 0 and a random index ik drawn uniformly from [ n ] . The idea behind importance sampling for SGD is to instead sample the index ik according to a chosen distribution pk on [ n ] , and to perform the update ( Needell et al. , 2014 ) : xk+1 = xk − αk 1 npikk ∇fik ( xk ) ( 2 ) where pik is the i th component of the probability vector pk . It is immediate to verify that the importance sampling estimator of the gradient is unbiased as long as pk > 0 . The question that we address in this paper is how to design a sequence { pk } ∞k=0 that produces more efficient gradient estimators than the ones produced by uniform sampling . One way to design such a sequence is by adopting a greedy strategy : at each iteration k we choose pk to minimize the conditional variance of the gradient estimator , which is given by , up to an additive constant : σ2 ( xk , p ) : = E ik∼p [ ∥∥∥∥ 1npik∇fik ( xk ) ∥∥∥∥2 2 ∣∣∣∣∣ ( it ) k−1t=0 ] = 1 n2 n∑ i=1 1 pi ‖∇fi ( xk ) ‖22 , ( 3 ) This conditional variance is minimized at ( Zhao & Zhang , 2015 ) : arg min p∈∆ σ2 ( xk , p ) = ( ‖∇fi ( xk ) ‖2∑n j=1 ‖∇fj ( xk ) ‖2 ) n i=1 ( 4 ) Ideally , we would like to set pk to this minimizer . However , this requires knowledge of the gradient norms ( ‖∇fi ( xk ) ‖2 ) n i=1 , which , in general , requires n gradient evaluations per iteration . 3 ALGORITHM . In this section , we show how to design a tractable sequence of importance sampling distributions for SGD that approximate the conditional-variance-minimizing distributions ( 4 ) . First we construct efficient approximations of the conditional variances ( 3 ) . We then state a simple bound on the approximation errors and use it to motivate our choice of importance sampling distributions . To approximate the conditional variances , we follow the strategy of certain variance-reduced methods ( Roux et al. , 2012 ; Schmidt et al. , 2017 ; Defazio et al. , 2014 ) . These methods maintain a table ( gik ) n i=1 that tracks the component gradients ( ∇fi ( xk ) ) ni=1 and updates itself each iteration at the index ik used to update the iterates xk ( 2 ) . Our method instead maintains an array of gradient norms , from which we construct an approximation of the conditional variance ( 3 ) of the gradient estimator : σ̃2 ( xk , p ) : = 1 n2 n∑ i=1 1 pi ∥∥gik∥∥22 Algorithm 1 SRG 1 : Parameters : step sizes ( αk ) ∞k=0 > 0 , mixture coefficients ( θk ) ∞k=0 ∈ ( 0 , 1 ] 2 : Initialization : x0 ∈ Rd , ( ∥∥gi0∥∥2 ) ni=1 ∈ Rn 3 : for k = 0 , 1 , 2 , . . . do 4 : pk = ( 1− θk ) qk + θk/n { qk is defined in ( 5 ) } 5 : bk ∼ Bernoulli ( θk ) 6 : if bk = 1 then ik ∼ 1/n else ik ∼ qk 7 : xk+1 = xk − αk 1 np ik k ∇fik ( xk ) 8 : ∥∥gik+1∥∥2 = { ‖∇fi ( xk ) ‖2 if bk = 1 and ik = i∥∥gik∥∥2 otherwise 9 : end for which is minimized at : qk = arg min p∈∆ σ̃2 ( xk , p ) = ∥∥gik∥∥2∑n j=1 ∥∥∥gjk∥∥∥ 2 n i=1 ( 5 ) We do not directly use qk as an importance sampling distribution , because this approximation may be poor . In particular , we have the following bound on the conditional variance ( 3 ) : σ2 ( xk , p ) ≤ 2 n2 n∑ i=1 1 pi ∥∥∇fi ( xk ) − gik∥∥22 + 2σ̃2 ( xk , p ) ( 6 ) Recall that our goal is to pick pk that minimizes σ2 ( xk , p ) . qk minimizes the second term on the right-hand side , but we must ensure that both terms are small . Two conditions are needed to keep the first term small . The first is to control the terms ∥∥∇fi ( xk ) − gik∥∥22 , which we can do by making sure that the historical gradients gik are frequently updated . The second is to ensure that the probabilities pik are lower bounded . We achieve both of these properties , as well as approximately minimize the right-hand side of ( 6 ) by mixing qk with the uniform distribution over [ n ] . This yields our final importance sampling distribution for a given mixture coefficient θk ∈ ( 0 , 1 ] : pk = ( 1− θk ) qk + θk n ( 7 ) Our analysis clarifies the role of the sequence of mixture coefficients ( θk ) ∞k=0 , and relates it to both the step size sequence ( αk ) ∞k=0 and the asymptotic error of SRG in the constant step size setting . Curiously , our analysis requires performing the update of the array ( ∥∥gik∥∥2 ) ni=1 only when the index ik is drawn from the uniform mixture component . It is not clear to us whether this constraint is an artifact of the analysis or a property of the algorithm . We discuss this further after the statement of Lemma 1 in section 4 . The pseudocode for our method is given in Algorithm 1 . We briefly discuss how SRG can be efficiently implemented . To sample from qk , we store ( ∥∥gik∥∥2 ) ni=1 in a binary indexed tree using O ( n ) memory , along with the normalizing constant λk = ∑n i=1 ∥∥gik∥∥2 . The binary indexed tree can be updated and maintained in O ( log n ) operations , while the normalizing constant can be maintained in constant time . We can then sample from qk using inverse transform sampling : we multiply a uniform random variable u ∈ [ 0 , 1 ) by λk , then return the largest index whose corresponding prefix sum is less than λku . This procedure also requires only O ( log n ) operations . The total overhead of SRG when compared to SGD is therefore O ( n ) memory and O ( log n ) floating point operations per iteration . 4 THEORY . In this section , we analyze the convergence rate of SRG , and show that it can achieve a better asymptotic error than SGD . Two key constants are helpful in contrasting the asymptotic errors of SRG and SGD . Recall the definition of σ2 ( xk , pk ) in ( 3 ) , and define : σ2 : = σ2 ( x∗ , 1/n ) = 1 n n∑ i=1 ‖∇fi ( x∗ ) ‖22 σ2∗ : = min p∈∆ σ2 ( x∗ , p ) = 1 n2 ( n∑ i=1 ‖∇fi ( x∗ ) ‖2 ) 2 It is well known that the asymptotic error of SGD depends linearly on σ2 ( Needell et al. , 2014 ) . We here show that SRG reduces this to a linear dependence on σ2∗ , which can be up to n times smaller . To study the convergence rate of SRG , we use the following Lyapunov function , which is similar to the one used to study the convergence rate of SAGA ( Hofmann et al. , 2015 ; Defazio , 2016 ) : T k = T ( xk , ( g i k ) n i=1 ) : = αk θk a Lmax n∑ i=1 ∥∥gik −∇fi ( x∗ ) ∥∥22 + ‖xk − x∗‖22 ( 8 ) for a constant a > 0 that we set during the analysis . The proofs of this section are in Appendix B .

Summary: This paper considers minimizing the finite sum problem with SGD with non-uniform sampling.  The optimal sampling weight at a given point is achieved via minimizing the variance of stochastic gradient which needs the norm of the gradient of each function at that point. To approximate these sampling weights and also be efficient this paper proposes to take the approach similar to SAGA and use auxiliary memory with the size of number of functions and update one of the stored values in each iterate similar to SAGA. Using stale and stored gradient norms gives an upper bound to the stochastic gradient variance that could be minimized similarly. Since this new upper bound contains two terms, they propose a mixture of non-uniform and uniform sampling to make both terms small. They analyze SGD with the proposed sampling scheme and show that the sub-optimality upper bound is n (the number of functions)-times smaller than the upper bound of SGD with uniform sampling. They also propose another weighting scheme that mixes uniform and non-uniformes schemes proportioned with gradient norms and smoothness parameters of each function.