LoSparse: Structured Compression of Large Language Models based on Low-Rank and Sparse Approximation ¹¹1Published as a conference paper in ICML 2023.

A typical transformer architecture comprises several sequential layers, where each layer contains two sub-layers: a multi-head self-attention (MHA) and a fully connected feed forward network (FFN). Given the input $X\in\mathbb{R}^{n\times d}$, MHA computes the attention in parallel $h$ heads:

where $W_{q_{i}},W_{k_{i}},W_{v_{i}}\in\mathbb{R}^{d\times d_{h}}$ are query, key, and value projection matrices, $W_{o}\in\mathbb{R}^{d\times d}$ is an output projection matrix, and $d_{h}$ is typically set as $d/h$. FFN comprises two linear transformations and an activation: $\mathrm{FFN}(X)=\sigma(XW_{f_{1}}+b_{1})W_{f_{2}}+b_{2},$ where $W_{f_{1}}\in\mathbb{R}^{d\times d_{m}}$, $W_{f_{2}}\in\mathbb{R}^{d_{m}\times d}$, and $\sigma(\cdot)$ is the activate function. A residual connection is used and followed by a layer normalization.

Generally, we denote all the matrix multiplication in a transformer model as

where $W\in\mathbb{R}^{d_{1}\times d_{2}}$ denotes any weight matrix in the model.

We further denote the parameter set consisting of all trainable weight matrix by $\mathcal{W}=\{W_{m}\}_{m=1}^{M}$. Unless specified otherwise, we use $W$ to represent any weight matrix and $W^{(0)}$ is its pre-trained value. We use $i,j$ to index the entry of matrices and denote $w_{ij}$ as $ij$-th entry of $W$, $W_{i*}$ as the $i$-th row of $W$, and $W_{*i}$ as the $i$-th column of $W$.

Pruning methods zero out redundant parameters according to their importance estimation. Parameters with high importance scores are retrained for fine-tuning while the others with low importance are zeroed out. Popular importance metrics include magnitude (Han et al., 2015), sensitivity (Sanh et al., 2020; Molchanov et al., 2019) and uncertainty (Zhang et al., 2022). Sensitivity of parameters is essentially designed to approximate the change of training loss $\mathcal{L}$ when a parameter is zeroed out. If the removal of a parameter causes a large variation on the loss, then the model is sensitive to this parameter and we should retain it. Specifically, for a weight $w_{ij}$, its sensitivity score is defined by the gradient-weight product:

Note that the calculation of $I^{(t)}$ is conditioned on the sampled mini-batch at the $t$-th iteration. It can induce high uncertainty due to stochastic sampling. To reduce the variability in (2), Zhang et al. (2022) propose to smooth $I$ by:

using exponential moving of average.

As mentioned in Section 1, there are two types of pruning methods: unstructured and structured pruning. Sensitivity in (2), however, targets on unstructured pruning. We extend it to structured pruning and introduce neuron importance scores. For a linear projection represented as a weight matrix $W\in\mathbb{R}^{d_{1}\times d_{2}}$, we define the importance score of its $i$-th neuron $W_{*i}$ as

We further define $\Gamma(W)=[\Gamma(W_{*1}),...,\Gamma(W_{*d_{2}})]^{\top}\in\mathbb{R}^{d_{2}}$.

Given a weight matrix $W\in\mathbb{R}^{d_{1}\times d_{2}}$, a structured-pruned sparse matrix $S\in\mathbb{R}^{d_{1}\times d_{2}}$ is commonly applied to approximate $W$ for compression (Han et al., 2015; Lagunas et al., 2021). The sparse matrix approximation, however, results in poor performance especially when the remaining ratio is low (See experiments in Section 4). Therefore, we introduce a low-rank matrix to improve the approximation. Specifically, the weight matrix can be represented as

where the product of $U\in\mathbb{R}^{d_{1}\times r}$ and $V\in\mathbb{R}^{r\times d_{2}}$ represents the low-rank matrix of rank $r$.

Why low-rank matrices? First, they can approximate the coherent parts of neurons effectively, even if the rank is small. As shown in Figure 3, we observe that the spectrum of weight matrices in language models drops rapidly at the beginning. This indicates neurons in a weight matrix share a common subspace, which can be viewed as the coherent parts of these neurons. In addition, the common subspace can be recovered by the singular vectors of top singular values. Therefore, the coherent parts of neurons can be well approximated by the low-rank matrix computed by singular value thresholding.

Second, the decoupling of low-rank and sparse matrices makes it easy to prune. The heavy-tailed spectrum in Figure 3 indicates each neuron $W_{*i}$ spans their individual subspaces, which can represent the incoherent parts of these neurons. Since these subspaces are not shared, the incoherent parts cannot be captured by the low-rank approximation. Fortunately, the low-rank matrix is able to decouple the coherent parts from the incoherent parts of neurons. This enables us to approximate the remaining incoherent parts by adding a new matrix $S$ and then prune it to remove the non-expressive incoherent parts. As an example, Figure 4 demonstrates that most of the incoherent parts have low importance scores after decoupling, motivating us to remove these redundant parameters.

We then present our proposed algorithm. Given a pre-trained weight matrix $W^{(0)}$, we first initialize the low-rank matrix of rank $r$ based on the singular value decomposition (SVD) of $W^{(0)}$. Specifically, we choose

where $u_{1},u_{2},...,u_{r}\in\mathbb{R}^{d_{1}}$ are left-singular vectors and $v_{1},v_{2},...,v_{r}\in\mathbb{R}^{d_{2}}$ are right-singular vectors, with respect to the top $r$ singular values $\sigma_{1}\geq\sigma_{2}\geq...\geq\sigma_{r}$ in the SVD of $W^{(0)}$. Then, we initialize $S^{(0)}$ by

Notably, we replace the forward pass involving $W$ (e.g. $Y=XW$) with (3.2) to improve computational efficiency:

We apply such a decomposition to every weight matrix of the model and denote $\mathcal{S}=\{S_{m}\}_{m=1}^{M}$ as the set of all sparse matrices. After the initialization, we conduct the iterative structured pruning for $S$. Specifically, at $t$-th iteration, we first take a stochastic gradient decent step to update $U^{(t)},V^{(t)}$, and $S^{(t)}$. In particular, for $S^{(t)}$,

Then we evaluate the neuron importance scores of $S^{(t)}$ based on (4). Given the importance scores, $\widetilde{S}^{(t)}$ is pruned following:

with the $i$-th column of $\mathcal{T}(\widetilde{S}^{(t)},\Gamma(S^{(t)}))$ defined as

Here we retain $\widetilde{S}^{(t)}_{*i}$ only if its importance score is in top $p_{t}$% among all neurons in $\mathcal{S}^{(t)}$. $p_{t}$ is the percentage of remaining neurons at $t$-th iteration. We gradually decay $p_{t}$ following a cubic schedule:

where $T$ is the total training steps. $t_{i}$ is the number of initial warm-up steps. $t_{f}$ is the number of final fine-tuning steps. Finally, we summarize our algorithm in Algorithm 1.

Models and Datasets. We evaluate the performance of LoSparse when pruning DeBERTaV3-base models on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2019). GLUE includes two single-sentence classification tasks: SST-2 (Socher et al., 2013) and CoLA (Warstadt et al., 2019), and three similarity and paraphrase tasks: MRPC (Dolan & Brockett, 2005), STS-B (Cer et al., 2017), and QQP. There are also four natural language inference tasks in GLUE: MNLI (Williams et al., 2018), QNLI (Rajpurkar et al., 2016a), RTE (Dagan et al., 2007, 2006; Giampiccolo et al., 2007; Bentivogli et al., 2009), and WNLI (Levesque et al., 2012). Following previous works, we exclude WNLI in the experiments.

Implementation Details. We select the learning rates from $\{1\times 10^{-5},3\times 10^{-5},5\times 10^{-5},8\times 10^{-5},9\times 10^{-5},1\times 10^{-4}\}$. We select the proportion of the parameters of all low-rank matrices over all pre-trained parameters from $\{1\%,2\%,3\%,5\%\}$. We discuss the influence of different proportion later in Section 4.4. More implementation details, such as the training epochs and batch sizes, are presented in the Appendix B.

Main Results. We compare our method with the baseline methods under different remaining ratios. The results are shown in Table 1. We see that LoSparse achieves better or on par performance compared with existing approaches on all the datasets of GLUE under all remaining ratios. For example, when the remaining ratio is 10%, LoSparse achieves 81.7% accuracy on MNLI-m dataset, which surpasses the best-performing baseline (ITP) by 2%. In addition to the superior performance, our method is more stable than the baselines (e.g. ITP and Movement). This is because each weight matrix in LoSparse at least maintains a low-rank matrix $S$ and always updates it along training horizon. It prevents the dramatic variation of weight matrices from nonzero to zero. By contrast, weight matrices are possibly pruned to zero by other iterative pruning methods. The expressive parts in these weight matrices can alternate between being pruned and updated and finally leads to divergence.

Table 2 summarizes the results of pruning BERT-base on MNLI, RTE, and QNLI. Similar to Table 1, our methods outperforms all baselines under all sparsity level for all three datasets. For example, when the remaining ratio is 20%, LoSparse achieves 65.2% accuracy on RTE dataset, while ITP only achieves 64.4% accuracy. We remark that LoSparse is even more effective under high sparsity level. For instance, given the 10% remaining ratio, LoSparse outperforms ITP by 0.6% on MNLI-m dataset (78.3 v.s. 77.7), 1.2% on RTE (63.0 v.s. 61.8), and 0.9% on QNLI (84.8 v.s. 83.9).

Models and Datasets. We evaluate the performance of our method on the question-answering task (SQuADv1.1, Rajpurkar et al. (2016a)). In the SQuADv1.1, question answering is treated as a sequence labeling problem, where we predict the probability of each token being the start and end of the answer span. We compress DeBERTaV3-base and BERT-base on the SQuADv1.1.

Implementation Details. We compress all the backbone weight matrices in DeBERTaV3-base model and BERT-base except layernorm and final classification head. We use learning rates from $\{1\times 10^{-5},3\times 10^{-5},5\times 10^{-5},8\times 10^{-5}\}$ and pick the learning rate that performs the best. We also select the proportion of parameters of low-rank approximation from $\{1\%,2\%,3\%,5\%\}$. We choose AdamW as the optimizer and set the batch size as 16. Please refer to the Appendix C for more details.

Main Results. We compare our method with the baseline methods under different sparsity levels. The experimental results are shown in Table 3. We can see that LoSparse consistently surpasses baseline methods under all remaining ratios in terms of the two evaluation metrics: exact match (EM) and F1. Similar to our result in GLUE tasks, our method is especially effective with low remaining ratios. For example, LoSparse outperforms ITP by 3.0% in terms of F1 if removing 95% of parameters. Even for high remaining levels, our method still achieves considerable performance gain. For example, LoSparse outperforms ITP by 1.9% in terms of F1 if removing 60% of parameters.

Table 3 also summarizes pruning BERT-base with different methods on SQuADv1.1. Our method achieves substantial improvements compared to all baseline methods. For example, our method outperforms the best baseline ITP by 2.6% on F1 given 10% remaining ratio. We remark that ITP is also effective under low sparsity levels. For example, ITP achieves the best result over LoSparse and movement pruning given the remaining ratio as 50%. Our method, however, still behaves on par with ITP with high remaining ratios: both ITP and LoSparse achieve 84.2 on F1 under 40% remaining ratios.

Models and Datasets. In natural language generation (NLG) tasks, we compress BART-large model (Lewis et al., 2020) to compare LoSparse with baseline methods. We evaluate the performance on the XSum (Narayan et al., 2018) and CNN/DailyMail(Hermann et al., 2015) datasets.

Implementation Details. We apply our method to all weight matrices of both encoder and decoder layers. We report ROUGE 1/2/L scores, which are the metrics for summarization tasks (Lin, 2004). Given a fixed total remaining ratio, we try different allocations between the low-rank matrices and the sparse matrices. The best allocation of sparse matrices is 10%. We choose the learning rate from $\{6\times 10^{-6},1\times 10^{-5},2\times 10^{-5},4\times 10^{-5},6\times 10^{-5},1\times 10^{-4}\}$. The training epochs and batch sizes are set to 12 and 32 respectively. The beam search length is 8. Please see Appendix D for the detailed configuration.

Main Results. We compare LoSparse with baseline methods under 30%, 40%, and 50% remaining ratios. We do not report results on lower remaining ratios because the baseline methods fail to surpass the Lead-3 baseline. Table 4 summarizes experiment results on the XSum and CNN/DailyMail test datasets. Note that LoSparse consistently surpasses the baseline methods under all remaining ratios in terms of ROUGE scores. For instance, LoSparse outperforms ITP on XSum dataset by 2.99 in terms of ROUGE-1 score. We also remark that LoSparse is more efficient under extremely low remaining ratios. For example, the gain on ROUGE-1 increases from 0.76 to 2.99 when the ratio drops from 50% to 30%. Note that LoSparse is particularly effective on more difficult summarization tasks. For example, XSum is more abstract and hence more difficult than CNN/DailyMail, and LoSparse yields 2.99 gain on XSum compared to 0.86 on CNN/DailyMail.

Effectiveness of Sparse Approximations. We experiment the model compression without sparse approximation to study its effectiveness. Specifically, we compare LoSparse with two variants: (i) we discard the sparse matrices and only fine-tune the low-rank matrices $UV$ (Low-rank I); (ii) we follow the initialization as (8) but gradually prune the initialized $S$ into zero (Low-rank II). Figure 5 summarizes the performance of these two variants on MNLI, SQuAD, and XSum. The results show that our method outperforms two low-rank variants, which verifies the effectiveness of the sparse approximation. Moreover, we find Low-rank II is better than Low-rank I. We discuss it in Section 5.

Sparsity Allocation. We study how low-rank and sparse approximations cooperate with each other. Specifically, given a fixed remaining ratio, we change the proportion of low-rank matrices and accordingly the ratio of sparse matrices. Figure 6 summarizes the result under different allocations. We see low-rank and sparse approximations exhibit the nearly equal contribution to the performance on NLU tasks as the performance stays stable when changing the allocation.

Knowledge distillation is a popular technique to improve the performance of small models (Romero et al., 2014; Hinton et al., 2015). In knowledge distillation, the small model (student) is trained to mimic the output of a larger fine-tuned model (teacher) such that the performance of the small model can be improved.

We remark that compression methods are complementary to knowledge distillation. We show it by integrating knowledge distillation into LoSparse and other pruning methods. Specifically, we choose a DeBERTaV3-base model that is fine-tuned on specific tasks as the teacher model and a compressed DeBERTaV3-base model as the student model. Then we conduct layer-wise distillation for them. Please see Appendix E for more training details. Table 5 shows the results. We find that distillation can further improve the performance of LoSparse and other compression methods. It shows that compression and knowledge distillation are complementary to each other. Besides, LoSparse still achieves better performance than ITP when integrated with distillation. It demonstrates the effectiveness of LoSparse in the setting of distillation.

We conduct further investigation on the performance of combining LoSparse and knowledge distillation on BERT-base models. Similarily, we choose a BERT-base model that is fine-tuned on specific tasks as the teacher model and a compressed BERT-base model as the student model. We find out combining LoSparse and knowledge distillation can achieve a comparable or even better performance than popular compression method, such as PKD (Sun et al., 2019), MixKD (Liang et al., 2020), CoDIR (Sun et al., 2020), BERT-of-Theseus (Xu et al., 2020), Low-Rank BERT Feature Distillation + KD (Noach & Goldberg, 2020), Path2: BERT+FWSVD(Hsu et al., 2022b). Please see Appendix E for more training details.

Table 6 shows the comparison of the aforementioned methods. From the table , we can see that LoSparse combined with knowledge distillation can achieve an on-par or better result than most distillation methods, such as PKD, under different remaining ratio. Our method also excels the combination of low-rank approximation methods and knowledge distillation such as Low-Rank BERT Feature Distillation + KD.

LoSparse is a generic compression method. It can be embedded into other popular methods, such as CoFi (Xia et al., 2022). CoFi is a coarse to fine-grained compression approach. It uses 3-level masks to determine which layer, heads, and neurons should be pruned. In the first level, it adds masks to MHA sub-layers and FFN sub-layers as a coarse compression. In the second level, it adds masks to the attention heads inside the MHA. In the final level, it adds masks to every neuron as the third level compression. In addition, it utilizes knowledge distillation to enhance the performance.

To embed our method into CoFi, we replace the third level masks as our method and keep the first and second level masks. Specifically, we first decompose the pre-trained weight matrices into low-rank and sparse matrices. Then, we follow the same training approach as CoFi. As for distillation, since CoFi has not released the teacher models, we download all the teacher models from Text Attack ^******https://huggingface.co/textattack(Morris et al., 2020) except teachers for the MNLI task. To obtain the MNLI teacher, we fine-tune BERT-base using following hyperparameters: learning rate is $3\times 10^{-5}$, batch size is 32, number of training epochs is 3. See Appendix F for more experiment details.

Experiment results are listed in Table 7. We see that LoSparse can improve the performance of CoFi on all datasets. For example, our method improves the accuracy by around 1% on MRPC, RTE, and SST-2. This notable improvement shows LoSparse is complementary to existing compression methods.