Rethinking Network Pruning -- under the Pre-train and Fine-tune Paradigm

TL;DR

This work tackles the challenge of compressing large pre-trained language models under the pre-train-then-fine-tune paradigm by introducing SparseBERT, a knowledge-aware sparse pruning method. It prunes self-attention and feed-forward linear projections during distillation while preserving both general language knowledge and task-specific knowledge through a teacher-student framework. Empirical results on GLUE and SQuAD show substantial compression (up to ~20x weight/FLOPs) with negligible accuracy loss, outperforming existing pruning approaches. The study also discusses hardware implications and practical trade-offs in parameters, FLOPs, and training/inference time, highlighting SparseBERT's potential for efficient deployment.

Abstract

Transformer-based pre-trained language models have significantly improved the performance of various natural language processing (NLP) tasks in the recent years. While effective and prevalent, these models are usually prohibitively large for resource-limited deployment scenarios. A thread of research has thus been working on applying network pruning techniques under the pretrain-then-finetune paradigm widely adopted in NLP. However, the existing pruning results on benchmark transformers, such as BERT, are not as remarkable as the pruning results in the literature of convolutional neural networks (CNNs). In particular, common wisdom in pruning CNN states that sparse pruning technique compresses a model more than that obtained by reducing number of channels and layers (Elsen et al., 2020; Zhu and Gupta, 2017), while existing works on sparse pruning of BERT yields inferior results than its small-dense counterparts such as TinyBERT (Jiao et al., 2020). In this work, we aim to fill this gap by studying how knowledge are transferred and lost during the pre-train, fine-tune, and pruning process, and proposing a knowledge-aware sparse pruning process that achieves significantly superior results than existing literature. We show for the first time that sparse pruning compresses a BERT model significantly more than reducing its number of channels and layers. Experiments on multiple data sets of GLUE benchmark show that our method outperforms the leading competitors with a 20-times weight/FLOPs compression and neglectable loss in prediction accuracy.

Figures (6)

• Figure 1: How knowledge is transferred under different pruning strategies. (a) is the general pre-training and fine-tuning procedure (Section \ref{['sec:knowledgetransfer']}). $g$ is an encoder. $g_L$ and $g_{L_D}$ are the encoders well-trained on the pre-training and fine-tuning datasets respectively. $\mathcal{L}$ and $\mathcal{D}$ are the general-purpose language knowledge and the task-specific knowledge respectively. There is a domain error between pre-training and testing, and a generalization error between fine-tuning and testing. (b) and (c) are two basic pruning strategies (Section \ref{['sec:knowledgeaware-basic-prune']}). Both $\mathcal{L}_D$ and $\mathcal{L}^{pr}$ are subsets of knowledge $\mathcal{L}$. $\mathcal{L}_D$ is related to the downstream task. $\mathcal{L}^{pr}$ is preserved in a pruned encoder $g_{L^{pr}}$. (d) is the proposed pruning strategy (Sections \ref{['sec:knowledgeaware-proposed-prune']}-\ref{['sec:knowledgeaware-proposed-distill']}). $(\mathcal{L}^{pr})_D$ refers to the knowledge obtained by first pruning and then fine-tuning. $(\mathcal{L}_D)^{pr}$ corresponds to first fine-tuning and then pruning while distilling.
• Figure 2: Knowledge Analysis.
• Figure 3: Illustration of the proposed knowledge-aware compression. Pruning is performed in parallel with distillation, based on specific data from downstream tasks.
• Figure 4: Performance comparison of SparseBERT and the pruning approach described in Figure \ref{['fig:prune_comp2']}.
• Figure 5: Fitting ability comparison of SparseBERT and the pruning approach described in Figure \ref{['fig:prune_comp3']}.