Language model compression with weighted low-rank factorization

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. However, 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.

Background

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 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:

The estimated information $\hat{I}_{w}$ accumulates the squared gradients over the training data $d_{i}\in\mathcal{D}$, where ${\mathcal{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 and Hua et al. 2021 use it to avoid the model catastrophic forgetting in a continual learning scenario. Liu et al. 2021 and 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.

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 $\bar{W}_{i}$ with the $i$-th singular value be truncated:

where $u_{i}$ and $v_{i}$ 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.

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 as 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 $W_{ij}$ in $W$, we weigh the individual reconstruction error by multiplying with the estimated Fisher information $\hat{I}_{W_{ij}}$:

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., $\hat{I}_{W_{i}}=\sum\limits_{j}\hat{I}_{W_{ij}}$.

Define the diagonal matrix $\hat{I}=diag(\sqrt{\hat{I}_{W_{1}}},...,\sqrt{\hat{I}_{W_{N}}})$, then the optimization problem of Equation (5) can be written as:

Equation 6 can be solved by the standard SVD on $\hat{I}W$. We use the notation $svd(\hat{I}W)=(U^{*},S^{*},V^{*})$, then the solution of Equation (6) will be $A={\hat{I}}^{-1}U^{*}S^{*}$, and $B=V^{*T}$. In other words, the solution is the result of removing the information $\hat{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={\hat{I}}^{-1}U_{r}^{*}S_{r}^{*}$, and $B=V_{r}^{*T}$, 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 $\hat{I}$ in Equation (6).

Experiments

This section describes how we obtain a compressed model under the popular pre-training schemes of language models. Figure 4 illustrates three paths that we examined for creating compressed language models. All the paths start from retraining a large transformer-based model pre-trained with a large language corpus in a self-supervised way, called the generic pre-training ($L\rightarrow L^{g}$).

The path-1 ($S\rightarrow S^{g}\rightarrow S^{t}$) is a popular scheme that creates a small model first, then performs the generic distillation for the small model to learn the knowledge of the large model. The resulting small generic model $S^{g}$ will be fine-tuned with the target task dataset to obtain the task-specific model $S^{t}$. The representative works of path-1 include DistilBERT (Sanh et al. 2019), TinyBERT (Jiao et al. 2019), MobileBERT (Sun et al. 2020), and MiniLM v1/v2 (Wang et al. 2020b; Wang et al. 2020a). Some previous methods may include task-specific distillation ($L^{t}\rightarrow S^{t}$) and data augmentation (Jiao et al. 2019), but we exclude those from the scheme (and all the experiments in this paper) to make a fair and clean comparison across methods. The task-specific distillation and data augmentation are orthogonal to all the methods and can be jointly applied with low-rank factorization to make further improvements.

The path-2 ($L^{g}\rightarrow L^{t}\rightarrow L^{tf}$) avoids the costly generic pre-training, directly compresses the task-specific model with factorization and task-specific fine-tuning (optional). Our analysis for the mismatched objectives phenomenon is based on this path. We also compare the models from path-1 and path-2, showing that path-2 can generate a model with a comparable performance under the same compression rate. Although path-2 requires much less training than path-1 (no generic pre-training for the compressed model).

The path-3 ($S^{t}\rightarrow S^{tf}$) is a challenging setting that aims to compress an already compact model. This setting examines whether FWSVD can further improve the compression rate on models obtained by path-1. Our experiments show the answer is yes.

With the three compression paths, we make four examinations as follows. Section 5.3: the comparison of path-1 versus path-2; Section 5.4: the compression of an already compact model (path-3); Section 5.5: the detailed comparison between FWSVD and vanilla SVD; Section 5.5.1: the empirical evidence for the schematic effects illustrated in Figures 2 and 3.

2 Experiment setup

We evaluate the methods of all three paths in Figure 4 on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al. 2019) and a token classification task. We include 2 single sentence tasks: CoLA (Warstadt et al. 2018) measured in Matthew’s correlation, SST2 (Socher et al. 2013) measured in classification accuracy; 3 sentence similarity tasks: MRPC (Dolan et al. 2005) measured in F-1 score, STS-B (Cer et al. 2017) measured in Pearson-Spearman correlation, QQP (Chen et al. 2018b) measured in F-1 score; and 3 natural language inference tasks: MNLI (Williams et al. 2018) measured in classification accuracy with the average of the matched and mismatched subsets, QNLI (Rajpurkar et al. 2016) measured in accuracy. The token classification task we used is the named entity recognition (NER) on the CoNLL-2003 dataset (Sang & De Meulder 2003). In summary, our evaluation includes 8 different natural language tasks.

2.2 Implementation details and the baseline models

First of all, we use the same training configuration for all the experiments in this paper and avoid any hyperparameter screening to ensure a fair comparison.

For the SOTA models on path-1 (MiniLMv2 and DistilBERT), we use the pre-trained generic compact models ($S^{g}$) provided by the original authors as the starting point, then directly fine-tune them with 3 epochs on the target task training data. The fine-tuning is optimized by Adam with learning rate of $2\times 10^{-5}$ and batch size of 32 on one GPU.

For the methods on path-2 (FWSVD and SVD), we start from the pre-trained generic large model ($L^{g}$), which is the standard 12-layer BERT model (Devlin et al. 2018). Then we fine-tune it with the training setting exactly the same as we used for the path-1 models to get the large task-specific models ($L^{t}$). The last step is applying the low-rank factorization (SVD or FWSVD) followed by another fine-tuning with the same training setting described above. The performance with and without fine-tuning will be both reported. We also note that we compress only the linear layers in the transformer blocks by reserving only 33% of the ranks in this work. The setup intentionally makes a fair comparison to the path-1 methods. In other words, we do not compress the non-transformer modules such as the token embedding. Previous works (Chen et al. 2018a) have shown significant success in using low-rank factorization to compress the embedding layer, which occupies 23.4M (21.3%) parameters in the standard BERT model. Therefore, the results we reported for the path-2 methods still have room for improvement by applying our method to non-transformer modules. Lastly, we add BERT-PKD (Sun et al. 2019) based on its reproduced results (Chen et al. 2018a) for comparison. BERT-PKD uses knowledge distillation instead of factorization in the path-2 process.

For the path-3 experiments, we use the pre-trained task-specific compact models ($S^{t}$) as the starting point. These pre-trained models have a much smaller size (TinyBERT-STSB, MiniLM-CoNLL, MobileBERT-MNLI) or a better performance (DistilBERT-MRPC) than the models we used in the path-1 and path-2, indicating they may contain denser knowledge in their compact models. Therefore, compressing these models introduces a significant challenge. In order to have better coverage for all tasks, we additionally use ALBERTlarge and ALBERTbase (Lan et al. 2019) as the already compact models to generate all 8 task-specific models ($S^{g}\rightarrow S^{t}$). Then follow path-3 to compress the compact models. All the training involved here has the same setting as described in path-1.

Lastly, our implementation and experiments are built on top of the popular HuggingFace Transformers library (Wolf et al. 2020). All other unspecified training settings use the default configuration of the library. Since no hyperparameter tuning is involved in our experiments, we directly report the results on the dev set of all the datasets, making the numbers convenient to compare and verify.

3 Path-1 versus Path-2

Table 1 reports the results of the GLUE benchmark and a NER task. Our FWSVD with fine-tuning achieves an average score of $83.6$, beating all other path-1 and path-2 methods. This is a non-trivial accomplishment since FWSVD with fine-tuning does not need the expensive generic pre-training. Furthermore, FWSVD has consistent performance retention for all the tasks; it contrasts the path-1 methods, which may have a more considerable variance. For example, DistilBERT is good at CoLA but poor at STS-B; oppositely, MiniLMv2 is a strong performer at STS-B but is weak with CoLA. In contrast, FWSVD+fine-tuning does not show an obvious shortcoming.

4 Compressing an already compact model

The setting of path-3 targets to further compress the lightweight models. This is challenging as the compact models are already 1.6x $\sim$ 7.8x smaller than the original BERT. The results in Table 2 demonstrate the effectiveness of FWSVD on further reducing the number of parameters. The original SVD is almost useless without fine-tuning, while our FWSVD can still retain a significant part of the performance. For example, SVD ends with a zero accuracy when compressing DistillBERT, while our FWSVD keeps a score of $67.9$ under the same setting. When combined with fine-tuning, FWSVD can cut off 30% redundancy for DistillBERT. Even for the highly compact model TinyBERT (only $14.4$M parameters), FWSVD+fine-tuning still successfully reduces 18% of the parameters without any performance loss. More interestingly, the TinyBERT, MiniLM, and DistillBERT-MRPC compressed by FWSVD+fine-tuning exceed the original performance slightly. The result suggests FWSVD+fine-tuning might introduce a small regularization effect to improve the model’s generalizability.

Lastly, Table 3 examines the compatibility between SVD/FWSVD and the parameter-sharing strategy of the ALBERT model. The average score of ALBERT-large is 84.7%. The performance of FWSVD (84.3%) is far better than that of SVD (77.1%) when both reducing 14% parameters, suggesting FWSVD is more robust than SVD in combining the parameter-sharing strategy.

5 FWSVD versus SVD

In Table 1, FWSVD consistently produces much better results than SVD on all tasks. On average, FWSVD without fine-tuning obtains an absolute improvement of 17.5% over SVD. To highlight, FWSVD without fine-tuning can maintain a significant portion of performance in challenging tasks such as CoNLL and STS-B, where SVD completely fails. With fine-tuning, FWSVD provides better initialization for fine-tuning and consistently achieves a better or comparable performance.

Figure 5 plots the performance trend with respect to the change of targeted rank ratio, where the full-rank reconstruction corresponds to the results at rank ratio $1.0$. These results demonstrate the apparent advantage of FWSVD over standard SVD. First, at each rank ratio, FWSVD shows significant improvements over SVD. Second, the performance of FWSVD keeps growing with the increase of rank ratio, while SVD shows fluctuations in its trend. Specifically, two tasks (COLA and STS-B) in Figure 5 show that SVD has abrupt performance drops at some points. On the STS-B task, the performance of SVD at rank ratio $0.3$ is significantly lower than having a smaller rank ratio of $0.2$. In contrast, FWSVD shows a much stable trend of increasing performance along with the rank ratio.

This section applies the same analysis of Section 3, but adds FWSVD to see if it matches the task’s objective better. In Figure 6(a), the red bars are significantly lower than the blue bars, especially for the tail groups, which will be truncated first. We specifically highlight group-10 in 6(a), which has the smallest 10% singular values. The height of the blue bar is equivalent to the size of the overlapped (green and meshed orange) region in Figures 2. Similarly, its red bar (close to zero) is equivalent to the overlapped region in Figure 3. In other words, the illustrations of Figures 2 and 3 are strongly supported by the results here. Although FWSVD shows a smaller performance drop shown by Figure 6(a), it has a more significant reconstruction error than SVD in many cases (see Figure 6(b)), especially with the rank groups that will be truncated first (e.g., groups 5 to 10). In other words, FWSVD’s objective (Equation 6) aligns with the task objective better by sacrificing the reconstruction error.

Limitation and future work

FWSVD has two limitations. First, FWSVD relies on a given task objective and a target task training dataset to compute the importance matrix; thus, it is more proper to compress a task-specific model (e.g., $L^{t}$ or $S^{t}$) than the pre-trained generic model (e.g., $L^{g}$). In contrast, the vanilla SVD can apply to any case. In other words, FWSVD trades the method’s applicability for the target task performance. Second, FWSVD only uses a simplified importance matrix that gives the same importance for the parameters on the same row of matrix $W$. Although this strategy is simple and effective, it does not fully utilize the Fisher information. Therefore, a future improvement can be made by directly seeking an element-wise factorization solution for Equation (5).

Conclusion

In this work, we investigate why using standard low-rank factorization (SVD) to compress the model may quickly lose most of its performance, pointing out the issue of the mismatched optimization objectives between the low-rank approximation and the target task. We provide empirical evidence and observations for the issue, and propose a new strategy, FWSVD, to alleviate it. Our FWSVD uses the estimated Fisher information to weigh the importance of parameters for the factorization and achieve significant success in compressing an already compact model. Furthermore, FWSVD reuses the existing SVD solver and can still implement its factorized matrices with linear layers; therefore, it is simple to implement and deploy. We believe FWSVD could be one of the most easy-to-use methods with good performance for language model compression.

References

Appendix A Supplementary