ChipNet: Budget-Aware Pruning with Heaviside Continuous Approximations

1 INTRODUCTION . Convolution Neural Networks ( CNNs ) have resulted in several breakthroughs across various disciplines of deep learning , especially for their effectiveness in extracting complex features . However , these models demand significantly high computational power , making it hard to use them on lowmemory hardware platforms that require high-inference speed . Moreover , most of the existing deep networks are heavily over-parameterized resulting in high memory footprint ( Denil et al. , 2013 ; Frankle & Carbin , 2018 ) . Several strategies have been proposed to tackle this issue , that include network pruning ( Liu et al. , 2018 ) , neural architecture search using methods such as reinforcement learning ( Jaafra et al. , 2019 ) and vector quantization ( Gong et al. , 2014 ) , among others . Among the methods outlined above , network pruning has proved to be very effective in designing small resource-efficient architectures that perform at par with their dense counterparts . Network pruning refers to removal of unnecessary weights or filters from a given architecture without compromising its accuracy . It can broadly be classified into two categories : unstructured pruning and structured pruning . Unstructured pruning involves removal of neurons or the corresponding connection weights from the network to make it sparse . While this strategy reduces the number of parameters in the model , computational requirements are still the same ( Li et al. , 2017 ) . Structured pruning methods on the other hand remove the entire channels from the network . This strategy pre- ∗All authors contributed equally . 1Code is publicly available at https : //github.com/transmuteAI/ChipNet serves the regular structure , thereby taking advantage of the high degree of parallelism provided by modern hardware ( Liu et al. , 2017 ; Gordon et al. , 2018 ) . Several structured pruning approaches have been proposed in the recent literature . A general consensus is that variational approaches using sparsity prior loss and learnable dropout parameters outperform the deterministic methods ( Lemaire et al. , 2019 ) . Some of these methods learn sparsity as a part of pretraining , and have proved to perform better than the three stage pretrain-prune-finetune methods . However , since such approaches need to train the model from scratch with pruning-related variables embedded into the network , they can not benefit from off-the-shelf pretrained weights ( Liu et al. , 2017 ; Alvarez & Salzmann , 2017 ) . Others require choosing hyperparameters based on the choice of the network , and can not be easily adapted for new models ( Gordon et al. , 2018 ) . Further , with most of these methods , controlled pruning can not be performed , and a resource-usage constraint can only be satisfied through trial-and-error approach . Recently , Lemaire et al . ( 2019 ) presented a budget-aware pruning method that includes the budget constraint as a part of the training process . A major drawback of this approach and other recent methods is that they are unstable for very low resource budgets , and require additional tricks to work . Overall , a robust budget-aware pruning approach that can be coupled with different budget constraints as well as maintains stability for very low target budgets , is still missing in the existing literature . In this paper , we present ChipNet , a deterministic strategy for structured pruning that employs continuous Heaviside function and crispness loss to identify a highly sparse network out of an existing pretrained dense network . The abbreviation ‘ ChipNet ’ stands for Continuous Heaviside Pruning of Networks . Our pruning strategy draws inspiration from the field of design optimization , where the material distribution task is posed as a continuous optimization problem , but only discrete values ( 0 or 1 ) are practically feasible . Thus , only such values are produced as final outcomes through continuous Heaviside projections . We use a similar strategy to obtain the masks in our sparsity learning approach . The flexible design of ChipNet facilitates its use with different choices of budget constraints , such as restriction on the maximum number of parameters , FLOPs , channels or the volume of activations in the network . Through experiments , we show that ChipNet consistently outperforms state-of-the-art pruning methods for different choices of budget constraints . ChipNet is stable for even very low resource budgets , and we demonstrate this through experiments where network is pruned to as low as 1 % of its parameters . We show that for such extreme cases , ChipNet outperforms the respective baselines by remarkable margins , with a difference in accuracy of slightly beyond 16 % observed for one of the experiments . The masks learnt by ChipNet are transferable across datasets . We show that for certain cases , masks transferred from a model trained on feature-rich teacher dataset provide better performance on the student dataset than those obtained by directly pruning on the student data itself . 2 RELATED WORK . As has been stated in the hypothesis by Frankle & Carbin ( 2018 ) , most neural networks are overparameterized with a large portion ( as much as 90 % ) of the weights being of little significance to the output of the model . Clearly , there exists enormous scope to reduce the size of these networks . Several works have explored the efficiency of network pruning strategies for reducing storage requirements of these networks and accelerating inference speed ( LeCun et al. , 1990 ; Dong et al. , 2017 ) . Some early works by Han et al . ( 2015a ; c ) ; Zhu & Gupta ( 2017 ) involve removal of individual neurons from a network to make it sparse . This reduces the storage requirements of these networks , however , no improvement in inference speed is observed . Recently , several works have focused on structured network pruning , as it involves pruning the entire channel/filters or even layers to maintain the regular structure ( Luo et al. , 2017 ; Li et al. , 2017 ; Alvarez & Salzmann , 2016 ) . The focus of this paper is on structured network pruning , thus , we briefly discuss here the recent works related to this approach . The recent work by Li et al . ( 2017 ) identifies less important channels based on L1-norm . Luo et al . ( 2017 ) ; He et al . ( 2017 ) perform channel selection based on their influence on the activation values of the next layer . Liu et al . ( 2017 ) perform channel-level pruning by imposing LASSO regularization on the scaling terms in the batchnorm layers , and prune the model based on a global threshold . He et al . ( 2018b ) automatically learn the compression ratio of each layer with reinforcement learning . Louizos et al . ( 2017 ) ; Alvarez & Salzmann ( 2017 ; 2016 ) train and prune the network in a single stage strategy . The above mentioned approaches can not optimize networks for a pre-defined budget constraint . Adding budget constraint to the pruning process can provide a direct control on the size of the pruned network . For example , Morphnet imposes this budget by iteratively shrinking and expanding a network through a sparsifying regularizer and uniform layer wise width multiplier , respectively , and is adaptable to specific resource constraints ( Gordon et al. , 2018 ) . However , it requires a modelspecific hyperparameter grid search for choosing the regularization factor . Another approach is BAR ( Lemaire et al. , 2019 ) that uses a budget-constrained pruning approach based on variational method . A limitation of this approach is that for low resource budgets , it needs to explicitly ensure that at least one channel is active in the downsample layer to avoid fatal pruning . The approach proposed in this paper does not require any such tweaking , and is stable for even very low resource budgets . 3 PROPOSED APPROACH . 3.1 LEARNING SPARSITY MASKS . Sparsity learning forms the core of our approach . It refers to learning a set of sparsity masks for a dense convolutional neural network ( parent ) . When designing the smaller pruned network ( child ) , these masks identify the parts of the parent that are to be included in the child network . We first describe here the general idea of learning these masks in the context of our method . The proposed approach falls in the category of structured pruning where masks are designed for channels and not individual neurons . Let f : Rd → Rk denote a convolutional neural network with weights W ∈ Rm and a set of hidden channels H ∈ Rp . We define z ∈ Rp as a set of sparsity masks , where zi ∈ z refers to the mask associated with the feature map hi ∈ H. To apply the mask , zi is multiplied with all the entries of hi . The optimization problem can further be stated as min W , z L ( f ( z H ( W ) ; x ) , y ) s.t . V ( z ) = V0 , ( 1 ) where denotes elementwise multiplication and { x , y } ∈ D are data samples used to train the network f . The desired sparsity of the network is defined in terms of the equality constraint , where V ( · ) denotes the budget function and V0 is the maximum permissible budget . Our proposed formulation of pruning is independent from the choice of the budget function . We later show this through experiments with volume budget as in Lemaire et al . ( 2019 ) , channel budget similar to Liu et al . ( 2017 ) , and budget defined in terms of parameters and FLOPs as well . Originally , zi ∈ z would be defined such that zi ∈ { 0 , 1 } , and a discrete optimization problem is to be solved . For the sake of using gradient-based methods , we convert it to a continuous optimization problem , such that zi ∈ [ 0 , 1 ] . Such reformulation would lead to intermediate values of z occurring in the final optimized solution . Any intermediate value of z , for example z = 0.4 , would imply that a fraction of the respective channel is to be used , and clearly such a solution is practically infeasible . We propose to overcome this challenge through the use of simple nonlinear projections and a novel loss term , and these are discussed in detail in the next section . 3.2 CONTINUOUS HEAVISIDE APPROXIMATION AND LOGISTIC CURVES . At the backbone of our pruning strategy lies three important functions : the commonly used logistic curves , continuous Heaviside function and crispness loss term . Figure 1 presents a graphical representation of these functions . We further provide below a brief motivation for the choice of these functions as well as their significance in our pruning approach . Logistic curves . A commonly used function for adding nonlinearity to a neural network ( LeCun et al. , 1998 ) , logistic curve projects an input from the real space to a range of 0-1 ( Figure 1a ) , and can be mathematically stated as z̃ = 1 1 + e−β ( ψ−ψ0 ) , ( 2 ) where ψ denotes the optimization parameter corresponding to the mask z , ψ0 is the midpoint of the curve , and z̃ denotes the resultant intermediate projection . The additional parameter β is used to control the growth rate of the curve , and forms an important ingredient of our approach . While low values of β can produce an approximately linear curve between -1 and 1 , higher values turn it into a step function . During the initial stages of training , we propose to keep β very low , and increase it to higher values at later stages of the optimization process . With increased values of β , the values further from 0.5 are made more favorable for z̃ . In our experience , the logistic curve alone can not be used to obtain approximately discrete ( 0-1 ) solutions for z in a continuous optimization scheme . The nonlinearity introduced by this function can not sufficiently penalize the intermediate values between 0 and 1 , and optimization algorithm can easily identify values of ψ for which the projected values are far from both . An example experiment demonstrating this issue is presented in Appendix C.2 . To circumvent this issue , we add another nonlinear projection using a continuous approximation of the Heaviside function . Continuous Heaviside function . A continuous approximation to the Heaviside step function , referred as continuous Heaviside function in this paper , is a commonly used projection strategy to solve continuous versions of binary optimization ( 0-1 ) problems in the domain of design optimization ( Guest et al. , 2004 ; 2011 ) . The generalized form of this function can be stated as : z = 1− e−γz̃ + z̃e−γ , ( 3 ) where , the parameter γ dictates the curvature of the regularization . Figure 1b shows the continuous Heaviside function for several values of γ . We see that z is linear for γ = 0 and approaches the Heaviside step function for very large values of γ . The advantages of our projection function are twofold . First , during projection , the values close to 0 and 1 are not affected irrespective of the choice of γ . This implies that the masks identified with most confidence in the early stage of training are not directly impacted by the continuation applied on the value of γ , thus helping in the convergence of the training process . Here , ‘ continuation ’ refers to slowly adapting the value of γ during the course of training . Second , even the values of z̃ which are slightly greater than 0 are also nonlinearly projected to close to 1 , and this effect is more prominent for larger values of γ . The projection adds higher penalty on values between 0 and 1 , and makes them extremely unfavorable when higher values of γ are chosen . While the continuous Heaviside function helps to obtain approximately discrete masks , there is still no explicit constraint or penalty function that can regulate this . To overcome this problem , we tie the outputs of logistic and continuous Heaviside functions to define a novel loss term , referred as crispness loss . Crispness Loss . This novel loss term explicitly penalizes the model performance for intermediate values of z , and drives the convergence towards crisp ( 0-1 ) masks . It is defined as the squared L2 norm of the difference between z̃ and z , stated as Lc = ‖z̃ − z‖22 , and from Figure 1c , we see that Lc achieves its minima when either z̃ = z = 0 or z̃ = z = 1 . Further , the trend of this loss function with respect to ψ for different values of β and γ is shown in Figure 1d . It can be seen that for lower values of β and γ , the loss value is low , and the crispness function plays little to no role in driving the pruning process . When the value of γ slowly increases , the peak of the graph shifts upwards as well as towards the left , thereby increasing the penalty associated with values of ψ . This drives the values of ψ to move farther from the origin . The left shift in the graph adds higher penalty on the negative values , forcing them to become even more negative , thus forcing the respective z to move closer to 0 . The additional loss function associated with the model generally favors values towards 1 . For example , cross-entropy loss used for classification would prefer to set all values in z to 1 to be able to maximize the classification accuracy . With increasing values of γ forcing the masks towards 0 , a balance between the two is identified during the training process . The term β acts as a regularizer that to some extent counteracts the too abrupt impact of γ and regulates the convergence of the training process .

Summary: This paper presents a new method for structure pruning called ChipNet. The ChipNet employs continuous Heaviside function with commonly used logistic curve and crispness loss to estimate sparsity masks. A combination of above three components is helpful to obtain approximately discrete solutions for a continuous optimization scheme.  As a result, it is possible to get a highly sparse network out of an existing pre-trained dense network. Experimentally, ChipNet outperforms other previous structured pruning methods by a large margin.