One-Shot Pruning for Fast-adapting Pre-trained Models on Devices

TL;DR

This paper introduces a scalable one-shot pruning method that uses knowledge from similar tasks to efficiently extract sub-networks from large pre-trained models, enabling quick adaptation to new tasks on resource-limited devices.

Contribution

The paper proposes a novel one-shot pruning approach leveraging pruned models of similar tasks to identify task-specific sub-networks, reducing resource requirements and training time for new tasks.

Findings

Outperforms popular pruning baselines in accuracy and efficiency.

Effective on CNNs and ViT architectures across various datasets.

Enables quick adaptation with minimal training iterations.

Abstract

Large-scale pre-trained models have been remarkably successful in resolving downstream tasks. Nonetheless, deploying these models on low-capability devices still requires an effective approach, such as model pruning. However, pruning the model from scratch can pose a practical challenge given the limited resources of each downstream task or device. To tackle this issue, we present a scalable one-shot pruning method that leverages pruned knowledge of similar tasks to extract a sub-network from the pre-trained model for a new task. Specifically, we create a score mask using the pruned models of similar tasks to identify task-specific filters/nodes in the pre-trained model for the new task. Based on this mask, we conduct a single round of pruning to extract a suitably-sized sub-network that can quickly adapt to the new task with only a few training iterations. Our experimental analysis demonstrates the effectiveness of the proposed method on the convolutional neural networks (CNNs) and vision transformers (ViT) with various datasets. The proposed method consistently outperforms popular pruning baseline methods in terms of accuracy and efficiency when dealing with diverse downstream tasks with different memory constraints.