ZipLM: Inference-Aware Structured Pruning of Language Models

Summary

This paper proposed a novel structured compression approach for the inference of LLM, which reduces the model size by removing entire sub-components, like rows or columns from the model’s weight matrices. Based on their algorithm, they manage to drop out some attention heads and shrink the size of FFN layers. Besides, the proposed method considers the runtime speedup while deciding the sparsity of the pruning, which leads to a better trade-off between the loss and the performance.

Strengths

- This work comes up with a new structured pruning method for the transformer models and achieved reasonable speedup in the experiments with multiple models. - The theoretical analysis is solid.

Weaknesses

- The performance of the proposed method seems strongly dependent on the data distribution of inference, which may make it not that practical in real-world inference settings where the incoming inputs can vary a lot. - In the latency-constrained scenario, the authors said "However, for the latter, the inputs are much smaller, and the size of weight matrices is no longer the primary bottleneck.", however, in my opinion, when the batch size is very small (like 1), loading weights from GPU HBM is exactly the bottleneck. Therefore, I guess sparsifying model parameters should lead to reasonable speedups? - In the throughput-constrained (aka. large batch size), it will be interesting to see how ZipLM could support an even larger batch size since we can take advantage of the memory saved by pruning to hold more inputs. In this way, we may further increase the throughput?

Questions

Please refer to the Weekness

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The potential social impact is not discussed in the paper.

Summary

This paper looks into structure pruning (e.g. neurons, or attention heads) of Transformer based models. Authors propose to formulate the pruning objective by requiring the output of each pruned layer to be as close as possible to that of an unpruned layer. Then they adopt optimal brain surgeon algorithm to come up with the weights to keep (mask) and the update to such weights. The proposed algorithm prunes each group (e.g. a neuron in a layer) one by one, ensuring that correlated groups are properly accounted for (which is hard to do when choosing a number of neurons to prune in one go). Additional extensions to the method include inclusion of inference aware criteria into the algorithm for choosing the groups to keep (using pre-computed knowledge of latency table) Suggested approach can be used for one shot pruning and gradual pruning. Experimental results indicate good performance on attention head pruning, fully connected layers pruning and removing full attention modules.

Strengths

- Really good experimental results - Well written and easy enough to follow (modulo my comments below)

Weaknesses

- Not very clear novelty - seems computationally expensive - Ablations are missing (is the improvement coming from the fact that you use formulation (1) for each layer? is it from line 155 of the algorithm? Is it the inference awareness inclusion?

Questions

1) How does pruning two consecutive layers happens? First pruning first layer (Algorothm 1) and then pruning the next layer? If yes you are essentially assuming that layers are independent (so assuming block diagonal structure of the hessian). 2) In your formulation (1) you seem to suggest that X is the input to the layer (whose weights W you are considering to prune). In the original Optimal Brain Surgeon, X is actually based on the gradients (and the global objective is considered, instead of local l2 objective of matching each layer's output). I think this all needs to be made clearer 2) for experiments, it would be nice to report flocs or any other time measurement for your method (not the inference) and competing methods. 3) Potentially you can apply your idea of pruning only one structure at a time (line 155 in Algorithm 1) to any of the existing methods, so ablation is needed (is your improvement comes from the fact that u do essentially more gradual pruning?) 4) The method you propose does not seem to be LLM specific - why not to compare it with methods on structured pruning for non LLM models - I am confused about contribution on line 56 (produces a family of compressed models). How is it any different from say using gradual pruning with various levels of sparsity (and at each step updating the copy for that sparsity level). You are still saving/updating multiple copies of the models right? Minor: - Algorithm 1: Mask Mr is not defined and is not updated in the body of the algorithm. I think you want to compile it based on R at the end of k iterations (the same also applies to Ms - please define how this mask looks like and its shape)

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

n/a

Summary

The paper proposes ZipLM, a structured pruning method that can achieve desired target runtime speedups.​ The idea of ZipLM is to solve a layer reconstruction problem with a structural constraint which minimizes the output changes on a set of calibration examples if the layer is reconstructed as it. The proposed ZipLM is shown to be effective on both decoder-only models and encoder-only models, outperforming prior distillation approaches and structured pruning approaches.

Strengths

* The paper focuses on structured pruning problems, where the pruned model can easily get real speedups. Moreover, the proposed approach takes as input a target speedup as well as the hardware environment and optimizes the model on this specific setup to ensure the model to achieve desirable speedup. I admire this realistic setting and believe that this can be practically impactful. * The paper has conducted extensive experiments studying the effectiveness of the proposed ZipLM approach, including on both encoder and decoder models, both one-shot pruning setting and gradual compression setting, and with different levels of sparsity. * The experimental results show the proposed approach is more effective compared to the existing baseline methods.

Weaknesses

* I believe the paper (especially the methodology part) can be presented better and a lot more details should be added (even in the appendix). For example, how do you solve the reconstruction problem and obtain the optimal mask and weight update (equation (2) & (3))? How do you consider the constraint C in equation (1) in the solution equation (2) & (3)? What exactly does a latency table look like and how do you get the table given a hardware environment? How many calibration inputs did you use in your experiments and how sensitive the final results will be to the number/quality of the calibration inputs? * The paper is motivated by structured compression of large language models (LLMs). However, all experiments are conducted on models with hundreds of millions of parameters. Given the current state-of-the-art language models (e.g., LLAMA-7B/13B/65B) contain a couple orders of magnitude more parameters compared to BERT, it is not clear how the proposed method works in larger models.

Questions

* Do you think the proposed approach can apply to larger language models such as LLAMA-7B? If not, what is the barrier? * How sensitive the proposed approach is to the calibration examples?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

2 fair

Contribution

3 good

Limitations

I didn’t see a clearly potential negative societal impact of this paper.

Summary

This paper proposes ZipLM, a structured pruning and reconstructing + layer-wise distillation + inference-aware pruning algorithm. The authors first extend the pruning formula of OBS to structured pruning and utilize estimation of inference for each structure and structured SPDY search to achieve more accurate inference-aware pruning. Then, the authors distill the model in a token-wise manner. Experimental results demonstrate that ZipLM achieves state-of-the-art (SOTA) performance in both one-shot and pruning/knowledge distillation settings.

Strengths

1. ZipLM demonstrates exceptional results among existing compression algorithms. It not only outperforms retraining-constrained approaches but also surpasses the performance of some pruning + knowledge distillation algorithms. 2. The authors propose a systematic framework for structured pruning, encompassing aspects such as pruning and reconstruction, inference-aware structure search, and knowledge distillation for performance recovery. This comprehensive framework exhibits a well-designed structure that, in my opinion, contributes to the entire community.

Weaknesses

1. The methodology is somewhat incremental. The core contributions of the authors revolve around extending previous methods (OBS, SPDY) to structured pruning since those methods couldn't be directly applied. While the authors have made these extensions, the core essence of the method remains largely based on the previous framework. 2. The experimental section of this paper has some shortcomings. For instance, although the authors conducted a simple ablation experiment, they did not analyze which specific parts of the framework played a crucial role in improving performance. Additionally, in the majority of compression papers, the GLUE benchmark's eight datasets are evaluated as a whole since they represent the fundamental measure for assessing the performance of compressed models. However, the authors only evaluated four of these datasets.

Questions

1. Could you provide a more comprehensive ablation experiment, specifically to identify which module within the entire framework contributes significantly to performance improvement? 2. Since Kwon [1] also employed a pruning metric and the layer-wise reconstruction via LLS for structured pruning, have you directly compared your pruning metric and the performance after reconstruction of OBS with their method to determine which one performs better, apart from the effect of distillation? 3. Line 192-193: For example, a 95% sparse BERT produced by CoFi has 12x speedup on a V100 GPU, but only 5x on an A100 GPU. Since this is an interesting observation that FLOPs-based calculation != inference speed, can you explain what is the potential reason of this?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This paper proposes a structured compression method to optimize inference efficiency for language models. The proposed method uses the accuracy-efficiency trade-off under specific inference objectives as the importance measure of the model component( attention head, MLP neurons, entire FC/attention blocks.) Extensive evaluation results are provided to demonstrate the effectiveness of the proposed method.

Strengths

I appreciate the thoroughness of the evaluation, in which the authors consider the encoder/decoder and decoder-only model, retraining compression and zero-shot compression, throughput and latency objective, etc. Evaluation provides a better understanding of the proposed method.

Weaknesses

This paper focuses on improving the language model's inference efficiency. Other than distillation and pruning, quantization is also a popular direction[1], which is not discussed/compared. Quantization generally reduces the model size. At the same time, a reduction in model size would require less number of GPUs for large model inference, leading to latency speed up. Further, there is a line of work on dynamic sparsity[2] that improves inference latency, which is also not discussed/compared. [1]Frantar, Elias, et al. "Gptq: Accurate post-training quantization for generative pre-trained transformers." arXiv preprint arXiv:2210.17323 (2022). [2]Liu, Zichang, et al. "Deja vu: Contextual sparsity for efficient LLMs at inference time." (2023).

Questions

Table 1 shows a significant drop in PPL compared to the original GPT2, even at the smallest speed-up. To better understand the trade-off, can the authors comment on at what speedup we can expect no loss in performance?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors didn't discuss limiation or negative social impact.

Rebuttal Response

Thanks for the pointer to the quantization experiment and the additional perplexity of GPT2 trained under the same setup. A discussion of quantization LLM would be nice to include in Related Work, aside from distillation and pruning.

Discussion reminder

Dear Reviewers, Given that the discussion period is ending shortly, we wanted to send a gentle reminder regarding our review responses. We would really appreciate your feedback on our rebuttal, especially regarding the added experimental results and additional clarifications regarding our method. With best regards, The ZipLM Authors