BESA: Pruning Large Language Models with Blockwise Parameter-Efficient Sparsity Allocation

Thank you for your thoughtful feedback. We appreciate your engagement with our rebuttal and the points you've raised. We have carefully considered each of your concerns and addressed them as follows:

Q1: "Authors claimed BESA has different advantages ... does not answer these questions."

A1: Thank you for your insightful suggestion. In response, we have conducted experiments using three additional pruning techniques commonly employed in computer vision and natural language processing. These methods, namely iterative pruning (proposed in the Lottery Ticket Hypothesis [A]), threshold pruning [B], and direct mask learning, have been implemented in a blockwise manner for a fair comparison. For iterative pruning, we employ a step-wise approach where we iteratively prune the LLM block with a step of 10% according to the metric Wanda or Weight magnitude. During each iteration, we update the weights with the reconstruction loss. The results are reported in the table under "Iterative Prune". Then for threshold pruning, we learn a threshold value for each layer within the block while maintaining the block sparsity at 50%. The results are reported in the table under "Threshold Prune". For direct mask learning, we directly learn a sparse mask for the LLM block. The results are reported in the table under "Direct Mask Learning". We opt for iterative pruning as it stands out as one of the most representative methods for pruning, incorporating weight updating. Additionally, we select threshold pruning for its cutting-edge capabilities, showcasing state-of-the-art performance in computer vision models. Lastly, we choose direct mask learning as it represents the most straightforward and directly differentiable pruning method. Notably, the original “Iterative prune” and “Threshold prune” methods use the weight value as a measure of importance. In contrast, we provide the versions that use the weight or Wanda importance metric.

The perplexity results on WikiText2, C4, and PTB datasets are presented below:

These results underscore the unique advantages of our method, particularly its ability to achieve substantial reductions in perplexity. We believe that this thorough exploration of differentiable methods provides a comprehensive understanding of the strengths and capabilities of BESA compared to other baselines.

Additionally, we have re-implemented our algorithm in the style of structured pruning and compared it with LLM-Pruner [C], the latest structured pruning algorithm we found. The comparison includes WikiText2 perplexity and zero-shot downstream task performance and is conducted on LLaMA-7B with 20% sparsity without tuning. The table below shows the results, highlighting that our method significantly reduces WikiText2 perplexity by at least 40%, while maintaining comparable performance in zero-shot downstream tasks compared to LLM-Pruner's various weight importance settings.

If you have any further inquiries or require additional clarification, please don't hesitate to ask.

[A] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019.

[B] Kusupati, Aditya and Ramanujan, Vivek and Somani, Raghav and Wortsman, Mitchell and Jain, Prateek and Kakade, Sham and Farhadi, Ali. Soft Threshold Weight Reparameterization for Learnable Sparsity. Proceedings of the International Conference on Machine Learning, 2020.

[C] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LLM-Pruner: On the Structural Pruning of Large Language Models. Advances in Neural Information Processing Systems, 2023.

Thank you for your thoughtful feedback. We appreciate your engagement with our rebuttal and the points you've raised. We have carefully considered each of your concerns and addressed them as follows:

Q1: "The paper does not provide a detailed analysis of the trade-offs between different pruning rates and their impact on model performance, ... "

A1: We apologize for any confusion caused by the placement of ablation studies in our paper. On page 7, in Section 4.2, we discussed and conducted experiments with varying sparsities, which are visually represented in Fig. 3. However, we acknowledge that the figure is presented on page 12 due to the limit of space. From Fig.3, we can see that BESA consistently outperforms other baselines such as SparseGPT and Wanda. The gap is particularly pronounced at large sparsities.

Q2: "The paper does not provide a ... structured pruning ... of the proposed method."

A2: Your suggestion to include a comparison with structured pruning methods is valid, and we appreciate the insight. Therefore, we conducted experiments about comparing our method with LLM-Pruner [A], the latest structured pruning algorithm we found. After checking the experiments shown in LLM-Pruner, we observed a dramatic decrease in both WikiText2 perplexity and zero-shot performance of some downstream tasks with structured pruning. For a fair comparison, we reimplemented our method in the manner of structured pruning. Specifically, we block-wise pruned the LLM model in the dimension of the attention module's head and the MLP module's hidden dimension. The detailed comparison results are shown in Table 1. The experiment was performed in LLaMA-7B with a sparsity of 20% without tuning. From the table below, it is evident that our method can reduce WikiText2 perplexity by at least 40% while achieving comparable performance in zero-shot downstream tasks compared to LLM-Pruner's different weight importance settings.

Table 1: Evaluation of WikiText2 perplexity (ppl, lower is beter) and performance on zero-shot downstream tasks (higher is better) for LLaMA-7B at 20% sparsity without tuning the pruned model.

[A] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LLM-Pruner: On the Structural Pruning of Large Language Models. Advances in Neural Information Processing Systems, 2023.

Q3: "There has been a lot of research on pruning in CV and NLP. The baselines ..."

A3: Thanks for the suggestion. In response, we have conducted experiments using three additional pruning techniques commonly employed in computer vision and natural language processing. These methods, namely iterative pruning (proposed in the Lottery Ticket Hypothesis [A]), threshold pruning [B], and direct mask learning, have been implemented in a blockwise manner for a fair comparison. For iterative pruning, we employ a step-wise approach where we iteratively prune the LLM block with a step of 10%. During each iteration, we update the weights with the reconstruction loss. The results are reported in the table under "Iterative Prune". Then for threshold pruning, we learn a threshold value for each layer within the block while maintaining the block sparsity at 50%. The results are reported in the table under "Threshold Prune". For direct mask learning, we directly learn a sparse mask for the LLM block. The results are reported in the table under "Direct Mask Learning". We opt for iterative pruning as it stands out as one of the most representative methods for pruning, incorporating weight updating. Additionally, we select threshold pruning for its cutting-edge capabilities, showcasing state-of-the-art performance in computer vision models. Lastly, we choose direct mask learning as it represents the most straightforward and directly differentiable pruning method. Notably, the original “Iterative prune” and “Threshold prune” methods use the weight value as a measure of importance. In contrast, we provide the versions that use the weight or Wanda importance metric.

The perplexity results on WikiText2, C4, and PTB datasets are presented below:

These results underscore the unique advantages of our method, particularly its ability to achieve substantial reductions in perplexity. We believe that this thorough exploration of differentiable methods provides a comprehensive understanding of the strengths and capabilities of BESA compared to other baselines.

[A] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019.

[B] Kusupati, Aditya and Ramanujan, Vivek and Somani, Raghav and Wortsman, Mitchell and Jain, Prateek and Kakade, Sham and Farhadi, Ali. Soft Threshold Weight Reparameterization for Learnable Sparsity. Proceedings of the International Conference on Machine Learning, 2020.

Thank you for your thoughtful feedback. We appreciate your engagement with our rebuttal and the points you've raised. We have carefully considered each of your concerns and addressed them as follows:

Q1: "It would be better to involve the computational cost of attention weight in Table 4."

A1: We acknowledge the importance of providing information on the computational cost of running our pruned model in ViTCoD. In response to your suggestion, we have revised Table 4 to include the FLOPs (floating-point operations) of the pruned model. The updated table now offers a more comprehensive view of the efficiency gains achieved by our method. Please refer to the revised table below.

Table 4: Runtime (cycles), computational cost (FLOPs), and speedup across various layers in LLaMA-7B. The term "FLOPs" refers to floating-point operations while "cycles" denotes the number of instruction cycles necessary for the ViTCoD accelerator to perform the associated computational workloads.

Q2: "It would be better to have an ablation study of block-wise pruning. Maybe directly pruning all the models instead of layer by layer."

A2: Good suggestion. Your insightful question highlights the potential for valuable insights by examining the direct pruning performance across all models. Regrettably, conducting such an experiment proves challenging due to computational limitations, even with our exploration of state-of-the-art strategies such as microsoft's DeepSpeed, the int-8 training scheme, and the int8 optimizer provided by the bitsandbytes library. We believe that directly pruning all the models instead of layer by layer could lead to better performance, which can be derived from the experiments in Table 6 and Figure 5 of the Appendix. Table 6 shows that performing our BESA with a larger learning granularity leads to better performance, which is led by the lower quantization error when pruning the model with a larger learning granularity. Hence, directly pruning all the models would result in a lower quantization error and thus better performance.

Thank you for your thoughtful feedback. We appreciate your engagement with our rebuttal and the points you've raised. We have carefully considered each of your concerns and addressed them as follows:

Q1: "There may be a few things ... add the meaning of “W” together with “M, F, X, …”;"

A1: Sorry for any confusion caused by our paper. To clarify, "W" represents all linear weights in a transformer block, "M" refers to the learnable binary mask of these weights, "X" means the input tokens, and "F" is the mapping function that gets the forward result of given "W" and "X".

Q2: "Considering the abstract mentions existing methods ... this paper more convincing."

A2: Thank you for pointing it out. We provide a comprehensive ablation study of the hyperparameters used in our method on page 12, Appendix Sec. A, which is inclusive of the size of the calibration set, epochs, the step of sparsity candidates, learning granularity, and importance metric.

Q3: "Due to the uniqueness of the proposed method BESA, ... make model faster than others?"

A3: Good question. Comparing the runtime of the model pruned by our method and other methods in the ViTCoD accelerator can make our experiments more reasonable and convincing. An updated version of Table 4 in our original paper is shown below.

Table 4: Runtime (cycles) and speedup across various layers in LLaMA-7B. The term "cycles" denotes the number of instruction cycles necessary for the ViTCoD accelerator to perform the associated computational workloads.

Based on the results shown in the previous table, although our method shares the same sparsity with other methods, we obtain better performance in the ViTCoD accelerator because of our nonuniform sparsity within the block. As shown in Appendix Sec.B, ViTCoD incorporates a dense engine and a sparse engine to perform computation simultaneously and assign the computation workloads to different engines with respect to their sparsity, the model pruned by our methods can have a more balanced workload assignment.