Structured Pruning Learns Compact and Accurate Models

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 1513 - 1528 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Structured Pruning Learns Compact and Accurate Models Mengzhou Xia Zexuan Zhong Princeton UniversityDanqi Chen {mengzhou,zzhong,danqic}@cs.princeton.edu Abstract The growing size of neural language models has led to increased attention in model com- pression. The two predominant approaches are pruning, which gradually removes weights from a pre-trained model, and distillation, which trains a smaller compact model to match a larger one. Pruning methods can significantly reduce the model size but hardly achieve large speedups as distillation. However, distillation methods require large amounts of unlabeled data and are expensive to train. In this work, we propose a task-specific structured pruning method CoFi1 (Coarse- and Fine-grained Prun- ing), which delivers highly parallelizable sub- networks and matches the distillation methods in both accuracy and latency, without resort- ing to any unlabeled data. Our key insight is to jointly prune coarse-grained (e.g., layers) and fine-grained (e.g., heads and hidden units) modules, which controls the pruning decision of each parameter with masks of different gran- ularity. We also devise a layerwise distillation strategy to transfer knowledge from unpruned to pruned models during optimization. Our experiments on GLUE and SQuAD datasets show that CoFi yields models with over 10× speedups with a small accuracy drop, show- ing its effectiveness and efficiency compared to previous pruning and distillation approaches.2 1 Introduction Pre-trained language models (Devlin et al., 2019; Liu et al., 2019a; Raffel et al., 2020, inter alia) have become the mainstay in natural language pro- cessing. These models have high costs in terms of storage, memory, and computation time and it has motivated a large body of work on model com- pression to make them smaller and faster to use in real-world applications (Ganesh et al., 2021). The 1CoFi is pronounced as . 2Our code and models are publicly available at https: //github.com/princeton-nlp/CoFiPruning. U T ↗ Params MNLI BERTbase (teacher) ✗ ✗ 1.0× 85M 84.8 Distillation DistillBERT6 ✓ ✗ 2.0× 43M 82.2 TinyBERT6 ✓ ✓ 2.0× 43M 84.0 MobileBERT‡ ✓ ✗ 2.3× 20M 83.9 DynaBERT ✗ ✓ 6.3× 11M 76.3 AutoTinyBERT‡ ✓ ✓ 9.1× 3.3M 78.2 TinyBERT4 ✓ ✓ 11.4× 4.7M 78.8 Pruning Movement Pruning ✗ ✓ 1.0× 9M 81.2 Block Pruning ✗ ✓ 2.7× 25M 83.7 CoFi Pruning (ours) ✗ ✓ 2.7× 26M 84.9 CoFi Pruning (ours) ✗ ✓ 12.1× 4.4M 80.6 Table 1: A comparison of state-of-the-art distillation and pruning methods. U and T denote whether Unlabeled and Task-specific are used for distillation or pruning. The inference speedups (↗) are reported against a BERTbase model and we evaluate all the models on an NVIDIA V100 GPU (§4.1). The models labeled as ‡ use a different teacher model and are not a direct com- parison. Models are one order of magnitude faster.3 two predominant approaches to model compression are pruning and distillation (Table 1). Pruning methods search for an accurate subnet- work in a larger pre-trained model. Recent work …

Content truncated to 500 words