Language model compression with weighted low-rank factorization

1 INTRODUCTION . Language models built with transformers ( Devlin et al. , 2018 ) have attained extensive success in natural language tasks such as language modeling ( Radford et al. , 2018 ) , text classification ( Wang et al. , 2018 ) , question answering ( Rajpurkar et al. , 2016 ) , and summarization ( Liu , 2019 ) . The success is achieved by fine-tuning a big transformer model pre-trained with a large corpus . The target task for fine-tuning may only focus on a restricted scenario such as sentiment analysis ( Socher et al. , 2013 ) and multiple-choice question inference ( Zellers et al. , 2018 ) . Having a big transformer model is often overkill for the target task and prohibits the model deployment to resource-constrained hardware . Therefore , language model compression raises immense interest . The popular strategy creates a compact model from scratch ( Jiao et al. , 2019 ) or a subset of the big model ’ s layers ( Sun et al. , 2019 ; Sanh et al. , 2019 ) , then pre-trains with a large corpus and distills knowledge from the big model . This process is called generic pre-training ( Wang et al. , 2020b ; Sun et al. , 2019 ; Sanh et al. , 2019 ) and is necessary for a compact model to achieve good performance on the target tasks . The generic pre-training could still cost considerable computational resources . For example , it takes 384 NVIDIA V100 GPU hours to get the pre-trained TinyBERT ( Jiao et al. , 2019 ) on the Wiki corpus dataset . So it may not be affordable for everyone who wants to create a compact model . In contrast , another line of strategy , specifically low-rank factorization ( Golub & Reinsch , 1971 ; Noach & Goldberg , 2020 ) , can potentially reduce a big model ’ s parameters without the generic pre-training . Since the factorization aims to approximate the learned model parameters , the method has the nature of directly inheriting the knowledge of the big trained model . However , approximating the learned weights with standard factorization often loses most of the task performance . This work investigates this issue with the most popular strategy , which uses singular value decomposition ( SVD ) to compress the learned model weights . With SVD , the learned matrix is factorized into three matrices ( U , S , V ) . The portion associated with small singular values will be truncated to produce a smaller version of factorized matrices . The multiplication of these smaller matrices will approximate the original one with fewer total parameters to achieve the model compression . In other words , SVD minimizes the reconstruction error with fewer parameters as its objective . However , this objective does not necessarily correlate to the ultimate goal of keeping task performance . Specifically , the SVD algorithm is biased to reconstruct the parameters associated with large singular values . As a result , the parameters mainly reconstructed by the ranks with small singular values will become the victim in the compression process . Are these victimized parameters less critical to achieving a good task performance ? We argue that this is not true , and the optimization objective of SVD is not properly aligned with the target task objective . This paper is the first work to provide an empirical analysis of this issue , proposing a novel weighted SVD to mitigate it . Our weighted SVD addresses the above issue by assigning importance scores to the parameters . This score has to correlate to how much the task performance is affected by the parameter change . The Fisher information nicely fits into this purpose ( Pascanu & Bengio , 2014 ) . Besides , the calculation of Fisher information is usually simplified to accumulating a parameter ’ s squared gradient over the training dataset based on its task objective ( e.g.cross-entropy , regression error , etc . ) , conveniently providing the importance of each parameter in a model . Then we modify the optimization objective of factorization ( i.e. , reconstruction error ) by multiplying it with Fisher information , providing a new objective that jointly considers matrix reconstruction error and the target task objective . In summary , this work makes the following contributions : ( 1 ) we analyze the issue of mismatched objectives between factorization and the target task for model compression ; ( 2 ) we propose a novel compression strategy with the SVD weighted by the Fisher information ; ( 3 ) we perform extensive analysis on varied language tasks , showing our Fisher-weighted SVD can compress an already compact model , and it can achieve comparable compression rate and performance with methods that require an expensive generic model pre-training . 2 BACKGROUND . 2.1 MODEL COMPRESSION WITH LOW-RANK APPROXIMATION . Given a matrix W ∈ RN×M , the low-rank approximation is achieved via singular value decomposition ( SVD ) : W ≈ USV T , ( 1 ) where U ∈ RN×r , V ∈ RM×r , and k is the rank of matrix W . S is a diagonal matrix of nonzero singular values diag ( σ1 , , ... , σr ) , where σ1 ≥ σ2 ≥ · · ·σr ≥ · · ·σk > 0 . The low-rank approximation with targeted rank r is obtained by setting zeros to σr+1 , ... , σk . Given input data X ∈ R1×N , a linear layer in neural networks is represented below with the weight matrix W ∈ RN×M and bias b ∈ R1×M : Z = XW + b ≈ ( XUS ) V T + b . ( 2 ) Factorizing W with Equation ( 1 ) leads to Equation ( 2 ) , which can be implemented with two smaller linear layers : 1 ) The first layer has Nr parameters without bias . Its weight matrix is US . 2 ) The second layer has Mr parameters plus bias . Its weight matrix and bias are V and b , correspondingly . The total number of parameters for approximating W is Nr + Mr . In the case of full rank matrix and M = N , the model size is reduced when r < 0.5N . For example , if we set r to reserve the largest 30 % singular values , the method will reduce about 40 % of the parameters from W . In general , the reduced size will be NM − ( Nr +Mr ) . Low rank approximation in neural networks has been extensively studied ( Jaderberg et al. , 2014 ; Zhang et al. , 2015 ; Denton et al. , 2014 ) . In more recent works , SVD is often applied to compress the word embedding layer ( Chen et al. , 2018a ; Acharya et al. , 2019 ) . Noach & Goldberg ( 2020 ) applies SVD to the transformer layers , but it does not investigate why SVD gives a very poor result without fine-tuning . Our work explores this issue and provides a weighted version to address it . 2.2 FISHER INFORMATION . The Fisher information measures the amount of information that an observable dataset D carries about a model parameter w. The computation of its exact form is generally intractable since it requires marginalizing over the space of D , which includes data and its labels . Therefore , most of the previous works estimate its empirical Fisher information : Iw = E [ ( ∂ ∂w log p ( D|w ) ) 2 ] ≈ 1 |D| |D|∑ i=1 ( ∂ ∂w L ( di ; w ) ) 2 = Îw . ( 3 ) The estimated information Îw accumulates the squared gradients over the training data di ∈ D , where L is the target task objective ( e.g. , cross-entropy for a classification task , or mean squared error for a regression task ) . This approximation provides a straight intuition : the parameters that change the task objective with a large absolute gradient are important to the target task ; therefore , those parameters should be reconstructed better than others in the compression process . The above estimation of Fisher information computes only the first-order derivatives and has been shown to measure the importance of parameters effectively . Kirkpatrick et al . ( 2017 ) ; Hua et al . ( 2021 ) use it to avoid the model catastrophic forgetting in a continual learning scenario . Liu et al . ( 2021 ) ; Molchanov et al . ( 2019 ) use it or a similar variant to help the structured model pruning . However , no previous work has explored its potential in assisting SVD for model compression . 3 MODEL COMPRESSION WITH SVD MAY LOSE PERFORMANCE QUICKLY . The singular values in S implicitly give an importance score for a group of parameters . Since the small singular values will be truncated first , those parameters affected by the truncation are expected to be not important for the task performance . We verify the above assumption with a brute force attack : truncate one singular value at a time , then reconstruct the matrix , put it into a model , evaluate and get its performance . Ideally , we hope to see less performance drop when we truncate the smaller singular values . This process can be written as having the reconstructed model weights W̄i with the i-th singular value be truncated : W̄i = u1σ1v T 1 + ... + ui−1σi−1v T i−1 + ui+1σi+1v T i+1 + ... + ukσkv T k , ( 4 ) where ui and vi are the i-th column in U and V , correspondingly . Applying this brute force attack to test a deep neural network is not straightforward since a model can have hundreds of linear layers . Therefore , we truncate a group of singular values together instead of only one . Specifically , we split the singular values of a layer into 10 groups sorted by their values . The 1st group has the top 10 % singular values , while the 10th group contains the smallest 10 % values . When we truncate a specific group , e.g. , 5th group , the 5th group of all the layers in a model are truncated together . In other words , we observe the summed impact in a rank group . This results in a smoothed trend for the observation . Figure 1 plots the result of truncating the 10 groups separately in a standard 12-layer BERT model ( Devlin et al. , 2018 ) trained for STS-B task ( Cer et al. , 2017 ) . The red dashed line shows an ideal trend which has a smaller performance drop with the tail groups . The blue bars show the actual performance drop . The 10th group surprisingly caused a performance drop as large as the 2nd group . This means the parameters associated with the 10th group are as important as the 2nd group . However , the magnitude of singular value does not reflect this importance , causing a model to lose its performance quickly even when truncating only a small portion . 4 FISHER-WEIGHTED LOW-RANK APPROXIMATION . The issue in Section 3 has an intuitive cause : the optimization objective of SVD does not consider each parameter ’ s impact on the task performance . This issue is illustrated in Figure 2 , and we address it by introducing the Fisher information into SVD ’ s optimization objective , described below . In the generic low-rank approximation , its objective minimizes ||W − AB||2 . SVD can solve this problem efficiently by having A = US and B = V T . Since we can obtain the importance of each element Wij in W , we weigh the individual reconstruction error by multiplying with the estimated Fisher information ÎWij : min A , B ∑ i , j ÎWij ( Wij − ( AB ) ij ) 2 . ( 5 ) In general , weighted SVD does not have a closed-form solution ( Srebro & Jaakkola , 2003 ) when each element has its weight . To make our method easy to deploy and analyze , we propose a simplification by making the same row of the W matrix to share the same importance . The importance for the row i is defined to be the summation of the row , i.e. , ÎWi = ∑ j ÎWij . Define the diagonal matrix Î = diag ( √ ÎW1 , ... , √ ÎWN ) , then the optimization problem of Equation ( 5 ) can be written as : min A , B ||ÎW − ÎAB||2 . ( 6 ) Equation 6 can be solved by the standard SVD on ÎW . We use the notation svd ( ÎW ) = ( U∗ , S∗ , V ∗ ) , then the solution of Equation ( 6 ) will be A = Î−1U∗S∗ , and B = V ∗T . In other words , the solution is the result of removing the information Î from the factorized matrices . Figure 3 illustrates this process and its schematic effect of reducing the overlap between important parameters and poorly reconstructed parameters . We will measure this overlap with the performance drop analysis of Section 3 . Lastly , to compress W , we will have A = Î−1U∗r S ∗ r , and B = V ∗T r , where r denotes the truncated U∗ , S∗ , and V ∗ with reserving only r ranks . We call the above method FWSVD in this paper . One thing to highlight is that since we share the same optimization process with the standard SVD , any advantage we observed will be the result of a direct contribution from the Î in Equation ( 6 ) .

Summary: SVD of matrices in a language model is used to reduce model size. Instead of standard SVD that minimizes the MSE of reconstructing the matrix elements, a modified version that weights elements is used. Specifically, importance of each parameter (matrix element) is estimated based on the magnitude of gradients w.r.t. to that element. To have closed-form solution for weighted SVD, average row weights are used instead of element weights.