BWCP: Probabilistic Learning-to-Prune Channels for ConvNets via Batch Whitening

This work presents a probabilistic channel pruning method to accelerate Convolutional Neural Networks ( CNNs ) . Previous pruning methods often zero out unimportant channels in training in a deterministic manner , which reduces CNN ’ s learning capacity and results in suboptimal performance . To address this problem , we develop a probability-based pruning algorithm , called batch whitening channel pruning ( BWCP ) , which can stochastically discard unimportant channels by modeling the probability of a channel being activated . BWCP has several merits . ( 1 ) It simultaneously trains and prunes CNNs from scratch in a probabilistic way , exploring larger network space than deterministic methods . ( 2 ) BWCP is empowered by the proposed batch whitening tool , which is able to empirically and theoretically increase the activation probability of useful channels while keeping unimportant channels unchanged without adding any extra parameters and computational cost in inference . ( 3 ) Extensive experiments on CIFAR-10 , CIFAR-100 , and ImageNet with various network architectures show that BWCP outperforms its counterparts by achieving better accuracy given limited computational budgets . For example , ResNet50 pruned by BWCP has only 0.58 % Top-1 accuracy drop on ImageNet , while reducing 42.9 % FLOPs of the plain ResNet50 . 1 INTRODUCTION . Deep convolutional neural networks ( CNNs ) have achieved superior performance in a variety of computer vision tasks such as image recognition ( He et al. , 2016 ) , object detection ( Ren et al. , 2017 ) , and semantic segmentation ( Chen et al. , 2018 ) . However , despite their great success , deep CNN models often have massive demand on storage , memory bandwidth , and computational power ( Han & Dally , 2018 ) , making them difficult to be plugged onto resource-limited platforms , such as portable and mobile devices ( Deng et al. , 2020 ) . Therefore , proposing efficient and effective model compression methods has become a hot research topic in the deep learning community . Model pruning , as one of the vital model compression techniques , has been extensively investigated . It reduces model size and computational cost by removing unnecessary or unimportant weights or channels in a CNN ( Han et al. , 2016 ) . For example , many recent works ( Wen et al. , 2016 ; Guo et al. , 2016 ) prune fine-grained weights of filters . Han et al . ( 2015 ) proposes to discard the weights that have magnitude less than a predefined threshold . Guo et al . ( 2016 ) further utilizes a sparse mask on a weight basis to achieve pruning . Although these unstructured pruning methods achieve optimal pruning schedule , they do not take the structure of CNNs into account , preventing them from being accelerated on hardware such as GPU for parallel computations ( Liu et al. , 2018 ) . To achieve efficient model storage and computations , we focus on structured channel pruning ( Wen et al. , 2016 ; Yang et al. , 2019a ; Liu et al. , 2017 ) , which removes entire structures in a CNN such as filter or channel . A typical structured channel pruning approach commonly contains three stages , including pre-training a full model , pruning unimportant channels by the predefined criteria such as ` p norm , and fine-tuning the pruned model ( Liu et al. , 2017 ; Luo et al. , 2017 ) , as shown in Fig.1 ( a ) . However , it is usually hard to find a global pruning threshold to select unimportant channels , because the norm deviation between channels is often too small ( He et al. , 2019 ) . More importantly , as some channels are permanently zeroed out in the pruning stage , such a multi-stage procedure usually not only relies on hand-crafted heuristics but also limits the learning capacity ( He et al. , 2018a ; 2019 ) . To tackle the above issues , we propose a simple but effective probability-based channel pruning framework , named batch-whitening channel pruning ( BWCP ) , where unimportant channels are pruned in a stochastic manner , thus preserving the channel space of CNNs in training ( i.e . the diversity of CNN architectures is preserved ) . To be specific , as shown in Fig.1 ( b ) , we assign each channel with an activation probability ( i.e . the probability of a channel being activated ) , by exploring the properties of the batch normalization layer ( Ioffe & Szegedy , 2015 ; Arpit et al. , 2016 ) . A larger activation probability indicates that the corresponding channel is more likely to be preserved . We also introduce a capable tool , termed batch whitening ( BW ) , which can increase the activation probability of useful channels , while keeping the unnecessary channels unchanged . By doing so , the deviation of the activation probability between channels is explicitly enlarged , enabling BWCP to identify unimportant channels during training easily . Such an appealing property is justified by theoretical analysis and experiments . Furthermore , we exploit activation probability adjusted by BW to generate a set of differentiable masks by a soft sampling procedure with Gumbel-Softmax technique , allowing us to train BWCP in an online “ pruning-from-scratch ” fashion stably . After training , we obtain the final compact model by directly discarding the channels with zero masks . The main contributions of this work are three-fold . ( 1 ) We propose a probability-based channel pruning framework BWCP , which explores a larger network space than deterministic methods . ( 2 ) BWCP can easily identify unimportant channels by adjusting their activation probabilities without adding any extra model parameters and computational cost in inference . ( 3 ) Extensive experiments on CIFAR-10 , CIFAR-100 and ImageNet datasets with various network architectures show that BWCP can achieve better recognition performance given the comparable amount of resources compared to existing approaches . For example , BWCP can reduce 68.08 % Flops by compressing 93.12 % parameters of VGG-16 with merely accuracy drop and ResNet-50 pruned by BWCP has only 0.58 % top-1 accuracy drop on ImageNet while reducing 42.9 % FLOPs . 2 RELATED WORK . Weight Pruning . Early network pruning methods mainly remove the unimportant weights in the network . For instance , Optimal Brain Damage ( LeCun et al. , 1990 ) measures the importance of weights by evaluating the impact of weight on the loss function and prunes less important ones . However , it is not applicable in modern network structure due to the heavy computation of the Hessian matrix . Recent work assesses the importance of the weights through the magnitude of the weights itself . Specifically , ( Guo et al. , 2016 ) prune the network by encouraging weights to become exactly zero . The computation involves weights with zero can be discarded . However , a major drawback of weight pruning techniques is that they do not take the structure of CNNs into account , thus failing to help scale pruned models on commodity hardware such as GPUs ( Liu et al. , 2018 ; Wen et al. , 2016 ) . Channel Pruning . Channel pruning approaches directly prune feature maps or filters of CNNs , making it easy for hardware-friendly implementation . For instance , relaxed ` 0 regularization ( Louizos et al. , 2017 ) and group regularizer ( Yang et al. , 2019a ) impose channel-level sparsity , and filters with small value are selected to be pruned . Some recent work also propose to rank the importance of filters by different criteria including ` 1 norm ( Liu et al. , 2017 ; Li et al. , 2017 ) , ` 2 norm ( Frankle & Carbin , 2018 ) and High Rank channels ( Lin et al. , 2020 ) . For example , ( Liu et al. , 2017 ) explores the importance of filters through scale parameter γ in batch normalization . Although these approaches introduce minimum overhead to the training process , they are not trained in an end-to-end manner and usually either apply on a pre-trained model or require an extra fine-tuning procedure . Recent works tackle this issue by pruning CNNs from scratch . For example , FPGM ( He et al. , 2019 ) zeros in unimportant channels and continues training them after each training epoch . Furthermore , both SSS and DSA learn a differentiable binary mask that is generated by channel importance and does not require any additional fine-tuning . Our proposed BWCP is most related to variational pruning ( Zhao et al. , 2019 ) and SCP ( Kang & Han , 2020 ) as they also employ the property of normalization layer and associate the importance of channel with probability . The main difference is that our method adopts the idea of whitening to perform channel pruning . We will show that the proposed batch whitening ( BW ) technique can adjusts the activation probability of different channels according to their importance , making it easy to identify unimportant channels . Although previous work SPP ( Wang et al. , 2017 ) and DynamicCP ( Gao et al. , 2018 ) also attempt to boost salient channels and skip unimportant ones , they fail to consider the natural property inside normalization layer and deign the activation probability empirically . 3 PRELIMINARY . Notation . We use regular letters , bold letters , and capital letters to denote scalars such as ‘ x ’ , and vectors ( e.g.vector , matrix , and tensor ) such as ‘ x ’ and random variables such as ‘ X ’ , respectively . We begin with introducing a building layer in recent deep neural nets which typically consists of a convolution layer , a batch normalization ( BN ) layer , and a rectified linear unit ( ReLU ) ( Ioffe & Szegedy , 2015 ; He et al. , 2016 ) . Formally , it can be written by xc = wc ∗ z , x̃c = γcx̄c + βc , yc = max { 0 , x̃c } ( 1 ) where c ∈ [ C ] denotes channel index and C is channel size . In Eqn . ( 1 ) , ‘ ∗ ’ indicates convolution operation and wc is filter weight corresponding to the c-th output channel , i.e . xc ∈ RN×H×W . To perform normalization , xc is firstly standardized to x̄c through x̄c = ( xc − E [ xc ] ) / √ D [ xc ] where E [ · ] and D [ · ] indicate calculating mean and variance over a batch of samples , and then is re-scaled to x̃c by scale parameter γc and bias βc . Moreover , the output feature yc is obtained by ReLU activation that discards the negative part of x̃c . Criterion-based channel pruning . For channel pruning , previous methods usually employ a ‘ smallnorm-less-important ’ criterion to measure the importance of channels . For example , BN layer can be applied in channel pruning ( Liu et al. , 2017 ) , where a channel with a small value of γc would be removed . The reason is that the c-th output channel x̃c contributes little to the learned feature representation when γc is small . Hence , the convolution in Eqn . ( 1 ) can be discarded safely , and filter wc can thus be pruned . Unlike these criterion-based methods that deterministically prune unimportant filters and rely on a heuristic pruning procedure as shown in Fig.1 ( a ) , we explore a probability-based channel pruning framework where less important channels are pruned in a stochastic manner . Activation probability . To this end , we define an activation probability of a channel by exploring the property of the BN layer . Those channels with a larger activation probability could be preserved with a higher probability . To be specific , since x̄c is acquired by subtracting the sample mean and being divided by the sample variance , we can treat each channel feature as a random variable following standard Normal distribution ( Arpit et al. , 2016 ) , denoted as X̄c . Note that only positive parts can be activated by ReLU function . Proposition 1 gives the activation probability of the c-th channel , i.e.P ( X̃c ) > 0 . Proposition 1 Let a random variable X̄c ∼ N ( 0 , 1 ) and Yc = max { 0 , γcX̄c + βc } . Then we have ( 1 ) P ( Yc > 0 ) = P ( X̃c > 0 ) = ( 1 + Erf ( βc/ ( √ 2|γc| ) ) /2 where Erf ( x ) = ∫ x 0 2/ √ π · exp−t2dt , and ( 2 ) P ( X̃c > 0 ) = 0⇔ βc ≤ 0 and γc → 0 . Note that a pruned channel can be modelled by P ( X̃c > 0 ) = 0 . With Proposition 1 ( see proof in Appendix A.2 ) , we know that the unnecessary channels satisfy that γc approaches 0 and βc is negative . To achieve channel pruning , previous compression techniques ( Li et al. , 2017 ; Zhao et al. , 2019 ) merely impose a regularization on γc , which would deteriorate the representation power of unpruned channels ( Perez et al. , 2018 ; Wang et al. , 2020 ) . Instead , we adopt the idea of whitening to build a probabilistic channel pruning framework where unnecessary channels are stochastically disgarded with a small activation probability while important channels are preserved with a large activation probability .

Summary: This paper proposes a way to train compact neural networks. The feature channels are whitened, then combined with a sampling method. At the testing time, this sampling is discarded and replaced with a deterministic selection.   The benefits of whitening in improving the convergence of NNs is widely explored. This paper explores another useful property of whitening,  the sparsity in the whitened results. Since the representation is sparser,  one might use fewer weights to reach similar performance. This paper presents such a solution.