Probe Pruning: Accelerating LLMs through Dynamic Pruning via Model-Probing

We sincerely appreciate the reviewer's insightful feedback and positive comments on the methodology of our paper. Please find our responses below:

Response to Weakness 1, Question 1, and Question 3:

We appreciate the reviewer's insightful feedback. We address Weakness 1, Question 1, and Question 3 here, as these three questions are closely related.

First, we would like to discuss what happens when the batch size is small and how history-informed pruning affects the final performance when the sampling sequence length is small. Note that when generating probes, PP prunes both the sequence and batch dimensions by selecting only key samples and tokens. When the batch size is small, PP can reduce the probe's sequence dimension and still maintain comparable results (or even better), as we perform online structured pruning in a batch-wise manner. If there are not many samples, we can reduce the probe's sequence dimension and still achieve similar performance.

We conducted an experiment on LLaMA-2-7B where we decreased the batch size to 5 for both tasks (in the paper, we used a batch size of 20). Moreover, we set the default probe size to 1 sample and 15% of the sequence length (we used 50% in the paper), resulting in approximately 1.5% FLOPs cost. We refer to this scenario as PP-Limit. Furthermore, we conducted an experiment on PP without using history-informed pruning, denoted as PP-Limit w/o History-Informed Pruning, while PP-Limit w/ History-Informed Pruning indicates that we still use history-informed pruning. In the table, PP represents the default setting used in the main text of the paper.

From Table 1, we observe that with history-informed pruning, PP under the limited scenario can obtain results comparable to PP's default setting. For example, at a 40% pruning ratio, PP-Limit w/ History-Informed Pruning achieves a perplexity of 16.8 on WikiText2 and 62.4% accuracy on ARC-e, which is comparable to PP (or even better), which has a perplexity of 16.8 on WikiText2 and 61.7% accuracy on ARC-e. We also observe that without history-informed pruning, the performance is not stable and is worse than with history-informed pruning. For instance, at a 40% pruning ratio, PP-Limit w/o History-Informed Pruning achieves a perplexity of 21.5 on WikiText2 and 50.0% accuracy on ARC-e. Although it performs better than baselines such as FLAP (which has a perplexity of 38.9 on WikiText2) on WikiText2, its accuracy on ARC-e is lower than FLAP's 52.5%.

From these results, we can conclude that PP can handle the small batch size case, as we can reduce the probe's sequence dimension and still maintain comparable performance. Furthermore, history-informed pruning helps improve performance when the sampling sequence length is small, as it allows PP to leverage historical information.

Furthermore, our primary focus is online dynamic structured pruning during the prefilling stage, where the model processes hundreds of tokens before generating the first one. We leave the exploration of probe pruning for the online dynamic structured pruning during the decoding stage to future work.

Table 1: Zero-shot performance of LLaMA-2-7B after pruning attention and MLP blocks without fine-tuning.

We sincerely appreciate the reviewer's insightful feedback and positive comments on the methodology of our paper. Please find our responses below:

Response to Weakness 1:

In Table 5, we show a trade-off between performance improvements and inference speed. We demonstrate that PP achieves the best Performance Runtime Ratio, which quantifies the ratio of performance degradation per unit of runtime reduction, in comparison with FLAP and Wanda-sp. Specifically, PP's Performance Runtime Ratio values are 2.56× (95.65 compared to 37.37) and 2.85× (106.48 compared to 37.37) more efficient than those of FLAP and Wanda-sp, respectively, indicating a significantly lower rate of performance degradation. Conversely, regarding final speedups, the inference speeds of PP are slightly slower but comparable to other structured pruning baselines, such as FLAP and Wanda-sp. It is important to note that inference speed is highly hardware-dependent and implementation-dependent. We believe that the inference speed of PP could be enhanced through further implementation optimizations or other techniques, which we plan to explore in future research.

Furthermore, we have explored the feasibility of running the probe in parallel with the actual computation of earlier pruned blocks. We present these results in Table 1. Here, PP-Parallel represents the approach where, when the actual computation happens on earlier pruned blocks, we generate the probe from the residuals of these earlier pruned blocks and perform the probing. PP represents the default setting used in the main text of the paper. The results show that we can still obtain performance gains and achieve comparable results to PP. For example, at a 40% pruning ratio, PP-Parallel achieves a perplexity of 17.9 on WikiText2, which is close to that of PP and much lower than the 38.9 achieved by FLAP. Furthermore, PP-Parallel achieves 61.4% accuracy on ARC-e, which is close to that of PP and much higher than the 52.5% achieved by FLAP. However, we are just demonstrating the feasibility of further improving PP's inference speed here; the actual parallelism is hardware-dependent and implementation-dependent.

Table 1: Zero-shot performance of LLaMA-2-7B after pruning attention and MLP blocks without fine-tuning.

Response to Weakness 2:

We sincerely appreciate the reviewer's insightful feedback. We have added the sentence "Our implementation loads the full model for dynamic pruning, while other methods load only the pruned version." to the caption of Table 1. The revised sentences is highlighted in red in our edited PDF.

Response to Weakness 3:

We sincerely thank the reviewer for the valuable feedback. We will revise the paper accordingly to improve its readability for future readers.

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC

Response to Weakness 3:

We would like to respectfully point out that the reviewer may have overlooked the main novelty of our paper, which we will clarify here. Our key novelties are primarily presented in two parts: (1) Probing. We leverage the insight that not all samples and tokens contribute equally to the model’s output, and probing a small portion of each batch effectively identifies crucial weights, enabling tailored dynamic pruning for different batches. (2) History-informed pruning with importance-scaled fusion. We strategically merge the probing states with historical states using importance-scaled fusion to capture the essential characteristics of each batch, enhancing the pruning decision.

Regarding the second point, our key merit lies in the strategic combination of probing states with historical states, which captures the essential characteristics of each batch and improves pruning decisions utilizing historical information. The Exponential Moving Average (EMA) is merely a method to update historical states; the reviewer may choose to substitute it with another method.

Furthermore, we have not mentioned contrastive learning at all in the paper, and our paper is irrelevant to contrastive learning. The Memory Bank that the reviewer refers to is fundamentally different from our history-informed pruning. The samples from the Memory Bank are used for calculating contrastive loss in models applying contrastive learning, whereas we utilize importance-scaled fusion to integrate historical states and probing states to guide the pruning decision.

Response to Weakness 4:

We would like to respectfully clarify that our research encompasses evaluations of LLMs released in 2024. Specifically, in Tables 3 and 16, we report experimental results for LLaMA-3-8B, which was released in April 2024, as detailed on https://ai.meta.com/blog/meta-llama-3/. Additionally, our experiments involve multiple model backbones, such as LLaMA-2, LLaMA-3, and OPT, thus demonstrating the versatility of our method across various LLM architectures. Furthermore, we focus on the diversity of our experimental design to demonstrate the broad applicability of our approach; the selection of models is just one part of this initiative.

Regarding the models LLaMA 3.2 1B and 3B, they were released on September 25th, 2024, according to https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/. Given that the submission deadline for ICLR 2025 was October 1st, 2024, the proximity of these dates made it infeasible to include these models in our experiments. We had considered adding experiments with LLaMA-3.2-3B during the rebuttal period; however, this requires using a version of transformers that is greater than or equal to 4.43, while our current setup and implementation utilizes version 4.35.0. Given the substantial time required to adapt the code, we were unable to implement these experiments within the rebuttal phase. We believe the experiments presented in the paper adequately demonstrate the effectiveness of our methodology.

We sincerely appreciate the reviewer's insightful feedback on the methodology of our paper. Please find our responses below:

Response to Weakness 1 and Weakness 2:

We greatly appreciate the reviewer’s feedback and would find it helpful if the reviewer could provide specific paper references for the research areas mentioned, as this would allow us to make more focused and meaningful comparisons.

We compare our method to static pruning techniques because, to the best of our knowledge, there are no existing dynamic pruning methods for the pre-filling stage. Our work is the first to present such a method. Among the methods we identified, static pruning techniques are the closest, as they also perform pruning during the pre-filling stage. However, our approach represents a novel direction by conducting pruning during inference rather than after calibration, thereby offering an additional advantage over static calibration-based methods.

Regarding dynamic depth pruning methods and token-level dynamic decoding methods, we would like to clarify the fundamental distinctions between our method and existing work in these two research areas. Our method primarily focuses on online dynamic pruning during the prefilling stage, wheras main focus of most papers in dynamic depth pruning and token-level dynamic decoding areas is decoding stage.

We acknowledge that static pruning method can reduces the LLM storage overhead while dynamic pruning cannot. We have added the sentence "Our implementation loads the full model for dynamic pruning, while other methods load only the pruned version." to the caption of Table 1. The revised sentences is highlighted in red in our edited PDF.

In this work, we focus on the results concerning performance and acceleration. Regarding performance, we demonstrate that our method consistently outperforms all static pruning baselines across various models and pruning ratios (Tables 2, 3, and 6 in the main text), and we also demonstrate that PP achieves the best Performance Runtime Ratio, which quantifies the ratio of performance degradation per unit of runtime reduction, in comparison with FLAP and Wanda-sp (Table 5 in the main text). In terms of acceleration, we show that the inference speeds of PP are comparable to those of other static structured pruning baselines (Table 5 in the main text). It is important to note that inference speed is highly hardware-dependent and implementation-dependent. We believe that the inference speed of PP could be enhanced through further implementation optimizations or other techniques; we intend to explore these possibilities in future work. For example, we have explored the feasibility of running the probe in parallel with the actual computation of earlier pruned blocks. We present these results in Table 1. Here, PP-Parallel represents the approach where, when the actual computation happens on earlier pruned blocks, we generate the probe from the residuals of these earlier pruned blocks and perform the probing. PP represents the default setting used in the main text of the paper. The results show that we can still obtain performance gains and achieve comparable results to PP. For example, at a 40% pruning ratio, PP-Parallel achieves a perplexity of 17.9 on WikiText2, which is close to that of PP and much lower than the 38.9 achieved by FLAP. Furthermore, PP-Parallel achieves 61.4% accuracy on ARC-e, which is close to that of PP and much higher than the 52.5% achieved by FLAP. However, we are just demonstrating the feasibility of further improving PP's inference speed here; the actual parallelism is hardware-dependent and implementation-dependent.

Table 1: Zero-shot performance of LLaMA-2-7B after pruning attention and MLP blocks without fine-tuning.

Response to Weakness 4:

We would like to clarify that PP consistently outperforms FLAP and Wanda-sp across nearly all tasks, as detailed in the experimental section of the main text. In Table 5, we demonstrate that PP achieves the best Performance Runtime Ratio, which quantifies the ratio of performance degradation per unit of runtime reduction, in comparison with FLAP and Wanda-sp. In terms of performance boosts, PP's Performance Runtime Ratio values are 2.56× (95.65 compared to 37.37) and 2.85× (106.48 compared to 37.37) more efficient than those of FLAP and Wanda-sp, respectively, indicating a significantly lower rate of performance degradation compared to FLAP and Wanda-sp. With regard to the final speedups, the inference speeds of PP are on par with other structured pruning baselines, such as FLAP and Wanda-sp. It is important to note that inference speed is highly hardware-dependent and implementation-dependent. We believe that the inference speed of PP could be enhanced through further implementation optimizations or other techniques; we intend to explore these possibilities in future work.

Furthermore, our method is orthogonal to the FLAP and Wanda-sp, and it can be applied with different metrics and varying pruning ratios across layers while still achieving performance gains, we conduct experiments on LLaMA-2-7B by applying our method to both the Wanda-sp and FLAP methods. The results are presented in the table below. Please note that w/ PP means we added probing and history-informed pruning with importance-scaled fusion on top of each method without modifying any part of the original framework. PP represents the default setting used in the main text of the paper. From the results, we can clearly see that PP generalizes well on both methods and shows significant improvements when integrated. For example, when integrated with the FLAP method under a 40% pruning ratio, PP reduces the perplexity on WikiText2 from 38.9 to 19.9 and increases the accuracy on ARC-e from 52.5% to 58.6%. Similarly, when integrated with the Wanda-sp method under a 40% pruning ratio, PP lowers the perplexity on WikiText2 from 43.8 to 19.2 and improves the accuracy on ARC-e from 54.4% to 57.9%.

Table 1: PP enhances the performance of baselines on LLaMA-2-7B.

We sincerely appreciate the reviewer's insightful feedback on the methodology of our paper. Please find our responses below:

Response to Weakness 1:

Following baseline methods in our paper, we define the "zero-shot" as a method can utilize a calibration dataset but cannot prompt the model with task-specific examples at inference time to guide its responses [1, 2]. Regarding the calibration dataset, we use the exact same subset of the C4 dataset (more specifically, c4-validation.00000-of-00008.json.gz) for all methods (including ours), in line with the baseline methods in our paper. Furthermore, the C4 dataset is not the final dataset on which we test performance. We discuss the C4 dataset in Section 5 under "Baseline" and provide detailed information in Appendix A, "Implementation Details." Furthermore, we do not prompt the model with any task-specific example during evaluation for any method (including ours).

We ensure a fair comparison between the baseline methods and our method by utilizing the same calibration dataset and maintaining the same experimental settings.

[1] Wei, Jason, et al. "Finetuned language models are zero-shot learners." arXiv preprint arXiv:2109.01652 (2021).

[2] Brown, Tom B. "Language models are few-shot learners." arXiv preprint arXiv:2005.14165 (2020).

Response to Weakness 2:

We would like to respectfully clarify that the probe sequence ratio is not zero, but rather 0.05. In the 'Effect of Probe Combination on Performance' section, we conduct experiments to investigate how different probe sizes affect PP’s performance by varying the probe batch size from 1 to 20 (specifically, 1, 5, and 20) and the probe sequence ratio from 0.05 to 1.0 (specifically, 0.05, 0.1, 0.3, 0.5, 0.8, and 1.0). We present the experimental results in Figure 3 of the main text, demonstrating that even a small probe can significantly improve model performance. For example, when applying a small probe with a batch size of 1 and a probe sequence ratio of 0.05 on LLaMA-2-7B, the perplexity drops from 29.8 to 21.7; for OPT-13B, it decreases from 35.4 to 27.7, compared to scenarios where neither probing nor history-informed pruning of PP is applied. Note that all experiments utilize the PPsp metric for a fair comparison.

Response to Weakness 3:

We thank the reviewer for the insightful comment. While our method can be applied to per-sentence dynamic pruning, we have chosen to conduct online dynamic structured pruning in a batch-wise manner due to constraints associated with modern GPUs and the inference pipeline. Specifically, all samples within a batch need to use the same pruned weight matrix. Performing per-sentence dynamic pruning would require generating a unique pruned weight matrix for each sentence (i.e., N pruned weight matrices for N sentences), which would incur significant latency and make it impractical for real-time inference. Therefore, we conduct online dynamic structured pruning in a batch-wise manner to achieve a balance between computational efficiency and model performance.

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC