Exploring extreme parameter compression for pre-trained language models

1 INTRODUCTION . Pre-trained Language Models such as BERT ( Devlin et al. , 2018 ) and ALBERT ( Lan et al. , 2019 ) have significantly improved various NLP tasks with significant improvement . Much recent work ( Brown et al. , 2020 ; Narayanan et al. , 2021 ; Fedus et al. , 2021 ) explores the potential of super large-scale PLMs . However , such large-scale PLMs are both economically and ecologically unfriendly ( Bender et al. , 2021 ; Patterson et al. , 2021 ) . Furthermore , deployment of large-scale PLMs is also challenging since ( 1 ) a model can not be fully deployed or stored in a single GPU server , model parallelism would consume extra time for network communication among many servers ; ( 2 ) edge devices may not have enough space for storing models ; ( 3 ) the long inference time can not support real-time feedback . Scaling down a model with negligible performance drop would facilitate the real-world applications of PLMs in a smaller size , faster inference time , and less network communication cost . For example , recent work explores quantization ( Zhang et al. , 2020 ; Bai et al. , 2020 ) , weights pruning ( Hou et al. , 2020 ) , and knowledge distillation ( Jiao et al. , 2020 ; Sanh et al. , 2020 ) for BERT ( one of the most popular PLMs ) . We argue that existing methods can not largely compress large-scale PLMs as stated in Sec . 2 . In this paper , we aim to explore extreme parameter compression ( i.e. , bigger compression ratios ) although they are by definition challenging . The parameter redundancy in PLMs was demonstrated by ( Kovaleva et al. , 2019 ; Michel et al. , 2019 ; Voita et al. , 2019 ; Cordonnier et al. , 2021 ) , for which we divide into two groups : intramatrix redundancy and inter-matrix redundancy . The former happens in different heads that are ∗Benyou and Yuxin contributed to this work equally . †Lifeng is the corresponding author . 1https : //github.com/twinkle0331/Xcompression calculated separately , e.g. , attentions among heads act on a similar subspace and are therefore lowrank ( Cordonnier et al. , 2021 ) – we relate this phenomenon to the so-called ‘ decomposability ’ defined in this paper . Like self-attention layers , decomposability also holds in FFN layers – each FFN layer could be decomposed to many independent sub-FFNs ( as explained in Appendix B ) . One example of inter-matrix redundancy happens across different layers , e.g. , attention maps among layers might be similar ( Clark et al. , 2019 ; Vig , 2019 ; Rogers et al. , 2020 ) . Exploration of main weight matrices in Transformer layers finds that these weight matrices are possible to be approximated in a low-rank manner – evidencing the possible intra-matrix redundancy and inter-matrix redundancy . We comprehensively analyze and compare different decomposition methods for parameter compression including matrix decomposition ( denoted as II ) , tensor train decomposition ( Oseledets , 2011 ) ( denoted as III ) and Tucker decomposition ( De Lathauwer et al. , 2000 ) ( denoted as IV ) . The fundamental difference between them is as below . II conducts matrix factorization ( e.g. , SVD ) for each weight matrix thanks to intra-matrix redundancy . Regarding intermatrix redundancy , III shares the head and tail matrices while keeping the core matrix individual ; IV introduces ‘ matrix bank ’ to make parameter scale being nearly constant w.r.t . the number of layers . It is concluded that Tucker decomposition ( IV ) is more parameter-efficient than others in terms of compression ratios . ALBERT ( Lan et al. , 2019 ) and III can be considered as special cases of IV . The practical challenges of matrix/tensor decomposition for compression are twofold . First , the decomposition may result in a discrepancy between the raw weights and approximated weights , and exact decomposition is impossible with large compression ratios . Instead , Knowledge Distillation ( KD ) is used on the compressed model to simulate the predictions of the raw model in a loss-aware manner . Second , reconstruction may lead to additional computation costs . An efficient reconstruction protocol is implemented by reordering multiplication operations that also preserve the same results . The contributions of this work are ( 1 ) we propose a formal framework with standardized terminology to comprehensively discuss matrix/tensor decomposition methods to compress Transformer-based language models ; ( 2 ) we adopt tensor decomposition for compressing PLMs which is also faster , while existing work ( Ma et al. , 2019 ; Liu et al. , 2021 ) did not show the potential for speedup in PLMs ; ( 3 ) our compressed BERT with 1/7 parameters in Transformer layers performs on-par with the original BERT in GLUE benchmark . Also , a tiny version achieves 96.7 % performance of BERT-base with only 1/48 parameters in Transformer layers and 2.7× faster on inference . We directly use the proposed methods on TinyBERT ( Jiao et al. , 2020 ) that is purely based on KD , since our work is complementary to existing compression methods like KD . 2 RELATED WORK . Compressing PLMs Although various work was proposed to design a new efficient Transformer ( Tay et al. , 2020 ) , e.g. , ( Ma et al. , 2019 ; Choromanski et al. , 2020 ; Wang et al. , 2020 ; Kitaev et al. , 2020 ; Zaheer et al. , 2020 ; Cao et al. , 2020 ) , in this paper , we are focusing on the compression of Transformer-based pre-trained language models . The difference is that the latter expects to reuse well-trained models , e.g. , BERT and GPT ( and even GPT3 ( Radford et al. , 2019 ) and PanGu-α ( Zeng et al. , 2021 ) ) with manageable computing resources , which typically does not change the original Transformer architecture . Taking BERT , one of the most commonly-used pre-trained language models , as an example . Existing work explored quantization ( Zhang et al. , 2020 ; Bai et al. , 2020 ) , weights pruning ( Lagunas et al. , 2021 ) , knowledge distillation ( Jiao et al. , 2020 ; Sanh et al. , 2020 ) , progressive module replacing ( Xu et al. , 2020 ) , neural architecture search ( Xu et al. , 2021 ; Yin et al. , 2021 ) and matrix decomposition ( Noach & Goldberg , 2020 ) . We argue that existing compression methods ( see Tab . 1 ) may be inadequate for extreme parameter compression , which is under-investigated . The reasons are manifold , first , the knowledge distillationbased method generally learns a new student model from scratch , which can not inherit too much knowledge from the teacher model before distillation . Second , some methods have upper bounds of compression ratio . For example , layer-sharing ALBERT ( Lan et al. , 2019 ) shares parameters in L layers with maximum L times compression . Quantization replaces existing 32-bit parameters with binary parameters with a maximum of 32 times reduction . Moreover , quantization needs further hardware support , which is usually ad hoc to specific platforms . Weight pruning arguably can not achieve a big compression ratio ( McCarley et al. , 2019 ; Michel et al. , 2019 ) . Matrix/tensor decomposition for compression Tensor/matrix decomposition aims to approximate a given tensor using a set of smaller tensors/matrices . It has been investigated to compress and speed up CNNs , RNNs , and Transformers for many years ( Lebedev et al. , 2014 ; Yang et al. , 2017 ; Ye et al. , 2018 ; Denton et al. , 2014 ; Winata et al. , 2019 ) . Matrix decomposition ( e.g. , ALBERT ( Lan et al. , 2019 ) in the embedding layer and ( Noach & Goldberg , 2020 ) ) could decrease parameter scale with a linear factor depending on the selected rank . More advanced tensor decomposition approaches can be implemented by tensor network , which has recently been used to compress general neural networks ( Gao et al. , 2020 ; Novikov et al. , 2015 ) , compress embedding layer ( Khrulkov et al. , 2019 ; Hrinchuk et al. , 2020 ; Panahi et al. , 2019 ) . Recently , Ma et al . ( 2019 ) redesigned a new ‘ Self-Attention Network ’ ( SAN ) in Transformer architecture inspired by block-term tensor decomposition . The compression ratio is limited since the majority of parameters in Transformer comes from another module called ‘ Feed-Forward Network ’ ( FFN ) instead of SAN ; moreover , the model does not have the potential for speedup regardless of compression ratios since the time complexity is closed to vanilla Transformer . Noach & Goldberg ( 2020 ) reparameterized each weight matrix using matrix decomposition and further distill the compressed models , which nearly achieves a compression ratio of 1.7 . Liu et al . ( 2021 ) adopted matrix product operators to reparameterize each group of weight matrices in embedding , SAN , and FFN , and only a small part of tensors of MPO ( called ‘ auxiliary tensors ’ ) are fine-tuned in downstream tasks . The compression ratio of the total parameter is negligible and the inference might be slow . Those works inspire us to explore in depth extreme parameter compression for large-scale PLMs . We argue that compression ratio of existing work using matrix/tensor decomposition ( Ma et al. , 2019 ; Liu et al. , 2021 ; Noach & Goldberg , 2020 ) for PLMs is relatively-small ; most of them do not have speedup effect , limiting their applications in large-scale PLMs . The potential to compress PLMs with matrix/tensor decomposition is under-investigated . In this work , we adopt tensor decomposition , to cubically compress the parameters of PLMs . 3 MOTIVATIONS FOR PARAMETER COMPRESSION . Pre-trained language models are typically a stack of multiple Transformer ( Vaswani et al. , 2017 ) layers that consist of a Self-Attention Network ( SAN ) module and a Feed-Forward Network ( FFN ) module , see App . A for Transformer . Sec . 3.1 will introduce an important property called ‘ decomposability ’ for SAN and FFN , which indicates each sub-component in SAN or FFN is independently calculated without interactions between sub-components and they may therefore be redundant . 3.1 DECOMPOSABILITY IN TRANSFORMER . A computing module f is decomposable if its sub-components { g1 , g2 , ⋯gH } could be independently calculated without interactions : f ( x ) = δ ( g1 ( x ) , g2 ( x ) , ⋯ , gH ( x ) ) . Usually , δ is a simple operation that has negligible computing cost compared to { gh } . Especially , backpropagation between sub-components is independent if δ is concatenation or addition . Sub-components in f could be calculated in parallel without interactions . We will examine if SAN and FFN are decomposable . Decomposability in SAN Following ( Hou et al. , 2020 ) , SAN could be decomposed as a sum of the output of every head . For the query/key/value/output transformations parameterized by W Q/WK/W V /WO , we divide them as Nh heads : { WQh } Nh , { WKh } Nh , { W Vh } Nh , { WOh } Nh . A single head is calculated as Atth ( X ) = Softmax ( 1√ d XWQh W K h T XT ) XW Vh WOh T ( 1 ) Then SAN ( X ) = ∑Nhh=1 Atth ( X ) , indicating SAN is decomposable . Cordonnier et al . ( 2021 ) argued that the attentions among heads learn redundant key/query projections due to independent calculation . Decomposability in FFN The two weight matrices in FFN , denoted as W In and WOut . FFN ( X ) = 4D ∑ h=1 GeLU ( XW In⋅ , h + bInh ) WOuth , ⋅ + bOut ( 2 ) As Geva et al . ( 2020 ) points out , FFN also operates as a key-value mechanism similar to SAN : one can consider the input of FFN as a query vector , and two linear layers of FFN as keys and values , respectively . There may exist some similar key-value pairs that may introduce redundancy . Remarks on decomposability Decomposability does not necessarily make sub-components being complementary , especially without orthogonality-like constraints between them . Sub-components may learn similar patterns , leading to some redundancy ( Michel et al. , 2019 ) . For example , Hou et al . ( 2020 ) demonstrated that BERT could be pruned in both width and height without performance drop . Moreover , BERT is over-parameterized in the sense that the number of training examples in downstream tasks is much fewer ( 393k training examples for the biggest task in GLUE , i.e. , MNLI ) comparing to the enormous number of parameters ( 110M parameters for BERT-base ) – this was thought problematic according to the traditional machine learning theory ( Vapnik , 2013 ) . We believe that the pre-trained language models in downstream tasks could be compressed .

Summary: After author rebuttal: thank you very much for the detailed response to both my and the other reviewers' comments. I have updated my score.  ------------ This paper proposes an approach to compressing the Transformer family of pretrained language models via tensor decomposition. Compared to existing work, the main differences are: (1) performing global tensor decomposition which takes into account matrices across different layers, and (2) using a matrix bank to enable even greater model compression. The approach, when combined with distillation, outperforms existing methods for compressing BERT models on the popular GLUE benchmark.