Rapid Structural Pruning of Neural Networks with Set-based Task-Adaptive Meta-Pruning

TL;DR

This paper introduces STAMP, a rapid, task-adaptive neural network pruning method that reduces training costs and improves compression rates by generating pruning masks based on target datasets, without requiring pretraining.

Contribution

The paper proposes a novel meta-learning based pruning approach that adapts to new datasets quickly, outperforming existing methods in speed and compression efficiency.

Findings

Significantly improved compression rates over baselines.

Orders of magnitude faster training speed.

Effective generalization to unseen datasets.

Abstract

As deep neural networks are growing in size and being increasingly deployed to more resource-limited devices, there has been a recent surge of interest in network pruning methods, which aim to remove less important weights or activations of a given network. A common limitation of most existing pruning techniques, is that they require pre-training of the network at least once before pruning, and thus we can benefit from reduction in memory and computation only at the inference time. However, reducing the training cost of neural networks with rapid structural pruning may be beneficial either to minimize monetary cost with cloud computing or to enable on-device learning on a resource-limited device. Recently introduced random-weight pruning approaches can eliminate the needs of pretraining, but they often obtain suboptimal performance over conventional pruning techniques and also does not allow for faster training since they perform unstructured pruning. To overcome their limitations, we propose Set-based Task-Adaptive Meta Pruning (STAMP), which task-adaptively prunes a network pretrained on a large reference dataset by generating a pruning mask on it as a function of the target dataset. To ensure maximum performance improvements on the target task, we meta-learn the mask generator over different subsets of the reference dataset, such that it can generalize well to any unseen datasets within a few gradient steps of training. We validate STAMP against recent advanced pruning methods on benchmark datasets, on which it not only obtains significantly improved compression rates over the baselines at similar accuracy, but also orders of magnitude faster training speed.