GDP: Stabilized Neural Network Pruning via Gates with Differentiable Polarization

TL;DR

This paper introduces GDP, a differentiable polarization-based gating method for neural network pruning that effectively removes channels without disrupting training, achieving state-of-the-art results on multiple benchmarks.

Contribution

GDP provides a novel, principled approach to channel pruning by using differentiable polarization to smoothly identify and remove unimportant channels during training.

Findings

Achieves state-of-the-art pruning performance on CIFAR-10 and ImageNet.

Maintains or improves performance on Pascal VOC segmentation with significant FLOPs reduction.

Abstract

Model compression techniques are recently gaining explosive attention for obtaining efficient AI models for various real-time applications. Channel pruning is one important compression strategy and is widely used in slimming various DNNs. Previous gate-based or importance-based pruning methods aim to remove channels whose importance is smallest. However, it remains unclear what criteria the channel importance should be measured on, leading to various channel selection heuristics. Some other sampling-based pruning methods deploy sampling strategies to train sub-nets, which often causes the training instability and the compressed model's degraded performance. In view of the research gaps, we present a new module named Gates with Differentiable Polarization (GDP), inspired by principled optimization ideas. GDP can be plugged before convolutional layers without bells and whistles, to control the on-and-off of each channel or whole layer block. During the training process, the polarization effect will drive a subset of gates to smoothly decrease to exact zero, while other gates gradually stay away from zero by a large margin. When training terminates, those zero-gated channels can be painlessly removed, while other non-zero gates can be absorbed into the succeeding convolution kernel, causing completely no interruption to training nor damage to the trained model. Experiments conducted over CIFAR-10 and ImageNet datasets show that the proposed GDP algorithm achieves the state-of-the-art performance on various benchmark DNNs at a broad range of pruning ratios. We also apply GDP to DeepLabV3Plus-ResNet50 on the challenging Pascal VOC segmentation task, whose test performance sees no drop (even slightly improved) with over 60% FLOPs saving.