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: This paper proposes to use tensor decomposition to compress the multi-head attention (MHA) and FFN layers in transformer architecture.   Premise:  Drawing from previously published research and their own experiments using PCA, the paper shows that the MHA and FFN layers are over-parameterized and exist in a lower dimensional subspace. Further, the authors talk about where these redundancy might come from and point to the decomposability of FFN and MHA layers and how that can lead to different parts learning similar behavior. Using PCA the authors show that these layers exhibit both inter-matrix and intra-matrix redundancy. Inter-matrix redundancy calls for low dimensional representation of each matrix in the layer, while intra-matrix redundancy implies possibility of parameter sharing between matrices.  Proposed Solution:  Based on this observation, the authors discuss the merit of various methods for decomposition of MHA and FFN matrices. Specifically, they discuss matrix decomposition, tensor train decomposition and tucker decomposition. The last technique will have the largest impact on compression and also has the most marginal cost of adding a new layer.   Results:  1. Impressive accuracy/parameter pareto frontier. Can compress BERT-base by ~50x with 1.5 to 2% loss in accuracy (Table 4) 2. Ablations on the need for 2 stage knowledge distillation - stage 1 to ensure attention maps and output of last layer are similar and stage two for task distillation