Data-aware Low-Rank Compression for Large NLP Models

1 INTRODUCTION . The representations learned by large-scale Natural Language Processing ( NLP ) models such as BERT have been widely used in various tasks ( Devlin et al. , 2018 ) . The pre-trained models of BERT and its variations are used as feature extractors for the downstream tasks such as question answering and natural language understanding ( Radford et al . ; Howard & Ruder , 2018 ) . The success of the pre-trained BERT relies on the usage of large corpus and big models . Indeed , researchers have reported better results of models with more parameters ( Shazeer et al. , 2018 ) and number of layers ( Al-Rfou et al. , 2019 ) . The increasing model size of the pre-trained models inhibits the public user from training a model from scratch , and it also brings the efficiency challenges , including the inference speed and the model size when deploying the model on devices . To deal with the efficiency issue , most existing works resort to adjusting the model structures or distillation . For instance , Kitaev et al . ( 2020 ) uses locality-sensitive hashing to accelerate dot-product attention , Lan et al . ( 2019 ) uses repeating model parameters to reduce the size and Zhang et al . ( 2018 ) applies a pre-defined attention pattern to save computation . A large body of prior work focuses on variants of distillation information has also been explored ( Sanh et al. , 2019 ; Jiao et al. , 2019 ; Sun et al. , 2020 ; Liu et al. , 2020 ; Xu et al. , 2020 ; Sun et al. , 2019 ) . However , all these methods either require a specific design of model architecture which is not generic , or require users to train the proposed structure from scratch which greatly reduces its practicality . In this work , we try to explore an acceleration method that is more generic . Note that as shown in Figure 1 , matrix multiplication ( Feed-forward layer ) is a fundamental operation which appears many times in the Transformer architecture . In fact , the underlying computation of both multi-head attention layers and feed-forward layers is matrix multiplication . Therefore , instead of resorting to the complex architecture redesign approaches , we aim to investigate whether low-rank matrix approximation , the most classical and simple model compression approach , can be used to accelerate Transformers . Despite being successfully applied to CNN ( Yu et al. , 2017 ; Sindhwani et al. , 2015 ; Shim et al. , 2017 ; You et al. , 2019 ) , at the first glance low-rank compression can not work for BERT . We could see in Figure 2 that regardless of layers , matrices in feed-forward layer , query and key transformation of attention layer are not low-rank . Therefore , even the optimal low-rank approximation ( e.g. , by SVD ) will lead to large reconstruction error and empirically the performance is limited . This is probably why low-rank approximation has not been used in BERT compression . In this paper , we propose a novel low-rank approximation algorithm to compress the weight matrices even though they are not low-rank . The main idea is to exploit the data distribution . In NLP applications , the latent features , indicating some information extracted from natural sentences , often lie in a subspace with a lower intrinsic dimension . Therefore , in most of the matrix-vector products , even though the weight matrices are not low-rank , the input vectors lie in a low-dimensional subspace , allowing dimension reduction with minimal degraded performance . We mathematically formulate this generalized low-rank approximation problem which includes the data distribution term and provide a closed-form solution for the optimal rank-k decomposition of the weight matrices . We propose DRONE method based on this novel Data-awaRe lOw-raNk comprEssion idea . Our decomposition significantly outperforms SVD under the same rank constraint , and can successfully accelerate the BERT model without sacrificing too much test performance . 2 RELATED WORK . The inference speed is important for NLP models when deployed in various applications . Generally speaking , inference efficiency could be enhanced by hardware ( Shawahna et al. , 2018 ) or lower-level instruction optimization ( Ning , 2020 ) . On the other hand , the main focus of the current research is on using algorithmic methods to reduce the computational complexity . It could be mainly categorized into two aspects : Attention Complexity Reduction and Model Size Reduction . Attention Complexity Reduction Attention mechanism is the building block of transformer model and it attracts most attentions of researcher recently in NLP field ( Vaswani et al. , 2017 ) . Pre-training on large courpus of BERT , a transformer-based model , has contributed to state-of-the-art performance on various tasks after fine-tuning ( Devlin et al. , 2018 ) . Attention on sequences of length L is O ( L2 ) in both computational and memory complexity . This would take long inference time when the sequence is long . Thus , researchers have focused on reducing the complexity of the attention module . Kitaev et al . ( 2020 ) uses uses locality-sensitive hashing to reduce the complexity to O ( LlogL ) . Zhang et al . ( 2018 ) ; Child et al . ( 2019 ) pre-defined an attention map to have a constant computational time . Goyal et al . progressively eliminates the redundant context vectors within the attended sequence to improve the efficiency of attention in last few layers of the model . Wang et al . ( 2020 ) proposes to train the low-rank attention by choosing a rank r L. This is similar with our work in the sense of leveraging low-rank structures . But our method do not require retraining the model and could be applied to different modules other than attention . In fact , most of these methods require special modules and thus we need to retrain the proposed models from scratch . This prohibits the usage of a large body of publicly available open models for faster research progress . More importantly , these methods only focus on the long sequence scenario . We found out that attention module is actually not the bottleneck of inference time in common usage as shown in Figure 1 . In most if not all models of common usages , two layers of large feed-forward layer are appended after attention module which incurs much more computational time . Attention complexity reduction only works when a large sequence is used but in current practice this is unusual . Thus , accelerating attention module itself is not contributing to a significant reduction of overall inference time . Model Size Reduction Efficiency of inference time is also related to model compression . In principle , smaller models would lead to smaller number of operations and thus faster inference time . Sanh et al . ( 2020 ) has explored pruning methods on BERT models to eliminate the redundant links , and there is a line of research on pruning methods ( Han et al. , 2015a ; b ; Chen et al. , 2020 ) . Quantization methods ( Zafrir et al. , 2019 ; Hubara et al. , 2016 ; Lin et al. , 2016 ) could convert the 32 bits float models into lower-bits fixed-point representation , and theoretically make model prediction faster with fixed point accelerator . Lan et al . ( 2019 ) reduces the model size by sharing of encoder parameters . A large body of prior work focuses on variants of knowledge distillation ( Sanh et al. , 2019 ; Jiao et al. , 2019 ; Sun et al. , 2020 ; Liu et al. , 2020 ; Xu et al. , 2020 ; Sun et al. , 2019 ; 2020 ) . These methods use different strategies to distill information from teacher network and reduce the number of layers ( Sanh et al. , 2019 ) or hidden dimension size ( Jiao et al. , 2019 ) . Further , A hybrid compression method by combining matrix factorization , pruning and knowledge distillation is proposed by Mao et al . ( 2020 ) . Among the above mentioned methods , Quantization requires hardware accelerator to reduce the inference time which is not applicable to general scenario . Pruning methods could only reduce the model size , but the inference speed might not be reduced due to the limitation of sparse operations . Only algorithmic method such as distillation could serve as a generic inference time accelerating method . We want to emphasize that our method is orthogonal to these distillation methods . In fact , the proposed method is a generic acceleration method applicable to all components in most NLP models . In section 4 , we show that DRONE can be combined with the distilled models to further improve the performance . 3 METHODS . We now introduce a generic algorithm for improving efficiency of matrix multiplication . The computation of Feed-Forward ( FF ) Layer in the attention models can be described as : h = Wx+ b , ( 1 ) where W ∈ Rd2×d1 and b ∈ Rd2 are model parameters , x ∈ Rd1 is the latent representation of a token , and h ∈ Rd2 is the output . Assume the sequence length is L , all the token representations x1 , . . . , xL ∈ Rd1 will pass through this same operation , so in practice the whole FF layer can be computed by a matrix-matrix product W [ x1 , . . . xn ] + b , and the computation of bias term b would be broadcasted to all L input tokens . In practice we will normally have L max ( d1 , d2 ) ( e.g. , L = 128 , d2 = 3 , 072 ) . A standard way to accelerate the computation is to perform low-rank approximation over W. A lowrank approximation can be acquired by using singular value decomposition ( SVD ) , which achieves the best rank-k approximation in terms of Frobenius norm and we could write W as : W = USV T ≈ UW , kVW , kT , with unitary matrices U ∈ Rd2×d2 , V ∈ Rd1×d1 and a diagonal matrix S ∈ Rd2×d1 . UW , k ∈ Rd2×k and VW , k ∈ Rd1×k are the rank-k approximation matrices by taking UW , k = US 1 2 k , VW , k = S 1 2 k V , where S 1 2 k is the square-root of the first k entries of the diagonal matrix S. Given such approximation , we could simplify the computation in equation 1 by writing it as : h = Wx+ b ≈ UW , kVW , kTx+ b . After the rank-k low-rank approximation , the computational complexity reduces from O ( d2d1 ) to O ( ( d1 + d2 ) k ) . When k is small enough , low-rank approximation could not only accelerate the computation ( Shim et al. , 2017 ) but also compress the model size ( Sainath et al. , 2013 ) . However , as we showed in Figure 2 , matrices in FF layer of BERT models do not show obvious low-rank structures . Ideally , we want a small percentage of the ranks which containing all large singular values such that sum of singular values connected to the selected ranks divided by sum of all singular values is large . But we could observe that choosing rank k to be larger than 50 % of the ranks ( e.g. , about 0.5 times min ( d1 , d2 ) ) could only accumulate 60 percent of the total singular values . This will lead to a large approximation error . In the meantime , the complexity is still about O ( d2d1 ) and there is no enhancement of speed . Despite the matrices in the model are not low-rank , here we provide an illustrative example to show that a low-rank computation could still exist when data distribution lies in a lower intrinsic dimension . Suppose we have a W defined as below and the input x lies in a subspace : W = 7 0 2 3 1 9 6 7 5 0 6 1 8 0 3 4 3 2 1 4 1 2 2 1 2 , x ∈ span ( 2 2 5 5 4 , 1 1 2 2 6 ) , In this case , W is a full-rank matrix so there won ’ t be a lossless low-rank approximation on W . On the other hand , input data x lies in a 2-dimensional subspace such that we could construct the following computation : 7 0 2 3 1 9 6 7 5 0 6 1 8 0 3 4 3 2 1 4 1 2 2 1 2 ︸ ︷︷ ︸ W 2 1 2 1 5 2 5 2 4 6 [ a b ] ︸ ︷︷ ︸ x = 43 23 90 39 66 41 45 37 29 21 ︸ ︷︷ ︸ U [ −1 −1 0.5 0.5 0 −0.5 0 0 0 0.25 ] ︸ ︷︷ ︸ V T 2 1 2 1 5 2 5 2 4 6 [ a b ] ︸ ︷︷ ︸ x , which gives a rank-2 matrix UV T where W 6= UV T but Wx = UV Tx for any x in the low dimensional space . This shows that even if we can ’ t approximate the W matrix , it is still possible to construct a good low-rank decomposition , and the key will be to exploit the space of input vectors .

Summary: This paper studies a technique to increase the inference speed and decrease model sizes of pretrained NLP models such as BERT. Since most operations in BERT consist of matrix multiplications, the authors conduct empirical experiments to show that while matrices themselves are not low-rank, the learned representation of each layer lies in a low-dimensional space. These empirical insights lead the authors to propose an approach based on data-aware low-rank compression of pretrained weight matrices and can be applied to fully-connected and self-attention layers. Their approach is able to improve both model size and inference speed