LaCo: Large Language Model Pruning via Layer Collapse

For the $l$-th layer of an LLM, we denote all its parameters, including those in self-attention (SAN) and MLP as $\bm{\theta}_{l}$. For the $m$ consecutive layers following it, we merge the parameters of $\bm{\theta_{l+1}},\bm{\theta_{l+2}},\cdots,\bm{\theta_{l+m}}$ into $\bm{\theta_{l}}$ to form $\bm{\theta^{*}_{l}}$:

where $(\bm{\theta_{l+k}}-\bm{\theta_{l})}$ is the layer-wise parameter difference. Given identical layer structures, we independently apply these processes to both SAN and MLP. Then, these $m$ consecutive layers will be discarded. Subsequent model pruning will continuously involve RDSC Layer Merge which can be regarded as the continual collapse of layers onto specific layers, hence the name Layer Collapse.

We dynamically merge adjacent layers from the topmost layer down, ensuring the pruned model’s output representation on few-shot calibration samples remains as similar as possible to the original model to minimize performance loss. Algorithm 1 summarizes the workflow of Layer Collapse:

For an LLM $\mathcal{M}$ to be pruned, we define the number of layers to be merged during each merging operation as $\mathcal{C}$. We configure the merging to operate within a certain range of layers, denoted as $[\mathcal{L},\mathcal{H}]$. As the layer merging operation inevitably leads to a performance loss, to prevent consecutive layer merging from causing a sharp decline in the model performance, we set a minimum interval of layers between two merging operations as $\mathcal{I}$. Few-shot calibration samples $\mathcal{D}$, typically a few plain sentences, are used during the pruning process. We perform forward computations on $\mathcal{D}$ with both the pruned and original models to obtain the output representations and ensure that the similarity of representations is not less than the threshold $\mathcal{T}$.

We present an illustration of Layer Collapse in Figure 2. We begin by initializing the model $\mathcal{M}^{*}$ with the model $\mathcal{M}$ and set a layer pointer $l$ to start from $\mathcal{H}-\mathcal{C}$. Then, the iterative process begins:

RDSC Layer Merge (line 4-5) During each iteration, our approach involves merging the $\mathcal{K}$ layers following layer $l$ into layer $l$ itself and then discarding the redundant $\mathcal{K}$ layers, where $\mathcal{K}$ is the minimum of $\mathcal{C}-1$ and the total layer count of $\mathcal{M}^{*}-l$, implying merging either the subsequent $\mathcal{C}-1$ layers or all layers following $l$, thus to prune the model $\mathcal{M}^{*}$, resulting in the interim model $\mathcal{M}_{tmp}$.

Calculate Similarity (line 6) We process each sentence in $\mathcal{D}$ using forward computations with $\mathcal{M}_{tmp}$ and $\mathcal{M}$ to derive their representations which are the output hidden-states of the last layer of the model. For every sentence, we then calculate the cosine similarity between these representations from both models, averaging these values to obtain the overall similarity score $s$.

Merge Evaluation and Adjustment (line 7-15) Then, we evaluate $s$ against the threshold $\mathcal{T}$. Should $s$ exceed $\mathcal{T}$, the current merge is considered successful. Then, $\mathcal{M}_{tmp}$ is updated to $\mathcal{M}^{*}$ for the next iteration, and the pointer $l$ is adjusted downwards by $\mathcal{I}$ layers. Conversely, $l$ is simply reduced by a single layer. It is important to highlight that the instances may occur where $l$ falls below the total layer count of $\mathcal{M}^{*}$ after a series of successive merges. Consequently, it is required to reset $l$ to the layer count in $\mathcal{M}^{*}-\mathcal{C}$, as illustrated in line 11.

We iterate through the above process until $l$ is less than $\mathcal{L}$ and output the pruned LLM.

The complexity of LaCo primarily depends on model inference. In the worst-case scenario, with $\mathcal{L}$ set to 0 and $\mathcal{H}$ to the total number of layers, if in each iteration the similarity $s$ is less than $\mathcal{T}$, all layers will be traversed. Thus, the worst-case time complexity is $O(\mathcal{H}\times||\mathcal{D}||)$. For example, for Llama2-13B with 40 layers and $||\mathcal{D}||$ consisting of 10 sentences, the maximum number of inference steps would be only 400 sentences, which can be completed within minutes on a single GPU.

To assess the effectiveness of the proposed LaCo, we conduct experiments on popular English LLMs, Llama2-7B and 13B Touvron et al. (2023). Additionally, we test the effectiveness on bilingual LLMs, specifically Baichuan2-7B and 13B Yang et al. (2023), which support both Chinese and English. We leverage the base versions of these LLMs. All these models are decoder-only models based on the transformer architecture.

To comprehensively evaluate the pruned model’s capabilities, we utilized the OpenCompass evaluation framework Contributors (2023). Specifically, following OpenCompass categorization, we conduct evaluations in five aspects: Reasoning, Language, Knowledge, Examination and Understanding. We select several benchmarks from each category. Reasoning: CMNLI Xu et al. (2020), HellaSwag (HeSw) Zellers et al. (2019), PIQA Bisk et al. (2019). Language: CHID Zheng et al. (2019), WSC Levesque et al. (2012). Knowledge: CommonSenseQA (CSQA) Talmor et al. (2018), BoolQ Clark et al. (2019). Examination: MMLU Hendrycks et al. (2021), CMMLU Li et al. (2023). Understanding: Race-High/Middle (H/M) Lai et al. (2017), XSum Narayan et al. (2018), C3 Sun et al. (2020).

We conduct evaluations using official scripts from OpenCompass, all zero-shot or few-shot, without additional training. Two evaluation modes are utilized: perplexity (PPL) and generation (GEN) ²²2opencompass.readthedocs.io/en/latest/get_started/faq.html. For CHID and XSum, we use the GEN mode. For the WSC dataset, we use both PPL (WSC${}_{\text{P}}$) and GEN (WSC${}_{\text{G}}$) modes. The remaining benchmarks are evaluated using the PPL mode. The evaluation results on each benchmark are converted to a score by OpenCompass, where a higher score indicates better performance. OpenCompass provides official evaluation results for the Baichuan2 and Llama2 series. However, to avoid discrepancies resulting from hardware and software environments, as well as potential errors in official results, we reproduce all results to ensure fairness.

Since LaCo involves structured pruning, which directly removes components from LLMs, we select two state-of-the-art (SOTA) structured pruning methods, LLM-Pruner (LLMPru.) Ma et al. (2023) and SliceGPT Ashkboos et al. (2024), as our baselines. These methods have surpassed the previous SOTA sparsity method, SparseGPT Frantar and Alistarh (2023). In our experiments, we set the pruning ratios of baselines to be equivalent to or slightly smaller than LaCo to ensure fairness.

Since previous work mostly set pruning ratios below 30%, we heuristically adjust the hyperparameters to bring the model pruning ratio close to 30%, as shown in Appendix A Table 8. We randomly select 5 sentences from both the English and Chinese Wikipedia datasets for Baichuan2 and 10 sentences from English Wikipedia for Llama2 as few-shot calibration samples. All experiments are conducted on a server with 8 Nvidia A100 80GB GPUs.

In Table 1, we present the results of four LLMs under different pruning methods across various benchmarks. “Dense” represents the official results of the unpruned LLMs in OpenCompass leaderboards, while “Dense^∗” represents our reproduction of the “Dense” results. "LLMPru." and "SliceGPT" correspond to the two baselines, respectively. “Ratio" refers to the overall pruning ratio, namely the proportion of the total number of pruned parameters to that of the unpruned model. “Lay.” denotes the total number of layers in the model.

Comparing Dense and Dense^∗, the results show not much difference, with most discrepancies within 5%. This indicates our experimental setup is error-free. To ensure fairness, we compare the results against Dense^∗ in the subsequent analyses.

Upon comparing LaCo with the baselines, from Table 1, it can be observed that LaCo achieves the best results on most benchmarks, despite our pruning ratio being slightly higher than the baselines.

To provide a more intuitive presentation of the results in Table 1, we compute the average scores of each pruner across all benchmarks (Avg.), the average scores per category (Reas., Lan., Know., Exam., Unde.), and the average performance percentages relative to Dense^∗ across all benchmarks (Per.) in Table 2. Overall, our average scores are significantly higher than the baselines. LaCo shows superior performance in four out of five categories. Though there is a slight dip in Reasoning, it remains comparable. Additionally, LaCo’s average performance percentage across all datasets, relative to Dense^∗, is far superior to the baselines. The average percentage surpasses 80% in three out of four models, with the lowest being 73% on Baichuan2-7B. In contrast, none of the baselines exceed 70%.

To demonstrate the stability of the pruned models by LaCo, we compute the performance percentage relative to Dense^∗ (Appendix D.4 Table 22). LaCo-pruned models maintain performance above 70% on most benchmarks and do not experience crashes, with no performance dropping below 30%.

Notably, on three benchmarks evaluated through GEN mode, CHID, XSUM, and WSC_G, the LLMs pruned by LaCo maintain relatively stable performance, while models pruned by baselines exhibit poorly, with even multiple results becoming 0.00. GEN mode scores are based on the model’s generated sentences, and the models pruned by baselines are prone to producing meaningless repetitive outputs. In Appendix D.5 Table 23, we showcase an example from the Xsum benchmark, where Llama2-7B, pruned by baselines, produces nonsensical repeated outputs, whereas our LaCo yields outputs resembling normal sentences.

We also conduct experiments with Llama2-70B on several benchmarks. The results in Appendix D.1 Table 19 show that LaCo still outperforms the baseline on larger-scale model.

In summary, LaCo is a superior pruner that preserves model performance and demonstrates exceptional stability across various benchmarks. It relies solely on parameter differences and additions, without altering the model’s internal structure, resulting in a concise and efficient pruning solution.

Since PPL is also a commonly used metric for evaluating model performance, we want to compare how LaCo differs from the baseline in maintaining the model’s PPL. We evaluate the PPL of the Llama2-7B model with 27% sparsity using different pruners. The evaluation is performed on 500 sentences selected from Wikipedia, each with a fixed length of 512 tokens.

The results in Table 3 show that the model pruned by LaCo also achieves a lower PPL compared to other baselines, further highlighting the advantage of our LaCo.

To verify that LaCo has lower time complexity and faster pruning speed than the baselines, we compare LaCo with them for 27% sparsity pruning of Llama2-7B on a single A100 GPU. For fairness, we only measure the main pruning process, excluding the time for loading models, loading data, and storing models. The results in Table 4 show LaCo pruning is more efficient compared to the baselines.

We also aim to investigate whether the model pruned by LaCo offers advantages in memory usage and inference speed compared to the models pruned by the baselines. In Table 5, we present the average memory consumption and inference speed of the Llama2-13B pruned models from Table 1 on the English Wiki dataset (The results for all models are in Appendix D.2 Table 20). All models are loaded in Bf16 on a single A100 GPU.

The results indicate that the LaCo-pruned models consume less memory and achieve faster inference speeds. Moreover, while existing baselines may decrease inference speeds compared to the dense model, LaCo does not have this issue.

This section discusses our motivation for merging adjacent layers. Our primary motivation comes from observing that the changes in parameters and output representations between adjacent layers in the LLMs are not particularly significant.

In Figure 4, we show the L2 similarities between the SAN q, k, v matrices of each layer and their counterparts in the subsequent layer, as well as the upscaling and downscaling matrices of the MLP for both Llama2-7B and Baichuan2-7B. The results indicate that the maximum L2 values between corresponding matrices in adjacent layers are generally no more than 200. Given the large sizes of the MLP upscaling (11008x4096) and SAN q, k, v (4096x4096) matrices, the change in each element between adjacent layers is minimal.

In Figure 5 (a), we randomly select 20 sentences from Wikipedia and calculate the cosine similarity between the hidden-states of adjacent layers outputs. The results show that for both Baichuan2-7B and Llama2-7B, the representation similarity between adjacent layers from layers 3 to 28 is typically very close to 1. The high similarity in parameters and representations between adjacent layers leads us to consider that a single layer might replace multiple subsequent layers.

Moreover, the similarity in parameters suggests focusing on the differences between layers. Inspired by previous model merging work Yu et al. (2023); Matena and Raffel (2022), we come up with collecting parameter differences between layers and merging them into a single layer. To verify that RDSC Layer Merge can replace multiple layers with one, we conduct the experiment: we merge every four consecutive layers into one within layers 10 to 19 and evaluate the cosine similarity between the merged layer’s output and the original last layer’s output, as in Figure 5 (b), where the lowest cosine similarity on the 4096-dimensional vectors is above 0.996, confirming the effectiveness of RDSC Layer Merge in preserving representations.

In this section, we explore the performance of LaCo at different pruning ratios. We conduct experiments on Llama2-7B and Llama2-13B, controlling the pruning ratios at approximately 10%, 25% (our main experiments), and around 50% by setting different hyperparameters (as shown in Appendix A Table 9)³³3Further ablation study on the hyperparameters $\mathcal{C}$, $\mathcal{I}$, $\mathcal{D}$, $\mathcal{T}$, different similarity metrics, different merging strategies, different calibration datasets can be found in Appendix B.. We evaluate pruned models accordingly. The average results are shown in Table 7 and the detailed results are shown in Appendix E Table 25.

As the pruning ratio increases, overall model performance decreases. However, from a pruning ratio of around 10-15% to about 25%, the performance does not significantly decline, indicating our method’s stability within this range. Furthermore, at a pruning ratio close to 50%, the model still maintains approximately 70% performance, demonstrating that our method prevents model crashes even with about half of the parameters removed.