Compressing Pre-trained Language Models by Matrix Decomposition (2020)

Abstract: 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 @xcite . 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 @xcite @xcite @xcite @xcite @xcite @xcite . 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 @xcite @xcite . 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 @xcite 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. ## related work in the past year, there have been many attempts to compress transformer models involving pruning @xcite @xcite @xcite , quantization @xcite and distillation @xcite @xcite @xcite @xcite . specifically, works on compressing pretrained transformer language models focused on pruning layers. @xcite suggested to prune layers while distilling information from the unpruned model layers. @xcite proposed to gradually remove layers during training. we also note that very recently a work similar to ours was uploaded to arxiv @xcite . 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 @xmath0, obtaining an approximated matrix @xmath1, @xmath2, @xmath3, where @xmath4. we seek a decomposition s.t. w is close to w in the sense that d(w x, w x) 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 , @xcite ). 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: l @xmath5, where p s is the probability for label y given by the decomposed student model. ## 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 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 @xmath9. bert-pkd trains a six layered bert model with @xmath10, 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; foot_1 , they are task-specific; foot_2 and they do not require the pretraining stage, which is expensive and not practical for most practitioners. ## 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.