KroneckerBERT: Significant Compression of Pre-trained Language Models Through Kronecker Decomposition and Knowledge Distillation

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2116 - 2127 July 10-15, 2022 ©2022 Association for Computational Linguistics KroneckerBERT: Significant Compression of Pre-trained Language Models Through Kronecker Decomposition and Knowledge Distillation Marzieh S. Tahaei Huawei Noah’s Ark Lab marzieh.tahaei@huawei.com Ella Charlaix Huawei Noah’s Ark Lab charlaixe@gmail.com Vahid Partovi Nia Huawei Noah’s Ark Lab vahid.partovinia@huawei.com Ali Ghodsi Department of Statistics Actuarial Science, University of Waterloo ali.ghodsi@uwaterloo.com Mehdi Rezagholizadeh Huawei Noah’s Ark Lab Mehdi.rezagholizadeh@huawei.com Abstract The development of over-parameterized pre- trained language models has made a significant contribution toward the success of natural lan- guage processing. While over-parameterization of these models is the key to their generaliza- tion power, it makes them unsuitable for de- ployment on low-capacity devices. We push the limits of state-of-the-art Transformer-based pre-trained language model compression us- ing Kronecker decomposition. We present our KroneckerBERT, a compressed version of the BERTBASE model obtained by compressing the embedding layer and the linear mappings in the multi-head attention, and the feed-forward net- work modules in the Transformer layers. Our KroneckerBERT is trained via a very efficient two-stage knowledge distillation scheme us- ing far fewer data samples than state-of-the-art models like MobileBERT and TinyBERT. We evaluate the performance of KroneckerBERT on well-known NLP benchmarks. We show that our KroneckerBERT with compression fac- tors of 7.7× and 21× outperforms state-of-the- art compression methods on the GLUE and SQuAD benchmarks. In particular, using only 13% of the teacher model parameters, it retain more than 99% of the accuracy on the majority of GLUE tasks. 1 Introduction In recent years, the emergence of Pre-trained Lan- guage Models (PLMs) has led to a significant break- through in Natural Language Processing (NLP). The introduction of Transformers and unsupervised pre-training on enormous unlabeled data are the two main factors that contribute to this success. Transformer-based models (Devlin et al., 2018; Radford et al., 2019; Yang et al., 2019; Shoeybi et al., 2019) are powerful yet highly over- parameterized. The enormous size of these models does not meet the constraints imposed by edge de- vices on memory, latency, and energy consumption. Therefore there has been a growing interest in de- veloping new methodologies and frameworks for the compression of these large PLMs. Similar to other deep learning models, the main directions for the compression of these models include low-bit quantization (Gong et al., 2014; Prato et al., 2019), network pruning (Han et al., 2015), matrix decom- position (Yu et al., 2017; Lioutas et al., 2020) and Knowledge distillation (KD) (Hinton et al., 2015). These methods are either used in isolation or in combination to improve compression-performance trade-off. Recent works have been relatively successful in compressing Transformer-based PLMs to a cer- tain degree (Sanh et al., 2019; Sun et al., 2019; Jiao et al., 2019; Sun et al., 2020; Xu et al., 2020; Wang et al., 2020; Kim et al., 2021); however, mod- erate and extreme compression of these models (compression factors >5 and 10 resepctively) is …

Content truncated to 500 words