K ronecker BERT : Significant Compression of Pre-trained Language Models Through Kronecker Decomposition and Knowledge Distillation (2022)

Abstract: The development of over-parameterized pre-trained language models has made a significant contribution toward the success of natural language processing. While over-parameterization of these models is the key to their generalization power, it makes them unsuitable for deployment on low-capacity devices. We push the limits of state-of-the-art Transformer-based pre-trained language model compression using Kronecker decomposition. We present our KroneckerBERT, a compressed version of the BERT_BASE model obtained by compressing the embedding layer and the linear mappings in the multi-head attention, and the feed-forward network modules in the Transformer layers. Our KroneckerBERT is trained via a very efficient two-stage knowledge distillation scheme using 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 factors of 7.7x and 21x 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.

## introduction in recent years, the emergence of pre-trained language models (plms) has led to a significant breakthrough 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 @xcite @xcite are powerful yet highly overparameterized. the enormous size of these models does not meet the constraints imposed by edge devices on memory, latency, and energy consumption. therefore there has been a growing interest in developing 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 @xcite , network pruning @xcite , matrix decomposition @xcite and knowledge distillation (kd) @xcite . 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 certain degree @xcite @xcite @xcite @xcite ; however, moderate and extreme compression of these models (compression @xmath0. in particular, several works @xcite @xcite that have tried to go beyond the compression factor of 10, have done so at the expense of a significant drop in performance. following the classical assumption that matrices often follow a low-rank structure, low-rank decomposition methods have been used for compression of weight matrices in deep learning models @xcite @xcite and especially transformer-based models @xcite . however, low-rank decomposition methods only exploit redundancies of the weight matrix in the horizontal and vertical dimensions and thus limit the flexibility of the compressed model. kronecker decomposition on the other hand exploits redundancies in predefined patches and hence allows for more flexibility in their representation. recent works prove kronecker product to be more effective in retaining accuracy after compression than svd @xcite . this work proposes a novel framework that uses kronecker decomposition for compression of transformer-based plms and provides a very promising compression-performance trade-off for medium and high compression levels, with 13% and 5% of the original model parameters respectively. we use kronecker decomposition for the compression of both transformer layers and the embedding layer. for transformer layers, the compression is achieved by representing every weight matrix both in the multi-head attention (mha) and the feed-forward neural network (ffn) as a kronecker product of two smaller matrices. we also propose a kronecker decomposition for compression of the embedding layer. previous works have tried different techniques to reduce the enormous memory consumption of this layer @xcite . our kronecker decomposition method can substantially reduce the amount of required memory while maintaining low computation. using kronecker decomposition for large compression factors leads to a reduction in the model expressiveness. this is due to the nature of the kronecker product and the fact that elements in this representation are tied together. to address this issue, we propose to distill knowledge from the intermediate layers of the original uncompressed network to the kronecker network during training. training of the state-of-the art bert compression models @xcite @xcite involve an extensive training which requires vast computational resources. for example in @xcite , first a specially designed teacher, i.e ib-bert large is trained from scratch on the en-tire english wikipedia and book corpus. the student is then pretrained on the same corpus via kd while undergoing an additional progressive kd phase. another example is tinybert @xcite which requires pretraining on the entire english wikipedia and also uses extensive data augmentation (20×) for fine-tuning on the downstream tasks. we show that our kronecker bert can out perform state-of-the-art with significantly less training requirements. more precisely, our kronecker-bert model undergoes a very light pretraining on only 10% of the english wikipedia for 3 epochs followed by finetuning on the original downstream data. note that, while our evaluations in this work are limited to bert, this proposed compression method can be directly used to compress other transformer-based nlp models. the main contributions of this paper are as follows: in this section, we first go through some of the most related works for bert compression in the literature and then review the few works that have used kronecker decomposition for compression of cnns and rnns. ## methodology in this section, we first introduce the background of kronecker decomposition and then explain our figure 2 : illustration of our proposed method for the compression of the embedding layer. left: conventional embedding stored in a lookup table. right: our proposed compression method where the original embedding matrix is represented as a kronecker product of a matrix and a row vector. the matrix is stored in a lookup table to minimize computation overhead. ## experiments in this section, we compare our kroneckerbert with the sate-of-the-art compression methods applied to bert on glue and squad. we also perform an ablation study to investigate the effect of pretraining and kd. ## conclusion we introduced a novel method for compressing transformer-based language models that uses kronecker decomposition for the compression of the embedding layer and the linear mappings within the transformer blocks. the proposed framework was used to compress the bert base model. we used a very light two-stage kd method to train the compressed model. we show that the proposed framework can significantly reduce the size and the number of computations while outperforming stateof-the-art. the proposed method can be directly applied for compression of other transformer-based language models. the combination of the proposed method with other compression techniques such layer truncation, pruning and quantization can be an interesting direction for future work.