Variance Reduction With Sparse Gradients

1 INTRODUCTION . Optimization tools for machine learning applications seek to minimize the finite sum objective min x∈Rd f ( x ) , 1 n n∑ i=1 fi ( x ) , ( 1 ) where x is a vector of parameters , and fi : Rd → R is the loss associated with sample i. Batch SGD serves as the prototype for modern stochastic gradient methods . It updates the iterate x with x− η∇fI ( x ) , where η is the learning rate and∇fI ( x ) is the batch stochastic gradient , i.e . ∇fI ( x ) = 1 |I| ∑ i∈I ∇fi ( x ) . The batch size |I| in batch SGD directly impacts the stochastic variance and gradient query complexity of each iteration of the update rule . In recent years , variance reduction techniques have been proposed by carefully blending large and small batch gradients ( e.g . Roux et al. , 2012 ; Johnson & Zhang , 2013 ; Defazio et al. , 2014 ; Xiao & Zhang , 2014 ; Allen-Zhu & Yuan , 2016 ; Allen-Zhu & Hazan , 2016 ; Reddi et al. , 2016a ; b ; Allen-Zhu , 2017 ; Lei & Jordan , 2017 ; Lei et al. , 2017 ; Allen-Zhu , 2018b ; Fang et al. , 2018 ; Zhou et al. , 2018 ; Wang et al. , 2018 ; Pham et al. , 2019 ; Nguyen et al. , 2019 ; Lei & Jordan , 2019 ) . They are alternatives to batch SGD and are provably better than SGD in various settings . While these methods allow for greater learning rates than batch SGD and have appealing theoretical guarantees , they require a per-iteration query complexity which is more than double than that of batch SGD . Defazio ( 2019 ) questions the utility of variance reduction techniques in modern machine learning problems , empirically identifying query complexity as one issue . In this paper , we show that gradient sparsity ( Aji & Heafield , 2017 ) can be used to significantly reduce the query complexity of variance reduction methods . Our work is motivated by the observation that gradients tend to be ” sparse , ” having only a small fraction of large coordinates . Specifically , if the indices of large gradient coordinates ( measured in absolute value ) are known before updating model parameters , we compute the derivative of only those coordinates while setting the remaining gradient coordinates to zero . In principle , if sparsity is exhibited , using large gradient coordinates will not effect performance and will significantly reduce the number of operations required to update model parameters . Nevertheless , this heuristic alone has three issues : ( 1 ) bias is introduced by setting other entries to zero ; ( 2 ) the locations of large coordinates are typically unknown ; ( 3 ) accessing a subset of coordinates may not be easily implemented for some problems like deep neural networks . We provide solutions for all three issues . First , we introduce a new sparse gradient operator : The random-top-k operator . The random-top-k operator is a composition of the randomized coordinate descent operator and the top-k operator . In prior work , the top-k operator has been used to reduce the communication complexity of distributed optimization ( Stich et al. , 2018 ; Aji & Heafield , 2017 ) applications . The random-top-k operator has two phases : Given a stochastic gradient and a pair of integers ( k1 , k2 ) that sum to k , the operator retains k1 coordinates which are most ” promising ” in terms of their ” likelihood ” to be large on average , then randomly selects k2 of the remaining coordinates with appropriate rescaling . The first phase captures sparsity patterns while the second phase eliminates bias . Second , we make use of large batch gradients in variance reduction methods to estimate sparsity patterns . Inspired by the use of a memory vector in Aji & Heafield ( 2017 ) , the algorithm maintains a memory vector initialized with the absolute value of the large batch gradient at the beginning of each outer loop and updated by taking an exponential moving average over subsequent stochastic gradients . Coordinates with large values in the memory vector are more ” promising , ” and the random-top-k operator will pick the top k1 coordinate indices based on the memory vector . Since larger batch gradients have lower variance , the initial estimate is quite accurate . Finally , for software that supports dynamic computation graphs , we provide a cost-effective way ( sparse back-propagation ) to implement the random-top-k operator . In this work we apply the random-top-k operator to SpiderBoost ( Wang et al. , 2018 ) , a recent variance reduction method that achieves optimal query complexity , with a slight modification based on the ” geometrization ” technique introduced by Lei & Jordan ( 2019 ) . Theoretically , we show that our algorithm is never worse than SpiderBoost and can strictly outperform it when the random-top-k operator captures gradient sparsity . Empirically , we demonstrate the improvements in computation for various tasks including image classification , natural language processing , and sparse matrix factorization . The rest of the paper is organized as follows . In Section 2 , we define the random-top-k operator , our optimization algorithm , and a description of sparse backpropagation . The theoretical analyses are presented in Section 3 , followed by experimental results in Section 4 . All technical proofs are relegated to Appendix A , and additional experimental details can be found in Appendix B . 2 STOCHASTIC VARIANCE REDUCTION WITH SPARSE GRADIENTS . Generally , variance reduction methods reduce the variance of stochastic gradients by taking a snapshot ∇f ( y ) of the gradient ∇f ( x ) every m steps of optimization , and use the gradient information in this snapshot to reduce the variance of subsequent smaller batch gradients ∇fI ( x ) ( Johnson & Zhang , 2013 ; Wang et al. , 2018 ) . Methods such as SCSG ( Lei & Jordan , 2017 ) utilize a large batch gradient , which is typically some multiple in size of the small batch gradient b , which is much more practical and is what we do in this paper . To reduce the cost of computing additional gradients , we use sparsity by only computing a subset k of the total gradients d , where y ∈ Rd . For d , k , k1 , k2 ∈ Z+ , let k = k1 + k2 , where 1 ≤ k ≤ d for a parametric model of dimension d. In what follows , we define an operator which takes vectors x , y and outputs y′ , where y′ retains only k of the entries in y , k1 of which are selected according to the coordinates in x which have the k1 largest absolute values , and the remaining k2 entries are randomly selected from y . The k1 coordinate indices and k2 coordinate indices are disjoint . Formally , the operator rtopk1 , k2 : R d → Rd is defined for x , y ∈ Rd as ( rtopk1 , k2 ( x , y ) ) ` = y ` if k1 > 0 and |x| ` ≥ |x| ( k1 ) ( d−k1 ) k2 y ` if ` ∈ S 0 otherwise , where |x| denotes a vector of absolute values , |x| ( 1 ) ≥ |x| ( 2 ) ≥ . . . ≥ |x| ( d ) denotes the order statistics of coordinates of x in absolute values , and S denotes a random subset with size k2 that is uniformly drawn from the set { ` : |x| ` < |x| ( k1 ) } . For instance , if x = ( 11 , 12 , 13 , −14 , −15 ) , y = ( −25 , −24 , 13 , 12 , 11 ) and k1 = k2 = 1 , then S is a singleton uniformly drawn from { 1 , 2 , 3 , 4 } . Suppose S = { 2 } , then rtop1,1 ( x , y ) = ( 0 , 4y2 , 0 , 0 , y5 ) = ( 0 , −96 , 0 , 0 , 11 ) . If k1 + k2 = d , rtopk1 , k2 ( x , y ) = y . On the other hand , if k1 = 0 , rtop0 , k2 ( x , y ) does not depend on x and returns a rescaled random subset of y . This is the operator used in coordinate descent methods . Finally , rtopk1 , k2 ( x , y ) is linear in y . The following Lemma shows that rtopk1 , k2 ( x , y ) is an unbiased estimator of y , which is a crucial property in our later analysis . Lemma 1 . Given any x , y ∈ Rd , E ( rtopk1 , k2 ( x , y ) ) = y , Var ( rtopk1 , k2 ( x , y ) ) = d− k1 − k2 k2 ‖ top−k1 ( x , y ) ‖ 2 , where E is taken over the random subset S involved in the rtopk1 , k2 operator and ( top−k1 ( x , y ) ) ` = { y ` if k1 > 0 and |x| ` < |x| ( k1 ) 0 otherwise . Our algorithm is detailed as below . Algorithm 1 : SpiderBoost with Sparse Gradients . Input : Learning rate η , inner loop size m , outer loop size T , large batch size B , small batch size b , initial iterate x0 , memory decay factor α , sparsity parameters k1 , k2 . 1 I0 ∼ Unif ( { 1 , . . . , n } ) with |I0| = B 2 M0 : = |∇fI0 ( x0 ) | 3 for j = 1 , ... , T do 4 x ( j ) 0 : = xj−1 , M ( j ) 0 : = Mj−1 5 Ij ∼ Unif ( { 1 , . . . , n } ) with |Ij | = B 6 ν ( j ) 0 : = ∇fIj ( x ( j ) 0 ) 7 Nj : = m ( for implementation ) or Nj ∼ geometric distribution with mean m ( for theory ) 8 for t = 0 , . . . , Nj − 1 do 9 x ( j ) t+1 : = x ( j ) t − ην ( j ) t 10 I ( j ) t ∼ Unif ( [ n ] ) with |I ( j ) t | = b 11 ν ( j ) t+1 : = ν ( j ) t + rtopk1 , k2 ( M ( j ) t , ∇fI ( j ) t ( x ( j ) t+1 ) −∇fI ( j ) t ( x ( j ) t ) ) 12 M ( j ) t+1 : = α|ν ( j ) t+1|+ ( 1− α ) M ( j ) t 13 xj : = x ( j ) Nj , Mj : = M ( j ) Nj Output : xout = xT ( for implementation ) or xout = xT ′ where T ′ ∼ Unif ( [ T ] ) ( for theory ) The algorithm includes an outer-loop and an inner-loop . In the theoretical analysis , we generate Nj as Geometric random variables . This trick is called ” geometrization ” , proposed by Lei & Jordan ( 2017 ) and dubbed by Lei & Jordan ( 2019 ) . It greatly simplifies analysis ( e.g . Lei et al. , 2017 ; Allen-Zhu , 2018a ) . In practice , as observed by Lei et al . ( 2017 ) , setting Nj to m does not impact performance in any significant way . We only use ” geometrization ” in our theoretical analysis for clarity . Similarly , for our theoretical analysis , the output of our algorithm is selected uniformly at random from the set of outer loop iterations . Like the use of average iterates in convex optimization , this is a common technique for nonconvex optimization proposed by Nemirovski et al . ( 2009 ) . In practice , we simply use the last iterate . Similar to Aji & Heafield ( 2017 ) , we maintain a memory vector M ( j ) t at each iteration of our algorithm . The memory vector is initialized to the large batch gradient computed before every pass through the inner loop , which provides a relatively accurate gradient sparsity estimate of x ( j ) 0 . The exponential moving average gradually incorporates information from subsequent small batch gradients to account for changes to gradient sparsity . We then use M ( j ) t as an approximation to the variance of each gradient coordinate in our rtopk1 , k2 operator . With M ( j ) t as input , the rtopk1 , k2 operator targets k1 high variance gradient coordinates in addition to k2 randomly selected coordinates . The cost of invoking rtopk1 , k2 is dominated by the algorithm for selecting the top k coordinates , which has linear worst case complexity when using the introselect algorithm ( Musser , 1997 ) .

Summary: This paper aims at improving the computational cost of variance reduction methods while preserving their benefits regarding the fast provable convergence. The existing variance reduction based methods suffer from higher per-iteration gradient query complexity as compared to the vanilla mini-batch SGD, which limits their utility in many practical settings. This paper notices that, for many models, as the training progresses the gradient vectors start exhibiting structure in the sense that only a small number of coordinates have large magnitude. Based on this observation, the paper proposes a modified variance reduction method (by modifying the SpiderBoost method), where a 'memory vector' keeps track of the coordinates of the gradient vectors with large variance. Let $d$ be the size of the model parameter. During each iteration, one computes the gradient for $k_1$ coordinates with the highest variance (according to the memory vector) and an additional $k_2$ random coordinates.