You Only Prune Once: Designing Calibration-Free Model Compression With Policy Learning

Authors: Ayan Sengupta, Siddhant Chaudhary, Tanmoy Chakraborty
Primary area: foundation or frontier models, including LLMs
Keywords: Model Compression, Large Language Models, Structured Pruning

Abstract:
The ever-increasing size of large language models (LLMs) presents significant challenges for deployment due to their heavy computational and memory requirements. Current model pruning techniques attempt to alleviate these issues by relying heavily on external calibration datasets to determine which parameters to prune or compress, thus limiting their flexibility and scalability across different compression ratios. Moreover, these methods often cause severe performance degradation, particularly in downstream tasks, when subjected to higher compression rates. In this paper, we propose *PruneNet*, a novel model compression method that addresses these limitations by reformulating model pruning as a policy learning process. PruneNet decouples the pruning process from the model architecture, eliminating the need for calibration datasets. It learns a stochastic pruning policy to assess parameter importance solely based on intrinsic model properties while preserving the spectral structure to minimize information loss. PruneNet can compress the LLaMA-2-7B model in just 15 minutes, achieving over 80\% retention of its zero-shot performance with a 30\% compression ratio, outperforming existing methods that retain only 75\% performance. Furthermore, on complex multitask language understanding tasks, PruneNet demonstrates its robustness by preserving up to 80\% performance of the original model, proving itself a superior alternative to conventional structured compression techniques.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730234514554

thread_type:
official_review

<thread_object>
creator: reviewer_onbh

<content>
summary:
This paper describes a technique -- PruneNet that avoids the usage of calibration data for structured pruning. They instead reformulate it as a policy learning process via (prun)ing-object(e)d (net)work. They does this by first making a corollary that the range of eigenvalues of the compressed model weights is a subset of the original model weights (corollary 3.3). Following that, the key idea is to identify (using an MLP/policy learner) and remove rows/columns in such a way that the distribution (and distance using Kolmogorov Smirnoff (KS) Distance) of the eigenvalues for original and compressed weights remain the same. 

The method seems to achieve robust results across various models/datasets and compression ratios without a calibration dataset.

strengths:
1. The removal of dependence on calibration dataset while pruning. As noted by authors in L115, calibration dataset has an impact on performance of models, so it has important effects for applications. 
2. The intrinsic model compression is an interesting observation. This helps practitioners understand the dependency of compression methods (whether method X is adding parameters or not) with respect to the compression ratios
3. The reusability of policy learner i.e train a learner (MLP) for one compression ratio and use it for another ratio as the authors have noted down is a huge strength (L412)
4. Very well written with clear presentation.

I believe the paper would be a great addition with some clarifications needed in design space and experiments.

weaknesses:
I don't have a list of "weakness" i.e substantial missing components. Please refer the questions section for potential ideas and clarifications needed.

questions:
***Model design***

1. The sampling is done on FFN1 and complementary columns are chosen in FFN2 to match the dimensions. This makes sense logically/intuitively but did the authors experiment with sampling FFN2 separately and/or any other design choices. Does the indices match the intuitive selection. Or does the authors have an explanation on why this choice is better compared to other intuitive selections? 
2. Did the authors perform experiments on usage of other penalty functions (KS)? That can be a good ablation study to have. 
3. L292: The assumption of discounted penalty works if the LLM has higher singular values at later layers. Has there been any (previous) analysis on existing LLMs whether that’s the scenario? A reference would be great if it exists. If not, having plots similar to Figure 1 for other LLMs at appendix would be great. 

***Experiments***

4. L361: Are the results for the one-sided test and additional details available? A reference would be great.
5. L375 (RFT): Can the authors add more insights on why RFT on model X might reduce the performance?
6. L375 (RFT): Can the RFT be done on SliceGPT as well (i.e Table 3 looking similar to Table 2). That will help to understand whether RFT helps SliceGPT more (>1.5%) and PruneNet less (marginal 1.5%) or less for both methods (in either case, it will be a useful finding to understand the effects of RFT)

***Suggestion for additional experiment on design choice***

1. Can this policy be extended to attention matrices as well? From my understanding, the key idea seems to be applicable on any matrix, so this experiment will be super helpful to understand the design space more clearly? There has been some studies [1] done to understand the effects of compression on different modules in a transformer, so the authors might take some inspiration to understand the effects on compressing on different modules.
2. Can the authors perform an experiment with >40% compression ratio? The performance drops most likely with higher compression ratios [1, 2, 3] but it would be great to know how much the PruneNet can achieve for reasonable compression. 

***Possible relevant citation***
1. The Cost of Compression: Investigating the Impact of Compression on Parametric Knowledge in Language Models - https://arxiv.org/abs/2312.00960
2. Are sixteen heads really better than one? - https://arxiv.org/abs/1905.10650
3. Compressing bert: Studying the effects of weight pruning on transfer learning - https://arxiv.org/abs/2002.08307

***Format***

1. Line 37 - Can be more descriptive about the model sizes rather than a simple statement (eg: Llama 2 can be of different sizes)
2. Can the authors give clear explanations on different between Effective Sparsity and Sparsity (maybe in Appendix if possible)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730475783579

thread_type:
official_review

<thread_object>
creator: reviewer_ooui

<content>
summary:
This paper introduces PruneNet, a structured pruning technique that deploys a policy learning process to identify important parameters. The proposed method is calibration-free, and once completed, can be applied to different pruning ratios without repeated assmentments of weights and data. The proposed method can be applied to popular Llama-2-7B in just 15 minutes, which is efficient enough for LLM pruning. The author conducted experiments on Llama-2 7B and Phi-2, with the compression ratio from 20% to 30%. Compared to SliceGPT, the method archives better accuracy on 5 zero-shot tasks, keeping >80% performance of the original LLMs. Besides, even without fine-tuning, the method is still able to preserve a good average accuracy.

strengths:
The idea of calibration-free pruning is interesting.  The policy learned on one compression ratio is scalable and can be directly transferred to higher or lower ratios. This is a useful feature if one would like to evaluate the performance of pruned models for a good trade-off between performance and efficiency. In addition, Extensive experiments were conducted to evaluate the effectiveness of the policy learning, with different finetuning settings and compression ratios, The proposed method is able to achieve better average accuracy (+3%~+4%) on zero-shot tasks.

weaknesses:
1. It appears that most of the experiments in this paper were conducted without fine-tuning. It would be insightful to examine how the proposed method compares to other baselines when fine-tuning is enabled, as this could provide a more complete picture of its performance under optimal conditions.
2. The efficiency of policy learning has not been studied. Providing more analysis on the efficiency of policy learning would enhance the understanding of its practical impact.
3. Based on point 3, if we allocate the same computational budget across all baselines, such as assigning the training cost of policy learning to LoRA fine-tuning in other baselines, does the proposed method still achieve superior performance?
4. Including pseudocode for the entire pipeline would improve clarity and reproducibility.
5. On of my main concerns is the lack of clarity behind the motivation for each design choice. For example, Equation 3 is introduced without sufficient context or explanation. The KS Distance introduced in Equation 6 also reveals weak insights about its motivation. Why KS distance instead of other distribution distance? In addition, other paragraphs could also benefit from refinement. In particular, the background section, especially sections 3.2 and 3.2, is densely packed with technical details and formulas about SliceGPT and eigenvalues. However, the relevance to the proposed method remains unclear until the final paragraph, creating a disconnect that might hinder reader understanding.

questions:
May I ask about the motivation behind Equation 3 (the importance score)? It appears that the equation is designed to project the parameter $W$ into a score $d$-dim vector. Could you clarify why this approach was chosen over other designs, such as directly optimizing a d-dimensional vector as the importance score? Additionally, will the learned matrix be shared across all linear weight layers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730582586308

thread_type:
official_review

<thread_object>
creator: reviewer_bk7b

<content>
summary:
This paper proposes a FFN-layer pruning algorithm for large language models that does not require a calibration dataset. The idea is to prune the weight matrices so that the spectrum of the matrices before and after pruning remains similar in the Kolmogorov-Smirnoff distance. To hasten the pruning procedure, one trains a policy that predicts the importance scores of the first layer of FFN. The empirical results show that the proposed method outperforms SliceGPT, in terms of prediction quality and the actual speedup.

strengths:
- The central research question of training a policy to prune LLMs is quite intriguing, with some practical importance. Although the paper mainly claims the benefit on the data side, I personally find that the method might have more future uses on low-memory or compute scenarios where it is very difficult to run usual pruning algorithms that require much gradient-based computations.

- The empirical performance seems to be reasonably good.

- The paper asks many interesting questions in section 5.2, including the transferability and the layerwise compression ratio.

weaknesses:
- One nitpick is that the comparison against SliceGPT, reported on Table 2, can be quite misleading. The figures seems to be the one for SliceGPT with WikiText2 calibration data, which is much worse than the SliceGPT with Alpaca calibration data. In fact, with Alpaca, I believe that the SliceGPT works much better than the proposed PruneNet. Of course, this is somewhat expected because it utilizes more data than PruneNet; making a comparison with both WikiText and Alpaca case does not diminish the usefulness of the proposed method, as they assume different amount of usable resources. Thus I recommend including both SliceGPTs, for a more comprehensive comparison.

- Another weakness is that the proposed method may not work better than the methods that use calibration data, as assuming no access to the calibration data also makes it impossible to use RFT. It seems like Table 3, ironically, can also be interpreted as a limitation of the proposed method, in a sense that the model pruned with PruneNet cannot recover much with RFT, while LLMs pruned with other methods can recover much with fine-tuning.

- Also, the motivation is quite unclear to me. At least for the text data, I am not sure why it is practical to assume no access to the calibration data. We already have quite abundant text data (crawled from web), even as our benchmark. Why should we suddenly assume that those data are not available? Can you give any further justifications?

- The idea of spectrum matching is also not justified well. First, I do not see how corollary 3.3 can be used to predict that the singular values of the matrix become more right-skewed. Doesn't the corollary simply states that the spectrum will be about subsampling, rather than predicting anything about the skewness? Also, the fact that there can be some spectrum shift does not fully mean that it will be an effective criterion for deciding how to prune. Perhaps more comparison with other criterion (e.g., minimizing Frobenius norm distortion) may help us understand whether the spectrum matching part is indeed an essential and nontrivial component.

questions:
Discussed in 'weaknesses'
</content>
</thread_object>

</thread>
<thread>
creation date:
1730622310338

thread_type:
official_review

<thread_object>
creator: reviewer_zgwt

<content>
summary:
This paper introduces PruneNet, a policy-based pruning method designed to learn the removal of redundant parameters without the need for calibration data. PruneNet’s policy is trained based on intrinsic model properties to minimize information loss, quantified using the KS distance between the singular value distributions of the uncompressed and compressed matrices. To make policy learning differentiable, the paper employs the reparameterization trick, introducing a random variable into the sampling process and training the policy module by evaluating various pruning outcomes. The proposed method is primarily evaluated on Llama-2 7B and Phi-2, with comparisons to SliceGPT, a technique that directly compresses parameter matrices by training on datasets. Results indicate that PruneNet achieves superior performance on several popular benchmarks.

strengths:
* One of the key strengths of this paper is a data-independent importance scorer, derived through a policy-learning process. The policy focuses mainly on individual linear layers, making the method theoretically efficient. This efficiency is demonstrated by its successful application to Llama-7B, where pruning requires only 15 minutes, illustrating its practicality for large-scale models.
* The authors conducted extensive experiments, exploring factors such as compression ratio, RFT dataset, and computational costs. 
* The policy learner can be applied across different compression ratios, which offers practical value for determining optimal pruning ratios in large language models (LLMs). This flexibility could be particularly advantageous for users looking to fine-tune pruning ratios without retraining the policy for each setting.

weaknesses:
* The pruning approach is focused on single layers, potentially overlooking structured interdependencies with subsequent layers. In structured pruning, parameters in one layer impact those in the next; by not accounting for this, the proposed method risks removing essential parameters from interconnected layers. This oversight could inadvertently affect model performance, as it might miss opportunities for more coordinated pruning across layers.
* Although the method is designed to optimize the KS distance before and after pruning, it lacks a detailed analysis of its policy-learning approach. A simple random sampling approach might also work: sampling multiple randomly pruned matrices and selecting the one with the lowest KS distance might be able to provide good results. Including more ablation studies to compare these methods would enrich the analysis. Additionally, further discussion on the design and benefits of the policy approach would add valuable insights.
* While the method shows promise on Llama-7B, its efficacy on larger models, such as Llama 13B, remains untested. An evaluation on larger LLMs would help establish its robustness and scalability, further highlighting its applicability to a broader range of large models.

questions:
Please see weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1732116268603

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Confirmation of Rebuttal Submission

comment:
Dear Reviewers,

We appreciate your effort to review our paper meticulously and provide us with your valuable suggestions. We have addressed all your comments and provided the additional results you suggested. We have also updated the paper (changes are highlighted in blue color) to accommodate these changes. We request you to kindly go through our responses and let us resolve any other query that you may have.

Thanks,
</content>
</thread_object>

</thread>
<thread>
creation date:
1733222732611

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comments by Authors

comment:
We sincerely thank all the reviewers for their valuable feedback during the discussion period and their comprehensive review of our work. We appreciate their recognition of the contributions of our work,

## The first calibration-free structured pruning technique

Our proposed method, PruneNet, is the *first* structured model pruning technique that does not need any calibration dataset for compression, wherein a data-independent importance scorer is used to prune the layers of an LLM. **All reviewers highlighted the importance of such calibration-free pruning methods**, particularly in data-private and low-resource settings, and a thorough discussion of the motivation behind such methods (**particularly with reviewer bk7b**) reaffirms the key contribution of our work. Completely removing calibration datasets also makes PruneNet a highly reusable and efficient pruning technique, as it bypasses the need to run expensive forward/backward passes of LLMs through the calibration data.

## Efficiency and transferability of the proposed method

With PruneNet, we can significantly reduce the model compression time by 50%. This efficiency, as highlighted in the paper, is a key advantage of PruneNet. It can compress an LLM at different ratios at once without retraining the policy learner, a feature that is particularly useful in practical applications. This adaptability to different edge device resources **has been acknowledged by both reviewers bk7b and ooUi**, further enhancing the method's practicality.

## Extensive comparison with baselines on a wide range of models and compression ratios

We have presented an extensive comparison of PruneNet with other competitive model pruning techniques, notably SliceGPT, in a variety of set-ups. This comparison, which includes scenarios with/without recovery fine-tuning (RFT) and using different calibration datasets (WikiText2, Alpaca), provides a comprehensive view of PruneNet's performance. We also compared the performance for a variety of model architectures, with the total number of parameters ranging from 125M to 13B. Such an extensive comparison showcases the robustness of our proposed method, even when baselines are calibrated with instruction-based calibration datasets such as Alpaca. 

**All the reviewers (ZGwT, bk7b, ooUI, onBH) have acknowledged the results and thoroughness of our experiments**, providing positive feedback that reaffirms the quality and significance of our work.

## Robustness of PruneNet 

We've also presented a detailed study of various components of the design choices in PruneNet. This includes the choice
of reward functions used in the policy learning process, the usage of FFN1 vs FFN2 layers for sampling, a comparison of a policy-based
pruner vs a static pruning method such as random sampling of indices and a comparison of stochastic and deterministic policies to be used for pruning. The results of these experiments motivate the design choices we made and highlight the robustness of PruneNet in various setups.

## Clarity of presentation

**Reviewer onBH commended the clarity of presentation in our paper.**

We've also thoroughly responded to every follow-up question for each reviewer. Here is a summary of all the responses:

* In response to reviewer ZGwT's follow-up questions, we presented an ablation study comparing PruneNet's performance against a simple random sampling approach (with PruneNet performing better). We highlighted the importance of sampling in our overall design (against simply choosing the top-k most important rows of a matrix).

* For reviewer bk7b, we presented a comparison of using different calibration datasets (WikiText2, Alpaca) with SliceGPT against PruneNet. While SliceGPT's performance is marginally better than PruneNet when it uses the Alpaca dataset for the LLaMa-2-7B model, the trend was the opposite for the Phi-2 model, exposing the inherent limitation of calibration-based pruning techniques (namely their sensitivity to the choice of calibration data).

* For reviewer ooUI, we presented the results of PruneNet with recovery fine-tuning. We conducted additional ablation experiments to highlight the importance of each design choice in PruneNet.

* For reviewer onBH, we presented a few more ablation results (FFN1 vs FFN2 layers, RFT), results of an implementation of PruneNet for matrices in the attention layer,  and the results of PruneNet for compression ratios larger than 40%.

Based on our responses, we observed an overall positive change in the scores:

* Reviewer onBH improved the soundness score of our paper from 3 to 4.
* Reviewers ooUI and bk7b improved the overall rating of our paper from 5 to 6.
* While reviewer ZGwT acknowledged that all of their concerns were addressed in our rebuttal, they maintained their original rating of 6.

We hope that our work motivates further research in developing more efficient model compression techniques in low-resource settings and makes a good contribution to the field.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces PruneNet, a novel structured pruning method for LLMs that eliminates the need for calibration datasets. By learning a stochastic policy to preserve the spectral structure of model weights, PruneNet achieves superior efficiency and robustness compared to methods like SliceGPT, retaining high performance even at significant compression ratios. The ability to transfer policies across compression ratios and compress models in minutes highlights its practical value for low-resource and privacy-sensitive settings.

The paper is well-written, with strong empirical validation and thorough comparisons, addressing reviewer concerns effectively. While minor limitations, such as limited tests on larger models, were noted, the authors’ detailed rebuttals and additional results strengthened the work. Overall, this paper provides a significant contribution to model compression, offering an efficient and scalable solution. I recommend its acceptance.