Venue: NAACL 2021

# Rethinking Network Pruning--

under the Pre-train and Fine-tune Paradigm

 Dongkuan Xu\({}^{1}\) &Ian En-Hsu Yen\({}^{2}\) &Jinxi Zhao\({}^{2}\) &Zhibin Xiao\({}^{2}\)

\({}^{1}\)The Pennsylvania State University, State College, PA, USA

\({}^{2}\)Moffett AI, Los Altos, CA, USA

{ducx19}@psu.edu, {ian.yan, jlx.zhao, zb.xiao}@moffett.ai

###### 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 accuracy1.

Footnote 1: Codes can be found in the authors’ website. Work done during an internship at Moffett AI.

_Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_, pages 2376–2382 June 6-11, 2021. ©2021 Association for Computational Linguistics

## 1 Introduction

Pre-trained language models, such as BERT Devlin et al. (2019), become the standard and effective methods for improving the performance of a variety of natural language processing (NLP) tasks. These models are pre-trained in a self-supervised fashion and then fine-tuned for supervised downstream tasks. However, these models suffer from the heavy model size, making them impractical for resource-limited deployment scenarios and incurring cost concerns Strubell et al. (2019).

In parallel, an emerging subfield has studied the redundancy in deep neural network models Zhu and Gupta (2017); Gale et al. (2019) and proposed to prune networks without sacrificing performance, such as the lottery ticket hypothesis Frankle and Carbin (2019). Common wisdom in CNN literature shows that sparse pruning leads to more compression rate than structural pruning. For example, for the same number of parameters (0.46M), the sparse MobileNets improve by 11.2% accuracy over the dense ones Zhu and Gupta (2017). However, similar conclusions are not observed for pre-trained language models.

The main question this paper attempts to answer is: how to perform sparse pruning under the pre-train and fine-tune paradigm? Answering this question correctly is challenging. First, these models adopt pre-training and fine-tuning procedures, during which the general-purpose language knowledge and the task-specific knowledge are learned respectively. Thus, it is desirable and challenging to keep the weights that are important to both knowledge during pruning. Second, unlike CNNs, pre-trained language models have a complex architecture consisting of embedding, self-attention, and feed-forward layers.

To address these challenges, we propose SparseBERT, a knowledge-aware sparse pruning method for pre-trained language models, with a special focus on the widely used BERT model. SparseBERT is executed in the fine-tuning stage. It preserves both general-purpose and task-specific language knowledge while pruning. To preserve the general-purpose knowledge learned during pre-training,

SparseBERT uses the pre-trained BERT without fine-tuning as the initialized model and prunes the linear transformations in self-attention and feed-forward layers, which is inspired by the recent findings that self-attention and feed-forward layers are overparameterized Michel et al. (2019); Voita et al. (2019) and are also the most computation consumption parts Ganesh et al. (2020). To learn the task-specific task knowledge during pruning while preserving the general-purpose knowledge at the same time, we apply knowledge distillation Hinton et al. (2015). We adopt the task-specific fine-tuned BERT as the teacher network and the pre-trained BERT that is being pruned as the student. We feed the downstream task data into the teacher-student framework to train the student to reproduce the behaviors of the teacher.

We summarize different types of BERT pruning approaches in Figure 1 (see Section 3.2 for detailed discussion) Experimental results on the GLUE benchmark demonstrate that SparseBERT outperforms all the leading competitors and achieves 1.4% averaged loss with down to only 5% remaining weights compared to BERT-base.

## 2 Related Work

A lot of efforts have been made on studying network redundancy and pruning networks without accuracy loss Gale et al. (2019); Renda et al. (2020). For example, the work on lottery ticket hypothesis Frankle and Carbin (2019) showed that there exist sparse smaller subnetworks capable of training to full accuracy in CNNs. Common wisdom in CNN literature shows that spare pruning leads to much more compression rate than structural pruning Gale et al. (2019); Elsen et al. (2020). For example, for the same number of parameters (0.46M), the sparse MobileNets achieve 61.8% accuracy while the dense ones achieve 50.6% Zhu and Gupta (2017). However, similar observations are not observed in existing approaches for pretrained language models Fan et al. (2019); Michel et al. (2019); Chen et al. (2020); McCarley et al. (2020); Jiao et al. (2020). Our method aims to fill the gap and summarize these pruning strategies. There are other compression approaches for pre-trained language models, such as quantization Zafrir et al. (2019) and weight factorization Wang et al. (2019), which are out of the scope of this work.

Figure 1: How knowledge is transferred under different pruning strategies. (a) is the general pre-training and fine-tuning procedure (Section 3.1). \(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 3.2.1). 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 3.2.2-3.2.3). \((\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.

 



## 3 SparseBERT

We first formalize the knowledge transfer involved in fine-tuning pre-trained language models. Then, we introduce our SparseBERT.

### Knowledge Transfer under the Pre-train and Fine-tune Paradigm

The practice of fine-tuning pre-trained language models has become prevalent in various NLP tasks. The two-stage procedure is illustrated in Figure 1(a). The language model is denoted by \(f\) = \(g\circ h\), where \(g\) is a text encoder and \(h\) is a task predictor head. Text encoders, like Transformers in BERT, are used to map input sentences to hidden representations and task predictors further map the representations to the label space. The pre-trained model is trained on a large amount of data examples \(({\bf x}^{p},y^{p})\) from the pre-training task domain via different tasks that resemble language modeling.

During pre-training, the general-purpose language knowledge, denoted by \(\mathcal{L}\), is learned based on \(({\bf x}^{p},y^{p})\). \(\mathcal{L}\) contains a subset that is related to the downstream task, denoted by \(\mathcal{L}_{D}\), and the amount of \(\mathcal{L}\) is far greater than that of \(\mathcal{L}_{D}\) (see Figure 2(a)). To transfer knowledge \(\mathcal{L}\) (especially \(\mathcal{L}_{D}\)) from pre-training domain to downstream domain, the well-trained encoder \(g_{L}\) is used to initialize the downstream encoder. In fine-tuning, downstream encoder is trained based on the task-specific knowledge \(\mathcal{D}\) preserved in a small amount of data examples \(({\bf x}^{d},y^{d})\) from downstream domain. Finally, the well-trained downstream encoder \(g_{L_{D}}\) is evaluated on test data.

### Knowledge-Aware Compression

#### 3.2.1 Two Basic Pruning Strategies

Intuitively, there are two pruning strategies. One is that pruning is applied to the downstream encoder \(g_{L}\) during fine-tuning (see Figure 1(b)). However, because the loss to update the weights during fine-tuning is exclusively based on the data examples \(({\bf x}^{d},y^{d})\) from the downstream task domain, this pruning strategy might destruct the knowledge \(\mathcal{L}_{D}\), which is learned based on \(({\bf x}^{p},y^{p})\) and encoded in the initialization of \(g_{L}\).

The other strategy is that pruning is executed during pre-training (see Figure 1(c)). The generated pruned network preserves a subset of knowledge \(\mathcal{L}\), denoted by \(\mathcal{L}^{pr}\). Unfortunately, because this strategy ignores the downstream task information and the amount of \(\mathcal{L}\) is extremely large, i.e., \(\mathcal{L}\gg\mathcal{L}^{pr}\), the knowledge \(\mathcal{L}^{pr}\) could be much different from \(\mathcal{L}_{D}\) that we hope to preserve (see Figure 2(a)).

#### 3.2.2 The Proposed Pruning Strategy

As shown in Figure 1(d), SparseBERT executes pruning at the distilling stage. It prunes the pre-trained encoder without fine-tuning, \(g_{L}\), while finetuning the pruned encoder based on the downstream dataset \(({\bf x}^{d},y^{d})\). Recent findings indicate that self-attention and feed-forward layers are over-parameterized and are the most computation consumption parts Michel et al. (2019); Voita et al. (2019); Ganesh et al. (2020). Thus, SparseBERT applies network pruning to the linear transformations matrices in self-attention and feed-forward layers (see Figure 3). The choice of pruning approach is flexible. We choose magnitude weight pruning Han et al. (2015) in this paper, mainly because it is one of the most effective and popular pruning methods. More details about the pruning strategy used in SparseBERT can be found in the codes.

#### 3.2.3 Knowledge Distillation Helps Pruning Preserve Task-Specific Knowledge

To mitigate the loss of \(\mathcal{L}_{D}\), we propose to utilize knowledge distillation while pruning. We use the task-specific fine-tuned BERT as the teacher network and the pre-trained BERT that is being pruned as the student (see Figure 1(d) and Figure 3). The motivation is that the task-specific fine-tuned BERT preserves \(\mathcal{L}_{\mathcal{D}}\). By feeding downstream task data \(({\bf x}^{d},y^{d})\) into the teacher-student framework, we help the student reproduce the behaviors of the teacher to learn both \(\mathcal{L}_{d}\) and \(\mathcal{L}\) as much as possible.

We design the distillation loss as

\[L_{distill}=L_{emb}+L_{att}+L_{hid}+L_{prd}. \tag{1}\]

\(L_{emb}\) **= MSE(E\({}^{S},{\bf E}^{T}\))** is the difference between the embedding layers of student and teacher. \(L_{att}\) = \(\sum\)MSE(A\({}^{S}_{i},{\bf A}^{T}_{i}\)) is the difference between attention matrices and \(i\) is the layer index. \(L_{hid}\) = \(\sum\)MSE(H\({}^{S}_{i},{\bf H}^{T}_{i}\)) is the difference between hidden representations. \(L_{prd}\) = -softmax(z\({}^{T}\)) \(\cdot\) log_softmax(z\({}^{S}\)/temp) is the soft cross-entropy loss between the logits of student and teacher. temp represents the temperature value. The proposed distillation loss is inspired by Jiao et al. (2020) and it helps the student imitate the teacher's behavior as much as possible. In addition, we perform the same data augmentation as Jiao et al. (2020) does to generate more task-specific data for teacher-student

learning. Notably, the choices of distillation loss and data augmentation method are flexible and we found the ones we adopted worked well in general.

## 4 Experiments

### GLUE Benchmark

We evaluate SparseBERT on four data sets from the GLUE benchmark Wang et al. (2018). To test if SparseBERT is applicable across tasks, we include the tasks of both single sentence and sentence-pair classification. We report the results on dev sets. We run 3, 20, 20, 50 epochs for QNLI, MRPC, RTE, CoLA separately. The baselines include BERTbase, ELMo Peters et al. (2018), BERT-PKD Sun et al. (2019), Bert-of-Theseus Xu et al. (2020), DistilBERT Sanh et al. (2019), MiniLM Wang et al. (2020), TinyBERT Jiao et al. (2020), BERTTickets Chen et al. (2020), CompressBERT Gordon et al. (2020), and RPP Guo et al. (2019).

The results are shown in Table 1. Compared to BERT-base, SparseBERT achieves 1.4% averaged performance loss with down to 5% weights. In addition, SparseBERT outperforms all leading competitors with the highest sparsity.

### SparseBERT v.s. Pruning at Downstream

We compare SparseBERT with the pruning described in Figure 1(b) on the question answer tasks of SQuAD v1.1 and v2.0 Rajpurkar et al. (2016, 2018). Given a question and a passage containing the answer, the two tasks are to predict the answer text span in the passage. The difference between them is that SQuAD v2.0 allows for the possibility that no short answer exists in the passage. We follow the general setting of SparseBERT, except that we only apply the logit distillation, i.e., \(L_{distil}\) = \(L_{prd}\), and do not perform data augmentation, which are the most common distillation strategies.

The results are shown in Figure 4. It is observed that SparseBERT consistently outperforms the baseline method, especially at high sparsity. The performance gain of SparseBERT decreases on SQuAD v2.0 mainly because SQuAD v2.0 is more challenging than SQuAD v1.1. These observations demonstrate advantage of SparseBERT compared to pruning at downstream.

Figure 3: Illustration of the proposed knowledge-aware compression. Pruning is performed in parallel with distillation, based on specific data from downstream tasks.

 



### SparseBERT v.s. Pruning at Pre-Training

To get more insights about the advantage of SparseBERT over the pruning described in Figure 1(c), we compare their fitting abilities. Specifically, we use TinyBERT as an example of the baseline pruning method. We compare SparseBERT with TinyBERT with 4 layers and 312 hidden dimensions, which has a similar number of parameters as SparseBERT (sparsity=95%). SparseBERT only distills knowledge from the same layers as TinyBERT does.

We vary the number of pruning epochs and report the results (loss on training set and accuracy on dev set) on RTE in Figure 5. It is observed that SparseBERT consistently shows smaller training loss while higher evaluation performance, which demonstrates that SparseBERT has a better fitting ability when pruning compared to the baseline.

## 5 Discussion

### Hardware Performance

Sparse networks were not hardware-friendly in the past. However, hardware platforms with sparse tensor operation support have been rising up. For example, the latest release of Nvidia high-end GPU A100 has native support of sparse tensor operation up to 2x compression rate, while startup company such as Moffett AI has developed computing platform with sparse tensor operation acceleration up to 32x compression rate.

Here we deployed SparseBERT of different sparse compression ratios (1, 2, 4, 8, 16, 20) on Moffett AI's latest hardware platform ANTOM to measure the real inference speedup induced by sparse compression, where '4' indicates the model is compressed by a factor of 4, with 75% of the parameters being zeros. As shown in Figure 6, the sparse compression has almost linear speedup up to 4x and leads to more than 10x speedup when compression rate is 20x.

### Reduction of Parameters and FLOPS

We studied the reduction of parameters and FLOPS. For example, on the MRPC dataset, BERT-base (backbone) vs SparseBERT (backbone) = 85.53 vs 4.84 (#parameters, M) and BERT-base vs SparseBERT = 10.87 vs 0.54 (GFLOPS).

### Inference/Training Time

We studied the time and convergence speed. For example, to get the reported 20x pruned result (Table 1), it needed 12 epochs of fine-tuning on MRPC and each epoch took 1.5 h (two RTX 2080 Ti). The inference time was around 20 s.

## 6 Conclusion

We introduce SparseBERT, a knowledge-aware sparse pruning method for pre-trained language models, with a focus on BERT. We summarize different types of BERT pruning approaches and compare SparseBERT with leading competitors. Experimental results on GLUE and SQuAD benchmarks demonstrate the superiority of SparseBERT.

## 7 Acknowledgements

We thank Xiaoqi Jiao for his valuable discussion and feedback on this work.

Figure 4: Performance comparison of SparseBERT and the pruning approach described in Figure 1(b).

Figure 5: Fitting ability comparison of SparseBERT and the pruning approach described in Figure 1(c).

Figure 6: Hardware performance under different compression ratios on the MRPC dataset, with 818, 1594, 3029, 5508, 9326, and 10826 SPS (sentences per second) respectively.





## References

* Chen et al. (2020) Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. 2020. The lottery ticket hypothesis for pre-trained bert networks. _NeurIPS_, 33.
* Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In _NAACL_, pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
* Elsen et al. (2020) Erich Elsen, Marat Dukhan, Trevor Gale, and Karen Simonyan. 2020. Fast sparse convnets. In _CVPR_, pages 14629-14638.
* Fan et al. (2019) Angela Fan, Edouard Grave, and Armand Joulin. 2019. Reducing transformer depth on demand with structured dropout. _arXiv preprint arXiv:1909.11556_.
* Frankle and Carbin (2019) Jonathan Frankle and Michael Carbin. 2019. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In _ICLR_.
* Gale et al. (2019) Trevor Gale, Erich Elsen, and Sara Hooker. 2019. The state of sparsity in deep neural networks. _arXiv preprint arXiv:1902.09574_.
* Ganesh et al. (2020) Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Deming Chen, Marianne Winslett, Hassan Sajjad, and Preslav Nakov. 2020. Compressing large-scale transformer-based models: A case study on bert. _arXiv preprint arXiv:2002.11985_.
* Gordon et al. (2020) Mitchell A Gordon, Kevin Duh, and Nicholas Andrews. 2020. Compressing bert: Studying the effects of weight pruning on transfer learning. _arXiv preprint arXiv:2002.08307_.
* Guo et al. (2019) Fu-Ming Guo, Sijia Liu, Finlay S Mungall, Xue Lin, and Yanzhi Wang. 2019. Reweighted proximal pruning for large-scale language representation. _arXiv preprint arXiv:1909.12486_.
* Han et al. (2015) Song Han, Jeff Pool, John Tran, and William Dally. 2015. Learning both weights and connections for efficient neural network. In _NeurIPS_, pages 1135-1143.
* Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. _arXiv preprint arXiv:1503.02531_.
* Jiao et al. (2020) Xiaoqi Jiao, Yichen Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. Tinybert: Distilling bert for natural language understanding.
* McCarley et al. (2020) J. S. McCarley, Rishav Chakravarti, and Avirup Sil. 2020. Structured pruning of a bert-based question answering model.
* Michel et al. (2019) Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? In _NeurIPS_, pages 14014-14024.
* Peters et al. (2018) Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. 2018. Deep contextualized word representations. In _NAACL_, pages 2227-2237, New Orleans, Louisiana. Association for Computational Linguistics.
* Rajpurkar et al. (2018) Pranav Rajpurkar, Robin Jia, and Percy Liang. 2018. Know what you don't know: Unanswerable questions for SQuAD. In _ACL_, pages 784-789, Melbourne, Australia. Association for Computational Linguistics.
* Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. SQuAD: 100,000+ questions for machine comprehension of text. In _EMNLP_, pages 2383-2392, Austin, Texas. Association for Computational Linguistics.
* Renda et al. (2020) Alex Renda, Jonathan Frankle, and Michael Carbin. 2020. Comparing rewinding and fine-tuning in neural network pruning. In _ICLR_.
* Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. DistilBERT, a distilled version of bert: smaller, faster, cheaper and lighter. _arXiv preprint arXiv:1910.01108_.
* Strubell et al. (2019) Emma Strubell, Ananya Ganesh, and Andrew McCallum. 2019. Energy and policy considerations for deep learning in NLP. In _ACL_, pages 3645-3650, Florence, Italy. Association for Computational Linguistics.
* Sun et al. (2019) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for BERT model compression. In _EMNLP_, pages 4323-4332, Hong Kong, China. Association for Computational Linguistics.
* Voita et al. (2019) Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. 2019. Analyzing multihead self-attention: Specialized heads do the heavy lifting, the rest can be pruned. _arXiv preprint arXiv:1905.09418_.
* Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_, pages 353-355, Brussels, Belgium. Association for Computational Linguistics.
* Wang et al. (2020) Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. 2020. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. _arXiv preprint arXiv:2002.10957_.
* Wang et al. (2019) Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2019. Structured pruning of large language models. _arXiv preprint arXiv:1910.04732_

Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. 2020. Bert-of-theseus: Compressing bert by progressive module replacing. _arXiv preprint arXiv:2002.02925_.
* Zafrir et al. (2019) Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. 2019. Q8bert: Quantized 8bit bert. _arXiv preprint arXiv:1910.06188_.
* Zhu and Gupta (2017) Michael Zhu and Suyog Gupta. 2017. To prune, or not to prune: exploring the efficacy of pruning for model compression. _arXiv preprint arXiv:1710.01878_.