Compressing Pre-trained Language Models by Matrix Decomposition

Large pre-trained language models reach state-of-the-art results on many different NLP tasks when fine-tuned individually; They also come with a significant memory and computational requirements, calling for methods to reduce model sizes (green AI). We propose a two-stage model-compression method to reduce a model’s inference time cost. We first decompose the matrices in the model into smaller matrices and then perform feature distillation on the internal representation to recover from the decomposition. This approach has the benefit of reducing the number of parameters while preserving much of the information within the model. We experimented on BERT-base model with the GLUE benchmark dataset and show that we can reduce the number of parameters by a factor of 0.4x, and increase inference speed by a factor of 1.45x, while maintaining a minimal loss in metric performance.


Introduction
Deep learning models have been demonstrated to achieve state-of-the-art results, but require large parameter storage and computation. It's estimated that training a Transformer model with a neural architecture search has a CO 2 emissions equivalent to nearly five times the lifetime emissions of the average U.S. car, including its manufacturing (Strubell et al., 2019). Alongside the increase in deep learning models complexity, in the NLP domain, there has been a shift in the NLP modeling paradigm from training a randomly initialized model to fine-tuning a large and computational heavy pre-trained language model (Howard and Ruder, 2018;Peters et al., 2018;Devlin et al., 2018;Radford, 2018;Radford et al., 2019;Lample and Conneau, 2019;Liu et al., 2019b;Raffel et al., 2019;Lan et al., 2019;Lewis et al., 2019).
While re-using pre-trained models offsets the training costs, inference time costs of the finetuned models remain significant, and are showstoppers in many applications. The main challenge with pre-trained models is how can we reduce their size while saving the information contained within them. Recent work, approached this by keeping some of the layers while removing others (Sanh et al., 2019;Sun et al., 2019;Xu et al., 2020). A main drawback of such approach is in its coarse-grained nature: removing entire layers might discard important information contained within the model, and working at the granularity of layers makes the trade-off between compression and accuracy of a model hard to control. Motivated by this, in this work we suggest a more finegrained approach which decomposes each matrix to two smaller matrices and then perform feature distillation on the internal representation to recover from the decomposition. This approach has the benefit of preserving much of the information while reducing the number of parameters. Alongside the advantage of preserving the information within each layer, there is also a memory flexibility advantage compared to removing entire layers; As a result of decomposing each matrix to two smaller matrices, we can store each of the two matrices in two different memory blocks. This has the benefit of distributing the model matrices in many small memory blocks, which is useful when working in shared CPU-based environments.
We evaluated our approach on the General Language Understanding Evaluation (GLUE) benchmark dataset (Wang et al., 2018) and show that our approach is superior or competitive in the different GLUE tasks to previous approaches which remove entire layers. Furthermore, we study the effects of different base models to decompose and show the superiority of decomposing a fine-tuned model compared to a pre-trained model or a ran-domly initialized model. Finally, we demonstrate the trade-off between compression and accuracy of a model.
We also note that very recently a work similar to ours was uploaded to arxiv (Mao et al., 2020). There are a few differences from their work to ours. Firstly, we distill different parts of the model (see Section 3 for details). Secondly, we focus on training the decomposed model and do not prune the model parameters. Thirdly, our base model, which is used for decomposition and as a teacher, is a fine-tuned model; This has the benefit of task-specific information as we show in our experiments in Section 4.2.

Method
Our goal is to decompose each matrix W ∈ R n×d as two smaller matrices, obtaining an approxi- is small for all x, where d is a distance metric between vectors. In practice, we require the condition to hold not for all x, but for vectors seen in a finite relevant sample (in our case, the training data). While one could start with random matrices and optimize the objective using gradient descent, we show that a two-staged approach performs better: we first decompose the matrices using SVD, obtaining A , B s.t. ||A B − W || 2 2 is small (SVD is guaranteed to produce the best rank-r approximation to W , (Stewart, 1991)). We then use these matrices as initialization and optimize d(W x, W x) (feature distillation), while also optimizing for task loss. We show that this process works substantially better in practice. Our loss function is thus composed of three different objectives: Cross Entropy Loss The cross entropy loss over an example x with label y is defined likewise: where p s is the probability for label y given by the decomposed student model.

Knowledge Distillation Loss
The goal of knowledge distillation is to imitate the output layer of a teacher model by a student model. The Knowledge Distillation Loss is defined likewise: where z s and z t are the logits of the decomposed and original models respectively and T is a temperature hyper-parameter.

Feature Distillation Loss
The goal of feature distillation is to imitate the intermediate layers of a teacher model by a student model. we use the following intermediate representations to distill the knowledge from 1 : • Query, Key and Value Layers -The dot product of a matrix of concatenated tokens representation vectors X by the query, key and value parameter matrices, • Attention Heads -The output of the attention heads.
• The Multihead Attention Layer Output -The dot product of the attention heads by the ma- • The first feed forward layer -The dot product of the multihead attention layer by the first feed forward layer.
• The second feed forward layer -The dot product of the first feed forward layer by the second feed forward layer.
We denote S i z and T i z as the intermediate representations which were described above of layer i for the decomposed student and original teacher models respectively. Our loss function then is defined by: Full Objective Our loss function is then defined by a weighted combination of these three loss functions likewise:

Experiments
We compare various variants of our compression method, corresponding to different subsets of our loss. All variants decompose the matrices using SVD, but differ in their objective functions. These correspond to the four last lines in Table  1. Low Rank BERT Fine-tuning (LRBF) corresponds to L = L CE . LRBF+KD corresponds to L = αL CE + (1 − α)L KD . LRBF+FD corresponds to L = L CE +L F D , while LRBF+FD+KD corresponds to the complete objective. The other lines in the table correspond to uncompressed model (first line) and to baselines which prune layers and distill. Fine-tuning finetunes a six layered BERT model. Vanilla KD trains a six-layered BERT model with L = αL CE +(1−α)L KD . BERT-PKD trains a six layered BERT model with L = αL CE + (1 − α)L KD while also adding an L F D objective, but on the hidden states between every consecutive layer. BERT-of-Theseus fine-tunes BERT model while gradually pruning half of the layers. We chose this baselines for several reasons: like our method they result in a practical reduction of parameters; 2 , they are task-specific; 3 and they do not require the pretraining stage, which is expensive and not practical for most practitioners.
Datasets We evaluate our proposed approach on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018), a collection of diverse NLP tasks.
Training Details We fine-tune a pre-trained BERT model (Devlin et al., 2018) for each task with a batch size of 8 and a learning rate of 2e−5 for 3 epochs with an early stop mechanism according to the validation set. We perform the matrix decomposition on every parametric weight matrix of the encoder (excluding the embedding matrix) in a fine-tuned model and train the decomposed model as the student model and the original fine-tuned model as the teacher. For each task we train for 3 epochs with an early stopping mechanism according to the task validation set, the maximum sequence length is 128 and we perform a grid search over the learning rates {2e −6 , 5e −6 , 2e −5 , 5e −5 , 2e −4 , 5e −4 } and 5 different seeds and choose the best model according to the validation set of each task. 4 For knowledge distillation hyper-parameters we used a temperature hyper-parameter T = 10 and α = 0.7. 5 Table 1 compares the results for validation and test of other compression approaches which prune layers, along with low rank models which were finetuned and trained with one or more of the distillation objectives described in Section 3. As can be seen, Low Rank BERT Feature Distillation + KD and Low Rank BERT Feature Distillation surpass all of results of all methods in both validation and test sets except BERT-of-Theseus method in the test set, in which Low Rank BERT Feature Distillation + KD surpasses the results in 5 of the tasks and reach comparable results in 2 of the tasks. Also, as can be seen knowledge distillation alone is not sufficient to compensate for the decomposition, but it slightly improves the results when incorporating feature distillation alone.

Further Analysis
Effect of Base Model and Decomposition In this experiment we test the importance of the base model we use to decompose and use as a teacher. We compared between three types of distillation sources: fine-tuned teacher, pre-trained teacher and no teacher. Furthermore, we compared between three types of model initializations: a decomposed fine-tuned model, a decomposed pretrained model and a randomly initialized model with the same architecture as the decomposed models. The results are shown in Table 2, on all tasks when training with no teacher distilla-  (Hinton et al., 2015) 45   tion, the results are best when decomposing a finetuned model and decomposing a pre-trained model is better than randomly initializing a model; This indicates that the decomposition saves the information within the model and when decomposing a fine-tuned model it saves some of the more task specific information. Furthermore, on all tasks and all initialization the best results are when using a fine-tuned model as a teacher.  Compression vs. Performance Trade-off Our method requires to determine a rank for the compression. But can we achieve better results when choosing a higher rank? Can we choose a lower rank for smaller models and still achieve satisfactory results? To determine this we experimented on three different ranks. As shown in Table 3, higher ranks achieve better results, while lower ranks achieve satisfactory results while compromising metric performance.

Run-time Savings
In this experiment we measured the average time in milliseconds it takes for BERT-base compared to its decomposed and six-layered counterparts to output predictions for a batch of samples with varying batch sizes. As shown in Figure 1, we still gain a significant time performance improvement when running on both CPU and GPU architectures over a BERT-base model. Models that are decomposed to a rank r = 245 are about 1.45 faster than their uncompressed counterpart for batches larger than one when running on a GPU and around 1.2 − 1.55 faster for batches 8, 16, 32, 64 when running on a CPU. Furthermore, higher ranks still benefit running time and lower ranks improve the running time further. Also, we note that although a six-layered BERT does achieve faster inference time, due to the coarse-grained compression, it losses more information contained within it and thus achieves inferior results; As shown in the results in Table 1, a six-layered model trained with distillation (e.g. BERT-PKD (Sun et al., 2019)) achieves significantly lower results and the BERT-of-Theseus model, which does improve upon BERT-PKD, requires many training iterations to achieve this to overcome the loss of information when gradually removing entire layers, which result in higher training times.

Conclusions
We presented a way to compress pre-trained large language models fine-tuned for specific tasks, while preserving much of the information contained within them, by using matrix decomposition to two small matrices. For future work it might be interesting to combine this approach with another approach such as pruning or quantization to achieve smaller models.