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: This paper proposes an algorithm called Sparse Momentum for learning sparse neural networks. They claim to maintain sparse weights throughout training while achieving dense network performance levels. They also show their method improves training speed up to 5.61x faster training. The provides a decent motivation for why sparse networks can be helpful. The related work section is well summarized and they emphasize that the current work's primary motivation is to reduce training time while maintaining performance. They compare their method with other methods that also maintain sparse neural networks throughout training and involve single training phase, which is fair. Their method consists of primarily 3 phases - i) pruning weights ii) redistribution of weights iii) regrowing of weights based on the exponentially smoothed momentum term for each layer. The method is well explained and motivation is clear. Edge case was also explained for more clarity. However, there is something that needs more clarification in Pg-3, 3rd line, they say the 3 components of the algorithm i) ii) and iii) can be tackled independently with a divide and conquer strategy to gain some computational benefits. In my understanding, the method performs these 3 steps sequentially after each epoch. Not sure how divide and conquer strategy can be used here ?