DRONE: 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 and its variations have been widely used in various tasks [ 8 , 2 , 25 , 3 , 24 ] . The successes of these large NLP models rely on the usage of large corpus and big models . Indeed , researchers have reported better results with models that have more parameters [ 31 ] and number of layers [ 1 ] . The increasing model size of the pre-trained models inhibits public users from training a model from scratch , and it also brings forth efficiency challenges , including inference speed and model size when deploying models on mobile devices . To deal with efficiency issues , most existing work resorts to adjusting the model structure or distillation . For instance , [ 17 ] used locality-sensitive hashing to accelerate dot-product attention , [ 18 ] used repeating model parameters to reduce the size and [ 44 ] applied a pre-defined attention pattern to save computation . A large body of prior work focused on variants of distillation has also been explored 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 27 , 15 , 35 , 20 , 34 , 35 , 41 , 45 , 4 ] . These methods require a specific design of model architecture , or a long training stage and thus it is less straightforward to combine these methods with each other . In this paper , we explore a simpler acceleration method to speed up inference time which can be applied to most existing architectures . As shown in Figure 1 , matrix multiplication ( feed-forward layer ) is a fundamental operation which appears many times in the Transformer [ 36 ] , the backbone architecture of the BERT model . 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 , a classical and simple model compression approach , can be used to accelerate Transformers . Despite its successful application to CNNs [ 42 , 33 , 32 ] , at first glance , low-rank compression does not appear to work for BERT since the matrices in both feed-forward layers and attention layers are not low rank ( see Figure 2 ) . Therefore , even the optimal low-rank approximation ( e.g. , by SVD ) will lead to very large reconstruction error . This is probably why low-rank approximation has not been successfully 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 ( features fed into each matrix mulitplication layer ) usually indicate some information extracted from natural sentences , and they often lie in a subspace with a low intrinsic dimension [ 5 , 32 , 22 ] . 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 lowrank approximation problem which includes the data distribution term and provide a closed-form solution for the optimal rank-k decomposition of the weight matrices . By leveraging the data distribution idea , we propose DRONE ( data-aware low-rank compression ) . Our decomposition significantly outperforms the SVD under the same rank constraint , and can successfully accelerate the BERT model without sacrificing too much test performance . In addition to compressing standard models , DRONE can also be used on distilled BERT models to further improve the compression rate . For example , DRONE alone achieves 1.92x speedup on the MRPC task with only 1.5 % loss in accuracy , and when combined with distillation , DRONE achieves over 12.3x speedup on various natural language inference tasks . 2 Related Work . Fast inference is important for deploying NLP models in various applications . Generally speaking , inference efficiency can be enhanced by hardware [ 30 ] or lower-level instruction optimization [ 23 ] . On the other hand , the main focus of the current research is on using algorithmic methods to reduce computational complexity . These methods can be mainly categorized into two aspects : attention complexity reduction and model size reduction . Attention Complexity Reduction Attention mechanism is the building block of transformer models and has attracted the most attention of researchers recently in the NLP field [ 36 ] . Pre-training on large corpus of BERT , a transformerbased model , has contributed to state-of-the-art performance on various tasks after fine-tuning [ 8 ] . Attention on sequences of length L is O ( L2 ) in both computational and memory complexity , which yields long inference time when the sequence is long . Thus , researchers have focused on reducing the complexity of the attention module . [ 17 ] used locality-sensitive hashing to reduce the complexity to O ( L logL ) . [ 44 , 7 ] pre-defined an attention map to have a constant computational time . [ 11 ] progressively eliminated the redundant context vectors within the attended sequence to improve efficiency of attention in the last few layers of the model . [ 38 ] proposed to train the low-rank attention by choosing a rank r L. This is similar to our work in the sense of leveraging low-rank structures . But our method does not require training the model from scratch and can be applied to different modules other than attention . In fact , most of the above methods require special modules and thus need to train 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 mainly focus on the long sequence scenario . As shown in Figure 1 , we have found out that attention module is actually not the main inference bottleneck of inference time in common usage . In most , if not all , models of common usages , two layers of large feed-forward layer are appended after the attention module which incurs much more computational time . Attention complexity reduction only works when a long sequence is used but in current practice this is unusual . Thus , in many tasks accelerating the attention module itself does not contribute to a significant reduction of overall inference time . Model Size Reduction Inference speed is also related to model compression . In principle , smaller models lead to reduction in the number of operations and thus faster inference time . [ 28 ] explored pruning methods on BERT models to eliminate redundant links , and there is a line of research on pruning methods [ 12 , 13 , 6 , 10 ] . Quantization methods [ 43 , 14 , 19 , 9 ] convert the 32 bits float models into fewer-bits fixed-point representation and make model prediction faster with fixed point accelerator . [ 18 ] reduce the model size by sharing encoder parameters . A large body of prior work focused on variants of knowledge distillation [ 27 , 15 , 35 , 20 , 34 , 35 , 41 , 45 , 4 ] . These methods use different strategies to distill information from a teacher network and reduce the number of layers [ 27 ] or hidden dimension size [ 15 ] . Further , a hybrid compression method by combining matrix factorization , pruning and knowledge distillation is proposed by [ 21 ] . Notice that [ 21 ] performed SVD for some components and in this paper we propose an improvement over SVD by leveraging input distribution to each layer . Idea of using input distribution to compress model has also been explored in PCN method [ 37 ] , which is perhaps the closest to our work . However , DRONE differs from PCN in following three aspects . First , PCN only considers input distribution but not weight matrix and thus it ’ s a special case of DRONE ( i.e. , W be an identity matrix in equation ( 3 ) ) . Second , PCN merely does dimension reduction whereas our formulation achieves dimension reduction and low-rank approximation simultaneously . Last , PCN does not guarantee the obtained transformation is the optimal ; whereas , DRONE formulates an approximation optimization problem and we provide the optimal solution . Other forms of low-rank learning strategies including initialization and structure pruning were also explored in the literature [ 39 , 16 ] , and we will compare to these baseline methods . Among the above-mentioned methods , quantization requires hardware accelerator to maximally reduce the inference time . Pruning methods can only reduce the model size , but the inference time might not be reduced due to the limitation of sparse operations . Only algorithmic methods such as distillation serve as a more generic inference time accelerating method . We want to emphasize that our method is orthogonal to these distillation methods . In fact , the proposed method is an acceleration method that is 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 Proposed Method . We now introduce an 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 ) o = σ ( h ) , ( 2 ) where W ∈ Rd2×d1 and b ∈ Rd2 are model parameters , x ∈ Rd1 is the latent representation of a token , and h ∈ Rd2 is the intermediate representation before the activation function , σ ( · ) is the activation function , and o ∈ Rd2 is the output . Assuming 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 , . . . xL ] + b , and the computation of the bias term b would be broadcast to all L input tokens . In practice we will normally have L max ( d1 , d2 ) ( e.g. , L = 128 , d2 = 3072 ) . Notice that applying σ ( · ) on h element-wisely costs O ( Ld2 ) , which is much smaller than the cost of computing Wx ( O ( Ld2d1 ) ) . Therefore , in this paper we focus on reducing the cost of computing Wx to accelerate the computation . A standard way to accelerate this computation is to perform low-rank approximation on W . A low-rank approximation can be obatined by using singular value decomposition ( SVD ) , which achieves the best rank-k approximation in terms of Frobenius norm and we can write W as : W = USV T ≈ UW , kVW , kT , with orthogonal 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 = V S 1 2 k , where S 1 2 k is the square-root of the first k entries of the diagonal matrix S. Given such an approximation , we can simplify the computation in ( 1 ) by h = Wx+ b ≈ UW , kVW , kTx+ b . After conducting rank-k approximation , the computational complexity reduces from O ( d2d1 ) to O ( ( d1+d2 ) k ) . When k is small enough , low-rank approximation not only accelerates the computation [ 32 ] but also compresses the model size [ 26 ] . However , as shown in Figure 2 , matrices in FF layer of BERT do not show obvious low-rank structures . We observe that choosing top 50 % rank ( e.g. , k = 0.5 min ( d1 , d2 ) ) can only achieve around 60 % of the accumulation ratio of singular values , which implies large matrix approximation error . In the meantime , the complexity is still about O ( d2d1 ) and there is no enhancement of speed . Even though the matrices in the model are not low-rank , we now 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 matrix 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 is no lossless low-rank approximation of W . On the other hand , the input data x lies in a 2-dimensional subspace so that we could construct the following low-rank approximation : 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 W can not be approximated , it is still possible to construct a good low-rank decomposition , and the key is to exploit the space of input vectors .

Summary: This paper proposes DRONE, a data-aware low-rank model compression method designed specifically for large language models, e.g., BERT. The method is motivated by an observation that the weight matrices in a pre-trained BERT model are generally not low rank, but the data/intermediate results usually lie in a low-rank space. DRONE leverages the above observation and arrives at a close form solution to approximate the layer weights using low-rank factorization. Theoretical analysis is provided and experimental results are shown to demonstrate that DRONE successfully accelerate the inference speed of BERT model on a various range of down stream tasks.