Gradient-based Intra-attention Pruning on Pre-trained Language Models

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 2775–2790 July 9-14, 2023 ©2023 Association for Computational Linguistics Gradient-based Intra-attention Pruning on Pre-trained Language Models Ziqing Yang†, Yiming Cui‡†, Xin Yao†, Shijin Wang†§ †State Key Laboratory of Cognitive Intelligence, iFLYTEK Research, Beijing, China ‡Research Center for SCIR, Harbin Institute of Technology, Harbin, China §iFLYTEK AI Research (Central China), Wuhan, China †{zqyang5,ymcui,xinyao10,sjwang3}@iflytek.com ‡ymcui@ir.hit.edu.cn Abstract Pre-trained language models achieve superior performance but are computationally expensive. Techniques such as pruning and knowledge dis- tillation have been developed to reduce their sizes and latencies. In this work, we propose a structured pruning method GRAIN (Gradient- based Intra-attention pruning), which performs task-specific pruning with knowledge distilla- tion and yields highly effective models. Differ- ent from common approaches that prune each attention head as a whole, GRAIN inspects and prunes intra-attention structures, which greatly expands the structure search space and enables more flexible models. We also propose a gra- dient separation strategy that reduces the inter- ference of distillation on pruning for a better combination of the two approaches. Experi- ments on GLUE, SQuAD, and CoNLL 2003 show that GRAIN notably outperforms other methods, especially in the high sparsity regime, and achieves 6 ∼ 7× speedups while main- taining 93% ∼ 99% performance. Under ex- treme compression where only 3% transformer weights remain, the pruned model is still com- petitive compared to larger models.1 1 Introduction Transformer-based (Vaswani et al., 2017) pre- trained language models (PLMs) have achieved great success and become the backbones of various natural language processing tasks. However, PLMs are computationally expensive and slow in infer- ence due to their large sizes, which limits their ap- plications in real-world scenarios. Hence, a grow- ing interest has been in developing compression and acceleration methodologies for PLMs. A common approach to model compression is structured pruning, which compresses the model by removing groups of consecutive parameters, namely the pruning units. In applying structured 1Code is available at https://github.com/airaria/ GRAIN. 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 22.5 Model Size (M) 76 78 80 82 84 86 Accuracy 5% 10% 20% Teacher MNLI-m GRAIN CoFi Block Pruning AutoTinyBERT DynaBERT TinyBERT MobileBERT MobileBERT-Tiny Figure 1: A comparison of GRAIN and other distillation and pruning methods on MNLI-m development set at different model sizes. More details are in Section 5. pruning on PLMs, recent works have investigated removing units such as hidden dimensions in feed- forward layers, attention heads in the multi-head attention (Michel et al., 2019; Li et al., 2022), and coarse-grained units such as multi-head attention layers and feed-forward layers (Xia et al., 2022). However, these pruning units only span a small space of model structures and limit the exploration for better structures. For example, in the pruning of BERTbase (Devlin et al., 2019), which contains 144 attention heads, the possible choices of attention heads for the pruned model are limited. Block Pruning (Lagunas et al., 2021) extends pruning units by considering blocks in the weight matrices, …

Content truncated to 500 words