Sparse Networks from Scratch: Faster Training without Losing Performance

1 INTRODUCTION . Current state-of-the-art neural networks need extensive computational resources to be trained and can have capacities of close to one billion connections between neurons ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Child et al. , 2019 ) . One solution that nature found to improve neural network scaling is to use sparsity : the more neurons a brain has , the fewer connections neurons make with each other ( Herculano-Houzel et al. , 2010 ) . Similarly , for deep neural networks , it has been shown that sparse weight configurations exist which train faster and achieve the same errors as dense networks ( Frankle and Carbin , 2019 ) . However , currently , these sparse configurations are found by starting from a dense network , which is pruned and re-trained repeatedly – an expensive procedure . In this work , we demonstrate the possibility of training sparse networks that rival the performance of their dense counterparts with a single training run – no re-training is required . We start with random initializations and maintain sparse weights throughout training while also speeding up the overall training time . We achieve this by developing sparse momentum , an algorithm which uses the exponentially smoothed gradient of network weights ( momentum ) as a measure of persistent errors to identify which layers are most efficient at reducing the error and which missing connections between neurons would reduce the error the most . Sparse momentum follows a cycle of ( 1 ) pruning weights with small magnitude , ( 2 ) redistributing weights across layers according to the mean momentum magnitude of existing weights , and ( 3 ) growing new weights to fill in missing connections which have the highest momentum magnitude . We compare the performance of sparse momentum to compression algorithms and recent methods that maintain sparse weights throughout training . We demonstrate state-of-the-art sparse performance on MNIST , CIFAR-10 , and ImageNet-1k . For CIFAR-10 , we determine the percentage of weights needed to reach dense performance levels and find that AlexNet , VGG16 , and Wide Residual Networks need between 35-50 % , 5-10 % , and 20-30 % weights to reach dense performance levels . We also estimate the overall speedups of training our sparse convolutional networks to dense performance levels on CIFAR-10 for optimal sparse convolution algorithms and naive dense convolution algorithms compared to dense baselines . For sparse convolution , we estimate speedups between 2.74x and 5.61x and for dense convolution speedups between 1.07x and 1.36x . In your analysis , ablations demonstrate that the momentum redistribution and growth components are increasingly important as networks get deeper and larger in size – both are critical for good ImageNet performance . 2 RELATED WORK . From Dense to Sparse Neural Networks : Work that focuses on creating sparse from dense neural networks has an extensive history . Earlier work focused on pruning via second-order derivatives ( LeCun et al. , 1989 ; Karnin , 1990 ; Hassibi and Stork , 1992 ) and heuristics which ensure efficient training of networks after pruning ( Chauvin , 1988 ; Mozer and Smolensky , 1988 ; Ishikawa , 1996 ) . Recent work is often motivated by the memory and computational benefits of sparse models that enable the deployment of deep neural networks on mobile and low-energy devices . A very influential paradigm has been the iterative ( 1 ) train-dense , ( 2 ) prune , ( 3 ) re-train cycle introduced by Han et al . ( 2015 ) . Extensions to this work include : Compressing recurrent neural networks and other models ( Narang et al. , 2017 ; Zhu and Gupta , 2018 ; Dai et al. , 2018 ) , continuous pruning and re-training ( Guo et al. , 2016 ) , joint loss/pruning-cost optimization ( Carreira-Perpinán and Idelbayev , 2018 ) , layer-by-layer pruning ( Dong et al. , 2017 ) , fast-switching growth-pruning cycles ( Dai et al. , 2017 ) , and soft weight-sharing ( Ullrich et al. , 2017 ) . These approaches often involve re-training phases which increase the training time . However , since the main goal of this line of work is a compressed model for mobile devices , it is desirable but not an important main goal to reduce the run-time of these procedures . This is contrary to our motivation . Despite the difference in motivation , we include many of these dense-to-sparse compression methods in our comparisons . Other compression algorithms include L0 regularization ( Louizos et al. , 2018 ) , and Bayesian methods ( Louizos et al. , 2017 ; Molchanov et al. , 2017 ) . For further details , see the survey of Gale et al . ( 2019 ) . Interpretation and Analysis of Sparse Neural Networks : Frankle and Carbin ( 2019 ) show that “ winning lottery tickets ” exist for deep neural networks – sparse initializations which reach similar predictive performance as dense networks and train just as fast . However , finding these winning lottery tickets is computationally expensive and involves multiple prune and re-train cycles starting from a dense network . Followup work concentrated on finding these configurations faster ( Frankle et al. , 2019 ; Zhou et al. , 2019 ) . In contrast , we reach dense performance levels with a sparse network from random initialization with a single training run while accelerating training . Sparse Neural Networks Throughout Training : Methods that maintain sparse weights throughout training through a prune-redistribute-regrowth cycle are most closely related to our work . Bellec et al . ( 2018 ) introduce DEEP-R , which takes a Bayesian perspective and performs sampling for prune and regrowth decisions – sampling sparse network configurations from a posterior . While theoretically rigorous , this approach is computationally expensive and challenging to apply to large networks and datasets . Sparse evolutionary training ( SET ) ( Mocanu et al. , 2018 ) simplifies prune-regrowth cycles by using heuristics : ( 1 ) prune the smallest and most negative weights , ( 2 ) grow new weights in random locations . Unlike our work , where many convolutional channels are empty and can be excluded from computation , growing weights randomly fills most convolutional channels and makes it challenging to harness computational speedups during training without specialized sparse algorithms . SET also does not include the cross-layer redistribution of weights which we find to be critical for good performance , as shown in our ablation study . The most closely related work to ours is Dynamic Sparse Reparameterization ( DSR ) by Mostafa and Wang ( 2019 ) , which includes the full prune-redistribute-regrowth cycle . However , DSR requires some specific layers to be dense . Our method works in a fully sparse setting and is thus more generally applicable . More distantly related is Single-shot Network Pruning ( SNIP ) ( Lee et al. , 2019 ) , which aims to find the best sparse network from a single pruning decision . The goal of SNIP is simplicity , while our goal is maximizing predictive and run-time performance . In our experiments , we compare against all four methods : DEEP-R , SET , DSR , and SNIP . 3 SPARSE LEARNING . We define sparse learning to be the training of deep neural networks which maintain sparsity throughout training while matching the predictive performance of dense neural networks . To achieve this , intuitively , we want to find the weights that reduce the error most effectively . This is challenging since most deep neural network can hold trillions of different combinations of sparse weights . Additionally , during training , as feature hierarchies are learned , efficient weights might change gradually from shallow to deep layers . How can we find good sparse configurations ? In this work , we follow a divide-and-conquer strategy that is guided by computationally efficient heuristics . We divide sparse learning into the following sub-problems which can be tackled independently : ( 1 ) pruning weights , ( 2 ) redistribution of weights across layers , and ( 3 ) regrowing weights , as defined in more detail below . Sparse Momentum Mean Momentum Contribution per layer 28 % 22 % 38 % 12 % 38 % 28 % 12 % 22 % 8 Weights 16 Weights16 Weights 8 Removed Weights 16 Weights 8 Weights Remove Regrow 8 Removed Weights Redistribute 1 . Calculate mean magnitude of momentum for existing weights 2 . Prune the smallest fraction of weights , p , for each layer 3 . Regrow weights according to momentum of missing weights Figure 1 : Sparse Momentum is applied at the end of each epoch : ( 1 ) take the magnitude of the exponentially smoothed gradient ( momentum ) of each layer and normalize to 1 ; ( 2 ) for each layer , remove p = 20 % of the weights with the smallest magnitude ; ( 3 ) across layers , redistribute the removed weights by adding weights to each layer proportionate to the momentum of each layer ; within a layer , add weights starting from those with the largest momentum magnitude . Decay p . 3.1 SPARSE MOMENTUM . We use the mean magnitude of momentum Mi of existing weights Wi in each layer i to estimate how efficient the average weight in each layer is at reducing the overall error . Intuitively , we want to take weights from less efficient layers and redistribute them to weight-efficient layers . The sparse momentum algorithm is depicted in Figure 1 . In this section , we first describe the intuition behind sparse momentum and then present a more detailed description of the algorithm . The gradient of the error with respect to a weight ∂E∂W yields the directions which reduce the error at the highest rate . However , if we use stochastic gradient descent , most weights of ∂E∂W oscillate between small/large and negative/positive gradients with each mini-batch ( Qian , 1999 ) – a good change for one mini-batch might be a bad change for another . We can reduce oscillations if we take the average gradient over time , thereby finding weights which reduce the error consistently . However , we want to value recent gradients , which are closer to the local minimum , more highly than the distant past . This can be achieved by exponentially smoothing ∂E∂W – the momentum Mi : Mt+1i = αM t i + ( 1− α ) ∂E ∂Wi t , where α is a smoothing factor , Mi is the momentum for the weight Wi in layer i ; Mi is initialized at t = 0 with 0 . Momentum is efficient at accelerating the optimization of deep neural networks by identifying weights which reduce the error consistently . Similarly , the aggregated momentum of weights in each layer should reflect how good each layer is at reducing the error consistently . Additionally , the momentum of zero-valued weights – equivalent to missing weights in sparse networks – can be used to estimate how quickly the error would change if these weights would be included in a sparse network . The details of the full training procedure of our algorithm are shown in Algorithm 1 . See Algorithm 2 in the Appendix for a more detailed , source-code-like description of sparse momentum . Algorithm 1 : Sparse momentum algorithm . Data : Layer i to k with : Momentum Mi , Weight Wi , binary Maski prune rate pi , density d 1 for i← 0 to k do 2 Wi ← xavierInit ( Wi ) 3 Maski ← createMaskForWeight ( Wi , d ) 4 applyMask ( Wi , Maski ) 5 end 6 for epoch← 0 to numEpochs do 7 for j← 0 to numBatches do 8 batch← getBatch ( j ) 9 ∂E ∂W = computeGradients ( W , batch ) 10 UpdateMomentum ( ∂E∂W ) 11 UpdateWeights ( M ) 12 for i← 0 to k do 13 applyMask ( Wi , Maski ) 14 end 15 end 16 totalMomentum← getTotalMomentum ( M ) 17 totalPruned← getTotalPrunedWeights ( W , p ) 18 for i← 0 to k do 19 mi ← getMomentumContribution ( Mi , Maski , totalMomentum ) 20 magnitudePruneWeight ( Wi , Maski , pi ) 21 regrowWeights ( Wi , Maski , mi · totalPruned ) 22 pi ← decayPrunerate ( pi ) 23 applyMask ( Wi , Maski ) 24 end 25 end Before training , we initialize the network with a certain sparsity s : we initialize the network as usual and then remove a fraction of s weights for each layer . We train the network normally and mask the weights after each gradient update to enforce sparsity . We apply sparse momentum after each epoch . We can break the sparse momentum into three major parts : ( a ) redistribution of weights , ( b ) pruning weights , ( c ) regrowing weights . In step ( a ) , we we take the mean of the element-wise momentum magnitude mi that belongs to all nonzero weights for each layer i and normalize the value by the total momentum magnitude of all layers ∑k i=0mi . The resulting proportion is the momentum magnitude contribution for each layer . The number of weights to be regrow in each layer is the total number of removed weights multiplied by each layers momentum contribution : Regrowi = Total Removed ·mi . In step ( b ) , we prune a proportion of p ( prune rate ) of the weights with the lowest magnitude for each layer . In step ( c ) , we regrow weights by enabling the gradient flow of zero-valued ( missing ) weights which have the largest momentum magnitude . Additionally , there are two edge-cases which we did not include in Algorithm 1 for clarity . ( 1 ) If we allocate more weights to be regrown than is possible for a specific layer , for example regrowing 100 weights for a layer of maximum 10 weights , we redistribute the excess number of weights equally among all other layers . ( 2 ) For some layers , our algorithm will converge in that the average weight in layer i has much larger momentum magnitude than weights in other layers , but at the same time , this layer is dense and can not grow further . We do not want to prune weights from such important layers . Thus , for these layers , we reduce the prune rate pi proportional to the sparsity : pi = min ( p , sparsityi ) . After each epoch , we decay the prune rate in Algorithm 1 in the same way learning rates are decayed . We use a cosine decay schedule that anneals the prune rate to zero on the last epoch . See Appendix A.1 for an analysis on how decay schedule and starting prune rate affects training .

Summary: In this paper, the authors propose a sparse momentum algorithm for doing efficient sparse training. The technique relies on identifying weights in a layer that do not have an effect on the error, pruning them, and redistributing and growing them across layers. The technique is compared against other recent algorithms on a range of models.