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: This paper studies the compression of the weights of a neural network via a modification of the standard paradigm of using low rank approximations via SVD on the weight matrices. Instead, they propose a weighted SVD objective where the weights are computed based on the Fisher information of the (trained) parameter with respect to the dataset. They simplify the method so it is more computationally tractable by having the same row of each weight matrix share the same weight. This translates to solving SVD on a modified reconstruction objective (multiply by a diagonal). Here the Fisher information for parameter w is approximated by the average squared l2 norm of the gradient with respect to w over the empirical data distribution.  They then conduct several experiments with this approach in the problem domain of compressing language models. They evaluate their method on compressing task specific models as well as the problem of compressing an already compressed model, and demonstrate their approach is generally better.