Learned Threshold Pruning

TL;DR

This paper introduces a differentiable, scalable weight pruning method for deep neural networks that learns per-layer thresholds via gradient descent, enabling efficient compression while maintaining high accuracy.

Contribution

The paper proposes a novel learned-threshold pruning (LTP) technique that optimizes pruning thresholds through gradient descent and incorporates a differentiable L0 regularization for architectures with batch-normalization.

Findings

LTP prunes ResNet50 by a factor of 9.1 in 30 epochs.

Achieves 26.4x compression on AlexNet with 79.1% Top-5 accuracy.

Effectively prunes modern compact architectures like EfficientNet and MobileNetV2.

Abstract

This paper presents a novel differentiable method for unstructured weight pruning of deep neural networks. Our learned-threshold pruning (LTP) method learns per-layer thresholds via gradient descent, unlike conventional methods where they are set as input. Making thresholds trainable also makes LTP computationally efficient, hence scalable to deeper networks. For example, it takes $30$ epochs for LTP to prune ResNet50 on ImageNet by a factor of $9.1$. This is in contrast to other methods that search for per-layer thresholds via a computationally intensive iterative pruning and fine-tuning process. Additionally, with a novel differentiable $L_0$ regularization, LTP is able to operate effectively on architectures with batch-normalization. This is important since $L_1$ and $L_2$ penalties lose their regularizing effect in networks with batch-normalization. Finally, LTP generates a trail of progressively sparser networks from which the desired pruned network can be picked based on sparsity and performance requirements. These features allow LTP to achieve competitive compression rates on ImageNet networks such as AlexNet ($26.4\times$ compression with $79.1\%$ Top-5 accuracy) and ResNet50 ($9.1\times$ compression with $92.0\%$ Top-5 accuracy). We also show that LTP effectively prunes modern \textit{compact} architectures, such as EfficientNet, MobileNetV2 and MixNet.