Waste not, Want not: All-Alive Pruning for Extremely Sparse Networks

1 INTRODUCTION . The state-of-the-art neural networks have shown remarkable performance gains on various downstream tasks such as computer vision , natural language processing , and speech recognition . Because neural networks are typically overparameterized , they require high computational cost and memory consumption . Such a nature inherently hinders the deployment of models with excessive parameters on low-end devices such as mobile , embedded , and on-device systems . Network pruning ( Reed , 1993 ) is the prevalent technique to compress high-capacity models by removing unnecessary units such as weights/filters while maintaining the performance with minimal accuracy loss . Existing pruning methods can be divided into two categories . 1 ) The first group enforces the sparsity as model regularization ( Chauvin , 1988 ; Weigend et al. , 1990 ; Ishikawa , 1996 ; Molchanov et al. , 2017a ; Carreira-Perpiñán & Idelbayev , 2018 ; Louizos et al. , 2018 ) . It is theoretically well-investigated and does not require network retraining . 2 ) Another group develops saliency criteria to prune less important units ( Mozer & Smolensky , 1988 ; LeCun et al. , 1989 ; Karnin , 1990 ; Hassibi et al. , 1993 ; Han et al. , 2015 ; Guo et al. , 2016 ; Lee et al. , 2019 ; Park et al. , 2020 ; Evci et al. , 2019 ) . Because of its simple operation and outstanding pruning performance , magnitude pruning ( MP ) ( Han et al. , 2015 ; Narang et al. , 2017 ; Zhu & Gupta , 2018 ) is the most popular saliency-based pruning method . Recently , the effectiveness of MP is highlighted by the success of the lottery ticket hypothesis ( Frankle & Carbin , 2019 ) and learning rate rewinding ( Renda et al. , 2020 ) , which achieves less than 1 % accuracy loss even after pruning 90 % of the parameters . However , all pruning methods show that the trade-off between sparsity and accuracy is significantly degraded , especially at the extremely high sparsity ( Gale et al. , 2019 ; Liu et al. , 2019 ; Blalock et al. , 2020 ) . Considering the explosive increase in the size of state-of-the-art models for downstream tasks ( e.g. , GPT-3 ( Brown et al. , 2020 ) for machine translation and FixEfficientNet-L2 ( Touvron et al. , 2020 ) for image classification ) , the effective pruning methods at extreme compression ratios must be accomplished , and it is particularly crucial for adopting high-capacity models to low-resource devices ( e.g . One Laptop per Child ) . To break the performance bottleneck at high compression ratios of the state-of-the-art pruning methods , we investigate the procedure of existing pruning methods . Surprisingly , it is revealed that all existing studies overlook the existence of dead neurons after pruning – dead neurons are the nodes 20 40 60 80 100 32 × 64 × 12 8× 25 6× 51 2× 10 24 × 0 20 40 60 T es t A cc u ra cy ( % ) % o f d ea d c o n n ec ti o n s Compression ratio LeNet-300-100 MNIST 20 40 60 80 100 12 8× 25 6× 51 2× 10 24 × 20 48 × 40 96 × 0 20 40 60 T es t A cc u ra cy ( % ) % o f d ea d c o n n ec ti o n s Compression ratio ResNet-18 CIFAR-10 0 15 30 45 60 32 × 64 × 12 8× 25 6× 51 2× 10 24 × 0 20 40 60 T es t A cc u ra cy ( % ) % o f d ea d c o n n ec ti o n s Compression ratio ResNet-50 Tiny-ImageNet IMP Accuracy % of dead connections Figure 1 : Model accuracy and percentage of dead connections under varying compression ratios : LeNet-300-100 on MNIST , ResNet-18 on CIFAR-10 , and ResNet-50 on Tiny-ImageNet . We used IMP with learning rate rewinding ( Renda et al. , 2020 ) as the base pruning method . The percentage of dead connections is calculated by the number of dead connections divided by the total number of connections in a given network . with no input connections or output connections . That is , dead neurons make all their connected weights useless , called dead connections . Although Han et al . ( 2015 ) have already raised the potential issue of dead neurons , they anticipated that the dead neurons could be negligible after multiple iterative pruning . Unfortunately , we find that this problem is not as simple as their expectation . Figure 1 depicts prediction accuracy ( green line ) and the percentage of dead connections ( red line ) over various compression ratios . For this empirical study , iterative magnitude pruning ( IMP ) with learning rate rewinding ( Renda et al. , 2020 ) is employed as the baseline pruning method . Then , we analyze the correlation between pruning accuracy and the occurrence of dead neurons on three benchmark datasets . We discover two meaningful findings : 1 ) The dead neurons can not be successfully removed by iterative pruning , especially at high compression ratios . 2 ) The severe performance degradation at high compression ratios is strongly correlated with the number of dead connections , i.e. , they are inversely proportional to each other . Based on these valuable findings , we propose all-alive pruning ( AAP ) , which improves the pruning performance by effectively eliminating dead connections . Specifically , we search dead neurons at the pruning stage by inspecting their gradient flows – when zero gradients are passing through the node , we regard them as dead neurons . Once identified , dead neurons and corresponding dead connections ( any weights linked from or to dead neurons ) are removed together during the pruning procedure . Note that detecting gradient-based dead neurons can be applied for complex model architectures with skip connections . If we remove more weights than a sparsity threshold , we can revive some weights with the highest saliency scores . As a result , AAP constitutes all-alive subnetwork , i.e. , all connections in the subnetwork are kept as trainable weights . To summarize , the key advantage of AAP is two-fold : 1 ) AAP is versatile – it is broadly applicable to various saliency-based pruning methods and model architectures by minimizing the loss of prediction accuracy . 2 ) AAP consistently improves the accuracy of the original pruning methods at high compression ratios and breaks the state-of-the-art performance on three benchmark datasets ( i.e. , MNIST , CIFAR-10 , and Tiny-ImageNet ) . 2 RELATED WORK . Network pruning utilizes the highly overparameterized nature of modern neural networks to compress the model . Also , it provides meaningful insight into what leads to the success of neural networks ; small subnetworks in the original network can achieve comparable performance and improve generalization effects ( Arora et al. , 2018 ) . Recent advances . In general , network pruning can be categorized into two groups . The first approach employs the loss function with regularization terms to enforce the sparsity ( Chauvin , 1988 ; Weigend et al. , 1990 ; Ishikawa , 1996 ) . Also , Molchanov et al . ( 2017a ) proposed variational dropout to produce highly sparse networks and Louizos et al . ( 2018 ) reparameterized network weights as the product of weight and a stochastic gate variable . Carreira-Perpiñán & Idelbayev ( 2018 ) attempted to minimize the loss using the stochastic version of projected gradient descent . The second approach utilizes the sensitivity of the loss for weights ( LeCun et al. , 1989 ; Karnin , 1990 ; Hassibi et al. , 1993 ) . For saliency-based criteria , Han et al . ( 2015 ) proposed iterative pruning based on the magnitude of weights . Despite its simplicity , magnitude pruning shows significant performance gain at high compression ratios ( Guo et al. , 2016 ; Narang et al. , 2017 ; Zhu & Gupta , 2018 ) . Recently , Park et al . ( 2020 ) proposed lookahead pruning ( LAP ) that considers the neighbors of connections . While iterative pruning requires pre-trained networks , SNIP ( Lee et al. , 2019 ) and GraSP ( Wang et al. , 2020 ) proposed to prune the network at initialization . Also , dynamic pruning methods ( e.g. , SET ( Mocanu et al. , 2018 ) , DSR ( Mostafa & Wang , 2019 ) , and RigL ( Evci et al. , 2019 ) ) adjust the sparsity level with prune-and-grow cycles . Comprehensive analysis and experimental results for network pruning can be found ( Gale et al. , 2019 ; Liu et al. , 2019 ; Blalock et al. , 2020 ) . Detecting dead connections in network pruning . Pruning neurons/weights may result in unwanted side-effects , producing useless connections in the pruned networks . In structured pruning , Li et al . ( 2017 ) and Liu et al . ( 2020 ) attempted to eliminate dead neurons by merely considering the connections between adjacent filters . However , it is not applicable for complex networks with shortcut connections . In unstructured pruning , Han et al . ( 2015 ) discussed the existence of dead weights but neglected them . Evci et al . ( 2019 ) proposed dynamic pruning , which iteratively updates the weights to identify more useful weights under given computing constraints . Although dynamic pruning eliminates the connections with relatively low importance , it still can not correctly handle dead connections in each pruning stage . In contrast , this paper devises a new pruning method that removes dead connections and keeps all trainable connections in general networks , including shortcut connections and batch normalization . 3 PROPOSED METHOD : ALL-ALIVE PRUNING ( AAP ) . Given a dataset D = { x ( i ) , y ( i ) } ni=1 and a sparsity threshold τ ( i.e. , the number of non-zero weights ) , unstructured pruning can be formulated as the constrained optimization problem . min c , w L ( c w ; D ) = min c , w 1 n n∑ i=1 ` ( c w ; D ) , s.t . c ∈ { 0 , 1 } m and ||c||0 ≤ τ , ( 1 ) where m is the number of parameters , and n is the number of samples in D. Also , ` ( · ) is the loss function , is the Hadamard product , and || · ||0 is the L0 norm . c ∈ { 0 , 1 } m is the indicator variable to determine whether or not w is pruned . In other words , cj indicates whether wj is active ( cj = 1 ) or inactive ( cj = 0 ) . In this work , we aim at improving saliency-based pruning , where the importance of weights in the network is measured by a pre-defined saliency criterion . Popular criteria include the magnitude of weights ( Han et al. , 2015 ) , the sensitivity of weights ( Lee et al. , 2019 ) , and the neighboring connections of weights ( Park et al. , 2020 ) . Given a saliency criterion , we calculate the importance of weight wj as a saliency score sj ; the higher the saliency score , the more important the weight is . Then , each element of c is expressed as cj = 1 if sj ∈ S and cj = 0 otherwise , where j ∈ { 1 , · · · , m } and S is a set of τ -largest scores associated with w. That is , τ connections are kept in the original network , and the rest of the connections are removed from the network .

Summary: The paper proposes to remove dead neurons and their connected parameters through a very simple check while reviving pruned (salient) parameters up to the prespecified sparsity level, such that the sparse network obtained could perform better. The main (and perhaps the single major) contribution of this work is in its demonstration that such a simple method is indeed effective for different pruning methods on various network architectures and datasets. The proposed method (AAP) can perhaps be considered as a generic post-processing step that could be equipped to any pruning method leaving dead neurons.