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: The paper provides a budget-aware regularizer method to train the network to be pruned. Existing methods based-on the regularizer method suffer from satisfying the user-specified constraints and resort to trial-and-error approach. The authors leverage logistic function and continuous heaviside function for continuous relaxation of the discrete variable to select the channels to be pruned.  There are four types of budget constraints (channel, volume, parameter, FLOPs).