Keep the Gradients Flowing: Using Gradient Flow to study Sparse Network Optimization

Training sparse networks to converge to the same performance as dense neural architectures has proven to be elusive . Recent work suggests that initialization is the key . However , while this direction of research has had some success , focusing on initialization alone appears to be inadequate . In this paper , we take a broader view of training sparse networks and consider various choices made during training that might disadvantage sparse networks . We measure the gradient flow across different networks and datasets , and show that the default choices of optimizers , activation functions and regularizers used for dense networks can disadvantage sparse networks . Based upon these findings , we show that gradient flow in sparse networks can be improved by reconsidering aspects of the architecture design and the training regime . Our work suggests that initialization is only one piece of the puzzle and a wider view of tailoring optimization to sparse networks yields promising results . 1 INTRODUCTION . Over the last decade , a “ bigger is better ” race in the number of model parameters has gripped the field of machine learning ( Amodei et al. , 2018 ; Thompson et al. , 2020 ) , primarily driven by overparameterized deep neural networks ( DNNs ) . Additional parameters improve top-line metrics , but drive up the cost of training ( Horowitz , 2014 ; Strubell et al. , 2019 ; Hooker , 2020 ) and increase the latency and memory footprint at inference time ( Warden & Situnayake , 2019 ; Samala et al. , 2018 ; Lane & Warden , 2018 ) . Moreover , overparameterized networks have been shown to be more prone to memorization ( Zhang et al. , 2016 ) . To address some of these limitations , there has been a renewed focus on compression techniques that preserve top-line performance while improving efficiency . A large amount of research focus has centered on pruning , where weights estimated to be unnecessary are removed from the network at the end of training ( Louizos et al. , 2017 ; Wen et al. , 2016 ; Cun et al. , 1990 ; Hassibi et al. , 1993a ; Ström , 1997 ; Hassibi et al. , 1993b ; Zhu & Gupta , 2017 ; See et al. , 2016 ; Narang et al. , 2017 ) . Pruning has shown a remarkable ability to preserve top-line metrics of performance , even when removing the majority of weights ( Hooker et al. , 2019 ; Gale et al. , 2019 ) . However , most pruning techniques still require training a large , overparameterized model before pruning a subset of weights . Due to the drawbacks of starting dense prior to introducing sparsity , there has been a recent focus on methods that allow networks which start sparse at initialization , to converge to similar performance as dense networks ( Frankle & Carbin , 2018 ; Frankle et al. , 2019b ; Liu et al. , 2018a ) . These efforts have focused disproportionately on trying to understand the properties of initial sparse weight distributions that allow for convergence . However , while this work has had some success , focusing on initialization alone has proven to be inadequate ( Frankle et al. , 2020 ; Evci et al. , 2019 ) . In this work , we take a broader view of why training sparse networks to converge to the same performance as dense networks has proven to be elusive . We reconsider many of the basic building blocks of the training process and ask whether they disadvantage sparse networks or not . Our work focuses on the behaviour of networks with random , fixed sparsity at initialization and we aim to gain further intuition into how these networks learn . Furthermore , we provide tooling tailored to the analysis of these networks . In order to effectively study sparse network optimization in a controlled environment , we propose an experimental framework , Same Capacity Sparse vs Dense Comparison ( SC-SDC ) . Contrary to most prior work comparing sparse to dense networks , where overparameterized dense networks are compared to smaller sparse networks , SC-SDC compares sparse networks to their equivalent capacity dense networks ( same number of active connections and depth ) . This ensures that the results are a direct result of sparse connections themselves and not due to having more or fewer weights ( as is the case when comparing large , dense networks to smaller , sparse networks ) . We go beyond simply comparing top-line metrics by also measuring the impact on gradient flow of each intervention . Historically , exploding and vanishing gradients were a common problem in neural networks ( Hochreiter et al. , 2001 ; Hochreiter , 1991 ; Bengio et al. , 1994 ; Glorot & Bengio , 2010 ; Goodfellow et al. , 2016 ) . Recent work has suggested that poor gradient flow is an exacerbated issue in sparse networks ( Wang et al. , 2020 ; Evci et al. , 2020 ) . To accurately measure gradient flow in sparse networks , we propose a normalized measure of gradient flow , which we term Effective Gradient Flow ( EGF ) – this measure normalizes by the number of active weights and thus is better suited to studying the training dynamics of sparse networks . We use this measure in conjunction with SC-SDC , to see where sparse optimization fails and to consider where this failure could be a result of poor gradient flow . Contributions Our contributions can be enumerated as follows : 1 . Measuring effective gradient flow We conduct large scale experiments to evaluate the role of regularization , optimization and architecture choices on sparse models . We evaluate multiple datasets and architectures and propose a new measure of gradient flow , Effective Gradient Flow ( EGF ) , that we show to be a stronger predictor of top-line metrics such as accuracy and loss than current gradient flow formulations . 2 . Batch normalization plays a disproportionate role in stabilizing sparse networks We show that batch normalization is more important for sparse networks than it is for dense networks , which suggests that gradient instability is a key obstacle to starting sparse . 3 . Not all optimizers and regulizers are created equal Weight decay and data augmentation can hurt sparse network optimization , particularly when used in conjunction with accelerating , adaptive optimization methods that use an exponentially decaying average of past squared gradients , such as Adam ( Kingma & Ba , 2014 ) and RMSProp ( Hinton et al. , 2012 ) . We show this is highly correlated to a high EGF ( gradient flow ) and how batch normalization helps stabilize EGF . 4 . Changing activation functions can benefit sparse networks We benchmark a wide set of activation functions , specifically ReLU ( Nair & Hinton , 2010 ) and non-saturating activation functions such as PReLU ( He et al. , 2015 ) , ELU ( Clevert et al. , 2015 ) , SReLU ( Jin et al. , 2015 ) , Swish ( Ramachandran et al. , 2017 ) and Sigmoid ( Neal , 1992 ) . Our results show that when using adaptive optimization methods , Swish is a promising activation function , while when using stochastic gradient descent , PReLU preforms better than the other activation functions . Implications Our work is timely as sparse training dynamics are poorly understood . Most training algorithms and methods have been developed to suit training dense networks . Our work provides insight into the nature of sparse optimization and suggests a wider viewpoint beyond initialization is necessary to converge sparse networks to comparable performance as dense . Our proposed approach provides a more accurate measurement of the training dynamics of sparse networks and can be used to inform future work on the design of networks and optimization techniques that are tailored explicitly to sparsity . 2 METHODOLOGY . 2.1 SAME CAPACITY SPARSE VS DENSE COMPARISON . Our goal is to measure what architecture and optimization choices favor sparse networks relative to dense networks . To fairly compare sparse and dense networks , we propose Same Capacity Sparse vs Dense Comparison ( SC-SDC ) , a simple framework which allows us to study sparse network optimization and identify what training configurations are not well suited for sparse networks . SC-SDC can be summarized as follows ( See Figure 1 for an overview ) : 1 . Initialize For a chosen network depth ( number of layers ) L and a maximum network width NMaxW , we compare sparse and dense networks at various widths , while ensuring they have the same paramater count . Initially , we mask the weights θS of sparse network S : alS = θ l S ml , alD = θlD , for l = 1 , . . . , L ( 1 ) , where θlS ml denotes an element-wise product of the weights θS of layer l and the random binary matrix ( mask ) for layer l , ml , alS is the nonzero weights in layer l of sparse network S and a l D is the nonzero weights in layer l of dense network D ( all the weights since no masking occurs ) . For a fair comparison , we need to ensure the same number of nonzero weights for sparse network S and dense network D , across each layer L. ||alS ||0 = ||alD||0 , for l = 1 , . . . , L ( 2 ) We provide more implementation details of how we achieve this in Appendix A.1 . 2 . Match active weight distributions Following prior work ( Liu et al. , 2018b ; Gale et al. , 2019 ) , we ensure the nonzero weights at initialization of the sparse and dense networks are sampled from the same distribution at each layer as follows : alS ∼ P l , alD ∼ P l , for l = 1 , . . . , L ( 3 ) , where P l refers to the initial weight distribution at layer l , for example Kaiming initialization ( He et al. , 2015 ) . This ensures that both sets of active weights ( sparse and dense ) are initially sampled from the same distribution . 3 . Train We then train the sparse and dense networks for 1000 epochs ( allowing for convergence ) . 4 . Evaluate the better architecture We gather the results across the widths/capacity levels and conduct a paired , one-tail Wilcoxon signed-rank test ( Wilcoxon , 1945 ) to evaluate the better architecture . Our null hypothesis ( H0 ) is that sparse networks have similar or worse test accuracy than dense networks ( lower or the same median ) , while our alternative hypothesis ( H1 ) is that sparse networks have better test accuracy performance than dense networks of the same capacity ( higher median ) . This can be formulated as : H0 : Sparse < = Dense , H1 : Sparse > Dense ( 4 ) Our goal of SC-SDC is to compare sparse and dense networks at the same capacity level . By same capacity , we are referring to the same number of active weights , but other notions of same capacity can also be used . We briefly discuss this in Appendix A.1 . We compare the average absolute Kendall Rank correlation between different formulations of gradient flow and generalization . The subscript denotes the p-norm ( l1 or l2 norm ) . We see that EGF has higher absolute correlation when compared to standard gradient flow measures . We also see that is consistent across Fashion MNIST ( see Appendix B.1 ) . 2.2 MEASURING GRADIENT FLOW . Gradient flow ( GF ) is used to study optimization dynamics and typically approximated by taking the norm of the gradients of the network ( Pascanu et al. , 2013 ; Nocedal et al. , 2002 ; Chen et al. , 2018 ; Wang et al. , 2020 ; Evci et al. , 2020 ) . We consider a feedforward neural network f : RD → R , with function inputs x ∈ RD and network weights θ . The gradient norm is usually computed by concatenating all the gradients of a network into a single vector , g = ∂C∂θ , whereC is our cost function . Then the vector norm is taken as follows : gfp = ||g||p , ( 5 ) where p denotes the pth-norm . Effective Gradient Flow Traditional measures of gradient flow take the l1 or l2 norm of all the gradients ( Chen et al. , 2018 ; Pascanu et al. , 2013 ; Evci et al. , 2020 ) . This is not appropriate for sparse networks , as this would include gradients of masked weights which have no influence on the forward pass . Furthermore , computing l1 or l2 across all weights in the networks gives disproportionate influence to layers with more weights . We instead propose a simple modification of Equation 5 , which we term Effective Gradient Flow ( EGF ) , that computes the average , masked gradient ( only gradients of active weights ) norm across all layers . We calculate EGF as follows : g = ( ∂C ∂θ1 m1 , ∂C ∂θ2 m2 , . . . , ∂C ∂θC mL ) for l = 1 , . . . , L ( 6 ) egfp = ∑L n=1 ||gi||p L , ( 7 ) where L is number of layers and ∂C ∂θl ml denotes an element-wise product of the gradients of layer l , ∂C ∂θl , and the mask ml applied to the weights of layer l. For a fully dense network , ml is a matrix of all ones , since no gradients are masked . EGF has the following favourable properties : • Gradient flow is evenly distributed across layers EGF distributes the gradient norm across the layers equally , preventing layers with a lot of weights from dominating the measure and also preventing layers with vanishing gradients from being hidden in the formulation , as is the case with equation 5 ( when all gradients are appended together ) . • Only gradients of active weights are used EGF ensures that for sparse networks , only gradients of active weights are used . Even though weights are masked , their gradients are not necessarily • Possibility for application in gradient-based pruning methods Tanaka et al . ( 2020 ) showed that gradient-based pruning methods like GRASP ( Wang et al. , 2020 ) and SNIP ( Lee et al. , 2018a ) , disproportionately prune large layers and are susceptible to layer-collapse , which is when an algorithm prunes all the weights in a specific layer . Due to the fact that EGF is evenly distributed across layers , maintaining EGF ( as opposed to standard gradient norm ) could possibly be used as a pruning criteria . Furthermore , current approaches measuring or approximating the change in gradient flow during pruning in sparse networks ( Wang et al. , 2020 ; Evci et al. , 2020 ; Singh Lubana & Dick , 2020 ) , could benefit from this new formulation . To evaluate EGF against other standard gradient norm measures , such as the l1 and l2 norm , we empirically compare these measures and their correlation to test loss and accuracy . We take the absolute average of the Kendall Rank correlation ( Kendall , 1938 ) , across the different experiment configurations . We follow a similiar approach to Jiang et al . ( 2019 ) , but unlike their work which has focused on correlating network complexity measures to the generalization gap , we measure the correlation of gradient flow to performance ( accuracy and loss ) . We measure gradient flow at 10 points evenly spaced throughout training , specifically at the end of epoch 0 , 99 . 199 , 299 , 399 , 499 , 599 , 699 , 799 , 899 and 999 . Our results from Table 1 shows that EGF has a higher average absolute correlation to both test loss and accuracy . This is also true of Fashion MNIST ( see Appendix B.1 ) . Due to the comparative benefits of EGF , we use it for the remainder of the paper to measure the impact of interventions . We include all measures of gradient flow in Appendix B for completeness .

Summary: Recently, initialization has been found critical to the model accuracy attained by sparse training [1]. In this paper, the authors studies the impact of factors other than initialization on the model accuracy attained by sparse training relative to dense training (under the same model parameter count). At the core of this paper, the authors argue that the effective gradient flow (grad norm from only activate model weight dimensions) is an effective indicator on the model accuracy attained by sparse training. Firstly, the authors show that the effective gradient flow attains higher correlation with model accuracy than the norm of full gradient (including gradients on sparsified weight dimensions) in sparse training. Secondly, the paper empirically demonstrate that in sparse training, 1) weight decay and data augmentation can hurt model accuracy, 2) batch normalization plays significantly role for model accuracy in sparse training and 3) non-saturating activations boost the magnitude of effective gradient flow and consequently improve model accuracy.