[Paper Review] Operation-Aware Soft Channel Pruning using Differentiable Masks

Min Soo Kang, Bohyung Han|arXiv (Cornell University)|Jul 8, 2020

Speech and Audio Processing32 citations

TL;DR

The paper introduces a differentiable, data-driven channel pruning method that jointly learns model parameters and soft masks by leveraging BN and ReLU to prune channels without extra fine-tuning, achieving better accuracy under similar resource budgets.

ABSTRACT

We propose a simple but effective data-driven channel pruning algorithm, which compresses deep neural networks in a differentiable way by exploiting the characteristics of operations. The proposed approach makes a joint consideration of batch normalization (BN) and rectified linear unit (ReLU) for channel pruning; it estimates how likely the two successive operations deactivate each feature map and prunes the channels with high probabilities. To this end, we learn differentiable masks for individual channels and make soft decisions throughout the optimization procedure, which facilitates to explore larger search space and train more stable networks. The proposed framework enables us to identify compressed models via a joint learning of model parameters and channel pruning without an extra procedure of fine-tuning. We perform extensive experiments and achieve outstanding performance in terms of the accuracy of output networks given the same amount of resources when compared with the state-of-the-art methods.

Motivation & Objective

• Motivate and achieve model compression by pruning channels in deep networks without a separate fine-tuning step.
• Leverage the interaction of batch normalization and ReLU to probabilistically prune channels with minimal accuracy loss.
• Develop a differentiable, end-to-end framework that jointly optimizes network parameters and channel masks.
• Provide empirical evidence that the proposed method outperforms prior structured pruning methods across datasets and backbones.

Proposed method

• Define a probabilistic mask for each channel based on BN parameters (beta, gamma) and the effect of ReLU.
• Relax the discrete pruning mask with a differentiable logistic function and employ Gumbel-Softmax for differentiable sampling.
• Modify the BN layer to incorporate the stochastic mask into its forward pass, enabling backpropagation through the mask.
• Introduce a sparsity loss that encourages pruning by increasing the CDF of a Gaussian distribution linked to BN parameters.
• Train the network end-to-end to jointly optimize model weights and masks without a separate fine-tuning stage.

Experimental results

Research questions

• RQ1 Can BN and ReLU together provide a reliable, probabilistic criterion for channel pruning?
• RQ2 Does joint, differentiable learning of masks and parameters yield better accuracy under the same resource constraints compared to prior pruning methods?
• RQ3 How does the sparsity loss with a Gaussian confidence interval influence channel reduction and performance?
• RQ4 Is fine-tuning unnecessary when pruning channels via a differentiable, jointly trained framework?

Key findings

• The proposed SCP method achieves lower accuracy drops than SFP, FPGM, Slimming, and Variational Pruning across CIFAR-10/100 with various backbones.
• On CIFAR datasets, SCP often matches or outperforms baseline accuracy while providing substantial reductions in channels, parameters, and FLOPs.
• On ILSVRC-12, SCP shows competitive top-1/top-5 drops and significant FLOPs reduction, with practical speedup (~24% inference time) without extra fine-tuning.
• Ablation shows that incorporating ReLU in pruning criteria improves performance over BN-only pruning (“SCP without ReLU”).
• Increasing the sparsity-λ in the loss generally yields greater FLOPs reduction with varying effects on accuracy, illustrating a controllable trade-off.