Streamlining Redundant Layers to Compress Large Language Models

LLMs primarily utilize a Transformer architecture, consisting of a series of Transformer decoder layers. These layers adopt a residual structure, so the effect of each Transformer layer can be viewed as a transformation of the input hidden states. Assuming that the parameters of the $\ell$-th layer $f$ are denoted as $\theta^{(\ell)}$, and its input hidden states are represented by $\bm{x}^{(\ell)}$, the layer $f$ can be expressed as

In Eq. 1, the $\ell$-th layer $f$ contributes a transformation $f(\bm{x}^{(\ell)},\theta^{(\ell)})$ to the input $\bm{x}^{(\ell)}$. Therefore, we assess the importance of each layer in LLMs by evaluating its impact on the input hidden states. We use the cosine similarity $\cos{(\cdot,\cdot)}$ between input $\bm{x}^{(\ell)}$ and output $\bm{x}^{(\ell+1)}$ as a metric. Essentially, a higher cosine similarity between the input and output of a layer indicates lower importance, and vice versa. This interpretation arises from the observation that a high cosine similarity suggests the layer’s transformation is minimal, making it more amenable to pruning.

To measure the importance of different layers in LLMs, we randomly select samples from the pre-training data (details in Section 4.1). We then record the hidden states generated by the LLMs for these samples and compute the cosine similarity between the input and output hidden states of each layer. The computation of cosine similarity can be formalized as follows,

where $\mathcal{D}$ denotes the recorded hidden states from different samples, $\bm{x}^{(\ell)}_{i}$, $\bm{x}^{(\ell+1)}_{i}\in R^{d\times L}$ denotes the input and output hidden states of the $i$-th sample respectively, $d$ denotes the hidden size and $L$ denotes the sequence length of each sample.

To mitigate the effects of model size and model structure, we conduct experiments on four models OPT-1.3B, OPT-2.7B, OPT-6.7B (Zhang et al., 2022), and Llama2-7B (Touvron et al., 2023). The results, illustrated in Fig. 2, show that there is high cosine similarity between the input and output of several consecutive layers in all models, indicating a low level of importance.

Discussion I: Why not use other similarity to measure the importance of layers? In deep learning, cosine similarity is widely employed to measure the similarity between two vectors (Chen et al., 2020; Chen & He, 2021; Reimers, 2019). Alongside it, dot product and Euclidean distance are also utilized, but they additionally consider vector magnitude. Current research suggests that the hidden states of Transformers with the Pre-Norm architecture tend to grow as the depth of layers increases (Liu et al., 2023). This trend leads to a bias where deeper layers in the model have higher dot product similarity, while earlier layers have smaller Euclidean distances. Consequently, we opt for cosine similarity, which is agnostic to the magnitude of the vectors.

Discussion II: Why not use perplexity as the metric to measure the importance of layers? Some concurrent layer pruning work uses perplexity as the metric to measure the importance of layers (Song et al., 2024; Kim et al., 2024). Specifically, they remove each layer one at a time, measuring the change in perplexity of the model on the pre-training data, and eliminate the layer that causes the least change. This process is repeated multiple times to remove several layers. However, we think perplexity is a highly data-sensitive metric, which results in different layers being removed when pruning with different pre-training data. This also results in a situation where, although the perplexity of the pruned model on the pre-training data used for pruning is low, it performs poorly on other datasets. In contrast, the cosine similarity is highly stable and always leads to the same pruned layers on different pre-training data. We conduct detailed experiments in the Appendix A to demonstrate that perplexity is a highly data-sensitive metric and performs poorly on downstream tasks.

As Fig. 2 shows, the less important layers are often contiguous. Hence, given number of pruned layers $n$ determined by a target sparsity, we remove $n$ contiguous layers by finding the initial layer $\ell^{*}(n)$ corresponding to the highest cosine similarity for pruning:

where we randomly select samples from the pre-training data to compute the cosine similarity between $\bm{x}^{(\ell)}$ and $\bm{x}^{(\ell+n)}$, as outlined in Eq. 2.

After the layer pruning process, we aim to replace the pruned layer with a lightweight network that has much fewer parameters. The rationale is that these layers contribute only minor transformations to the input. Therefore, we hypothesize that the cumulative effect of these layers can be approximated by a lightweight network. Specifically, after identifying the initial layer $\ell^{*}(n)$ for pruning, we use $(\bm{x}^{(\ell^{*})},\bm{x}^{(\ell^{*}+n)})$ as the training data to train the lightweight network using mean squared error (MSE) loss, which can be formalized as follows:

where $h$ denotes the lightweight network, $\mathcal{D}$ denotes the recorded hidden states of samples.

Discussion: Layer Replacement or Fine-Tuning Pruned LLMs? Here, we discuss why opt for layer replacement, instead of using common Parameter-Efficient Fine-Tuning (PEFT) methods such as LoRA  (Hu et al., 2021) and QLoRA  (Dettmers et al., 2023) after layer pruning.

First, from the perspective of resource overhead, layer replacement is more adaptable to hardware resource constraints compared to other methods. Fine-tuning the model using the PEFT methods requires storing the model’s weights, activation values, and the optimizer state of the PEFT module in the GPU during training. In contrast, layer replacement involves two stages: dataset construction and model training. The first stage only requires storing the model’s weight and the forward computation overhead, and the second stage only requires storing of the lightweight network’s weight, activation values of lightweight network, and the optimizer state of lightweight network. Therefore, layer replacement can also be implemented under conditions of hardware resource constraints.

Second, layer replacement uses a lightweight network to replace the pruned layer, and distills the knowledge of the pruned layer into the lightweight network using the MSE loss function. Unlike layer replacement, we speculate that training the model after pruning with LoRA is a process of redistributing the function of the pruned layers across the remaining layers. Therefore, substituting the pruned layers with a lightweight network could be a less challenging training task than redistributing the function of the pruned layers across the remaining layers. In the experiments of Section 4.7, we demonstrate that layer replacement has better performance compared to LoRA.

When evaluating the natural language understanding capabilities of LLMs, most existing benchmarks frame the task as a classification task. A classification task with $k$ choices and comprising $N$ samples is denoted as $\mathcal{T}=\{\left(x_{i},c_{i,1},c_{i,2},...,c_{i,k},y_{i}\right)\}_{i=1}^{N}$, where $x_{i}$ represents the question in the $i$-th sample, $c_{i,j}$ represents the $j$-th choices, and $y_{i}$ represents the correct choice. The input to the classification task consists of a question accompanied by multiple choices, and the LLM is required to select the correct answer from these choices. Typically, each choice is concatenated with the question to form multiple sentences, and the perplexity (PPL) of each sentence is computed. The choice corresponding to the sentence with the lowest PPL is selected as the answer.

Typically, model pruning results in decreased model performance. However, when we evaluate the model pruned by the method described in Sec 2.2, we unexpectedly observe the accuracy of the pruned model has been improved on some classification tasks. We define $\mathcal{M}$ to denote the original LLM, $\bar{\mathcal{M}}$ to denote the compressed LLM, and $\hat{y}(\mathcal{M})$ to denote the choice predicted by the model $\mathcal{M}$. To further investigate this phenomenon, we analyze the experimental results using the confusion matrix  (Li et al., 2024). Specifically, we count the number of samples and average standard deviation (std) for the PPL of the samples for each term of the confusion matrix. The calculation of the std for the PPL of the $i$-th sample is defined as follows:

where ${\rm PPL}_{i,j}$ denotes the PPL for the sentence created by question $x_{i}$ and choice $c_{i,j}$ of the model before pruning, ${\rm std}_{i}$ denotes the std for PPL of the $i$-th sample. A higher $\mathrm{std}_{i}$ value indicates the LLM exhibits greater confidence in answering the question $x_{i}$.

Each term of the confusion matrix is defined as follows,

Table 1 presents the counts of samples in TP, FN, FP, TN in sevaral datasets, and also the PPL standard deviation in Race-H dataset. We can observe that the std for ${\rm TP}$ and ${\rm TN}$ is significantly higher than that for ${\rm FN}$ and ${\rm FP}$. This indicates that the model is more uncertain about the ${\rm FN}$ and ${\rm FP}$ samples. In addition, the samples in ${\rm FP}$ constitute a considerable proportion of the total samples, implying that the model may guess the correct answer for a significant portion after pruning. This phenomenon suggests that the accuracy metric may overestimate the performance of the compressed model.

We propose a novel metric stability to evaluate the performance of LLMs after pruning, i.e.,

where the identifier $\mathbbm{1}_{\left[i\in{\rm TP}\cup{\rm TN}\right]}$ is used to indicate whether the $i$-th sample belongs to ${\rm TP}$ and ${\rm TN}$. We use ${\rm std}_{i}$ as the weight of the $i$-th sample. Because the ${\rm std}$ of different samples varies significantly, to mitigate the influence of samples with excessively large standard deviations, we apply the $\exp$ function to moderate the weight differences among samples. Different from accuracy, stability focuses on the model’s confidence in its answers and the consistency between the model before and after pruning on a task, aligning more closely with the goal of model pruning, i.e., ensuring the pruned model remains as similar as possible to the original model.

We conduct experiments on popular open-source LLMs, including Llama2-7B and Llama2-13B (Touvron et al., 2023). Following previous work (Men et al., 2024; Yang et al., 2024), we perform experiments with a 25% pruning rate and extract data from pre-training dataset SlimPajama which contains data from different domains for layer pruning and layer replacement. Sheared LLaMa  (Xia et al., 2023) finds that the performance degradation of pruned models varies across different domains, and proposes determining the distribution of data from different domains based on the degree of performance degradation. Therefore, we randomly sample the data based on the distribution used by Sheared LLaMa  (Xia et al., 2023), finally constructing the dataset containing 30,000 pieces of data. We randomly select 500 samples from this dataset and input them into LLMs, generating Fig. 2, and use these 500 data samples for layer pruning. All 30,000 pieces of data are used to train the lightweight network. We utilize two types of lightweight networks: a Feed-Forward Neural Network (FFN), referred to as Ours (FFN), and a Transformer Layer, referred to as Ours (Layer). The FFN is randomly initialized, while the Transformer Layer inherits the parameters from the first pruned layer. Additionally, we explore a purely pruning approach without incorporating any lightweight network, denoted as Ours (None). Further experimental details are available in the Appendix C.1.

We use 12 natural language understanding benchmarks for evaluation: CMNLI  (Xu et al., 2020), HellaSwag(HeSw)  (Zellers et al., 2019), PIQA  (Bisk et al., 2020),CHID  (Zheng et al., 2019), WSC  (Levesque et al., 2012),CommonSenseQA(CoQA)  (Talmor et al., 2018), BoolQ  (Clark et al., 2019),MMLU  (Hendrycks et al., 2020), CMMLU  (Li et al., 2023),Race-High/Middle  (Lai et al., 2017), C3  (Sun et al., 2020). The tasks in these benchmarks are formalized as classification tasks, so we refer to these benchmarks as classification benchmarks. For these benchmarks, we use both accuracy and stability as metrics for evaluating the models. Additionally, we include 3 benchmarks: XSum  (Narayan et al., 2018), GSM8K  (Cobbe et al., 2021) and StrategyQA  (Geva et al., 2021), to demonstrate the LLM’s performance on generation tasks after pruning. We refer to these tasks as generation benchmarks. Following the evaluation framework of OpenCompass (Contributors, 2023), we use accuracy to evaluate StrategyQA and GSM8K, and use ROUGE1 to evaluate Xsum.

We compare several pruning methods that prune the attention heads, the filters of the FFN layer, and the hidden dimension, as well as the concurrent layer-pruning methods LaCo. In addition, ShortGPT and UIDL (Men et al., 2024; Gromov et al., 2024) can be considered as a variant of our approach, i.e., Ours (None). We also discuss layer pruning methods which use perplexity as the metric, such as SLEB (Song et al., 2024), in Appendix A.

LLM-Pruner (Ma et al., 2023) prunes attention heads, FFN layer filters, and hidden dimensions by using gradients and activations to estimate the importance of these modules.

SliceGPT (Ashkboos et al., 2024) prunes hidden dimensions. It inserts dimensionality reduction matrices into the model and employs Principal Component Analysis (PCA) to initialize and compress the matrices, and then merge them with the original weight matrix to reduce the model’s size.

LaCo (Yang et al., 2024) prunes layers by dividing the layers into groups, each consisting of multiple consecutive layers, and compresses them separately, whereas our method simply compresses a piece of consecutive layers. LaCo merges consecutive layers by averaging their parameters whereas we train an additional lightweight network to replace these layers.

We present accuracy and stability for different methods on the classification benchmarks in Table 2 and Table 3, respectively, and Table 4 for the generative benchmarks. The results demonstrate that our proposed LLM-Streamline consistently outperforms the baseline methods. Specifically, in classification tasks, LLM-Streamline surpasses LLM-Pruner by 7% in accuracy and 12% in stability on Llama2-7B, and by 16% in accuracy and 19% in stability on Llama2-13B. LLM-Streamline also surpasses LaCo by 5% in accuracy on Llama2-7B. For generation tasks, LLM-Streamline retains nearly 77% of Llama2-7B and Llama2-13B’s capabilities, significantly outperforming other pruning methods. We find that almost all of the pruning methods fail on the GSM8K dataset. However, sufficient training can gradually restore the model’s performance on math tasks, and the specific experimental results are shown in Table 28 of Appendix E.8.

Additionally, comparing the average stability (Average) in Table 3 with the retrained performance (RP) in Table 2 reveals that stability is often much lower than accuracy. We also observe that accuracy on Race-M and Race-H even increases after model pruning. Furthermore, we find that without using any lightweight network, Llama2-13B achieves the highest accuracy on classification benchmarks, but its stability on classification benchmarks and performance on generation benchmarks are lower. These results indicate that pruned models tend to make correct guesses on some classification questions that they are uncertain, highlighting the limitations of accuracy as a sole measure of pruning method performance. We also conduct experiments on OPT-1.3B, OPT-2.7B, OPT-6.7B, Baichuan-7B, Baichuan-13B, Baichuan2-7B, Baichuan2-13B (yang2023baichuan), Llama3.1-8B, Llama3.1-70B (dubey2024llama) and Mixtral-8x7B-v0.1 (jiang2024mixtral). Details can be found in Appendix E.1, Appendix E.2, Appendix E.3, Appendix E.4, Appendix E.5 and Appendix E.6. In addition, we compare the performance of LLM-Streamline with other methods at a higher pruning ratio of approximately 50%, and the results can be found in Appendix E.7.

While FFN achieves the best result, Transformer layer still has performance potential. We perform experiments with Llama2-7B using various lightweight network, including Feed-Forward Neural Networks (FFN), SwiGLU-based Feed-Forward Neural Networks (SwiGLU), and Transformer layers. We also explore various initialization methods for the Transformer Layer, including random initialization (Layer-Random), inheritance of the first pruned layer (Layer-First), inheritance of the last pruned layer (Layer-Last), and averaging the pruned layers (Layer-Avg). The average accuracy and stability metrics across all the classification benchmarks are presented in Table 5, with detailed results on each benchmark available in Appendix E.9. The results show that FFN achieves the best results. Meanwhile, for the Transformer Layer, inheriting the pruned first layer yields the best results. In contrast, the performance of Layer-Avg, inspired by LaCo, shows that averaging weights does not achieve the same effectiveness as the pruned first layer. In addition, we plot the validation loss curves for different lightweight networks, as shown in Fig. 3. We can observe that FFN and SwiGLU have already converged by the 10th epoch, whereas the loss of Transformer Layer is still decreasing. This indicates that the Transformer layer still has potential, and further training could yield better results, but this would require more computing resources.

The performance of the pruned model is linearly correlated with the number of parameters at modest pruning ratios. To verify the model’s performance at various modest pruning ratios, we evaluate our method not only at the approximately 25% pruning ratio but also at ratios of around 2%, 11%, and 20% on Llama2-7B. The average stability and accuracy metrics across all the classification benchmarks are shown in Fig. 4, with details on each benchmark presented in Appendix E.10. By comparing the performance of the original Llama2-7B, TinyLlama-1.1B, OpenLlama-3B-v2, and Llama2-7B pruned at various ratios, we observe a linear correlation between the performance of both the pruned models and the pre-trained original models relative to the number of parameters. This suggests that the performance of models pruned using our method is comparable to that of pre-trained models with the same number of parameters.

Layer Replacement outperforms LoRA in both performance and GPU memory consumption. We compare the performance of layer replacement with LoRA. Since layer replacement is trained based on hidden states with a different training objective than LoRA, we additionally train one epoch using the language model loss for layer replacement when comparing it with LoRA. The training details can be found in the Appendix C.2. For layer replacement, we freeze the original model’s weights and train only the lightweight network. In the case of LoRA, we set the rank to 128 to align the number of parameters trained with those of the lightweight networks. We randomly extract 30,000 samples from SlimPajama-6B for layer replacement training and also test with the entire dataset to evaluate the limited impact of extensive data on performance (details in Appendix E.8). For LoRA, we use 300,000 samples from SlimPajama-6B. Table 6 presents the average accuracy and stability across all classification benchmarks, with detailed results available in Appendix E.11. The findings indicate that layer replacement surpasses LoRA in both accuracy and stability, while also requiring significantly less GPU memory and training data.

For both the FFN structure and the SwiGLU structure, the learning rate is set to 1e-3 and the weight decay is 1e-4. For the Transformer layer, the learning rate is set to 1e-5 and the weight decay is 1e-3. The model is trained using a batch size of 32 over 20 epochs. On a single A800 GPU, the training duration for the lightweight network is approximate 5 hours (for the Transformer layer).

For layer replacement, in order to have a fairer comparison with LoRA, we conduct one epoch of post-training with a learning rate of 5e-5, a weight decay of 1e-3, and a batch size of 32. This process takes less than an hour on a single A800 GPU. For LoRA, the model is trained one epochs with a learning rate of 1e-4, a weight decay of 1e-3, and a batch size of 32. Since the amount of training data used is ten times that of layer replacement, it take approximately 10 hours to complete the training on a single A800 GPU.

We conduct experiments on OPT-6.7B. The experimental results are shown in Table 12, Table 13 and Table 14. The results indicate that our proposed LLM-Streamline is superior to the previous SOTA method.

We conduct experiments on Baichuan-7B and Baichuan-13B. The experimental results are shown in Table 15, Table 16 and Table 17. The results indicate that our proposed LLM-Streamline is superior to the previous SOTA method.

We conduct experiments on Baichuan2-7B and Baichuan2-13B. The experimental results are shown in Table 18, Table 19 and Table 20. The results indicate that our proposed LLM-Streamline is superior to the concurrent SOTA method, LaCo.

We conduct experiments on Llama3.1-8B and Llama3.1-70B. The experimental results are shown in Table 21 and Table 22. The results indicate that our proposed LLM-Streamline is superior to the previous SOTA method.

We conduct experiments on Mixture of Experts(MoE) model Mixtral-8x7B-v0.1. The experimental results are shown in Table 23 and Table 24.

We also conduct experiments on small models (OPT-1.3B and OPT-2.7B). The experimental results are shown in Table 25. The results indicate that our proposed LLM-Streamline is superior to the previous SOTA method, across different pruning rates.

We conduct experiments on Llama2-7B at a higher pruning ratio of approximately 50%. The experimental results are shown in Table 26 and Table 27. The results indicate that our proposed LLM-Streamline is superior to the previous SOTA method.

Following the method outlined in Section 4.7, We conduct experiments using the entire SlimPajama-6B for post-training, and the results are presented in Table 28. As shown, using the entire dataset resulted in a slight improvement, but at a significant computational cost, requiring 100 times the computational time.

The detailed results of accuracy and stability of different lightweight networks on different classification benchmarks are shown in Table 29 and Table 30. We can observe that FFN achieves the best results. Meanwhile, for the Transformer Layer, inheritance of the pruned first layer yields the best results.

The detailed results of accuracy and stability on LLMs under different pruning ratios on different classification benchmarks are shown in Table 31 and Table 32. The experiment results show that the performance of the pruned model is linearly correlated with the number of parameters, demonstrating the effectiveness of our method.

The detailed results of accuracy and stability of different layer replacement strategies and LoRA on different benchmarks are shown in Table 33 and Table 34. The results show that layer replacement outperforms LoRA in both accuracy and stability.