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

An implementation of policy-based pruning for attention layers. We carry out a similar pruning strategy on attention layers, wherein we learn importance scores of the rows for the output matrix of the attention layer, and correspondingly slice off the key, query and value matrices to keep the output dimension consistent. We highlight the pruning results for the LLaMA-2-7B model with compressed self-attention layers in Table 5 (Table 20 of the updated paper). The performance drop with the compressed models suggests that compressing self-attention layers intrinsically is harder than compressing FFN layers. However, around $50$% of the performance drop can be recovered with recovery fine-tuning.

Table 5. Results with Llama-2-7B with pruned self-attention modules.

Performance drops for PruneNet with high compression ratios ($>40$%). Table 6 (Table 13 of the updated paper) highlights the performance of the LLaMA-2-7B model at $50$% compression ratio with PruneNet and SliceGPT. While both methods can regain only $60$% of the performance of the original uncompressed model, PruneNet demonstrates $2$% better than the SliceGPT, showcasing its effectiveness over the baseline, even at a very high compression rate.

Table 6. Results with 50% compression ratio on Llama-2-7B

Format: Sparsity and Effective Sparsity. We thank the reviewer for the suggestion. We have added the explanation in Appendix A.2 of the updated paper.

References

[1] Yuan, Zhihang, Yuzhang Shang, Yue Song, Qiang Wu, Yan Yan, and Guangyu Sun. "Asvd: Activation-aware singular value decomposition for compressing large language models." arXiv preprint arXiv:2312.05821 (2023).

[2] Li, Yuanzhi, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. "Textbooks are all you need ii: phi-1.5 technical report." arXiv preprint arXiv:2309.05463 (2023).

We thank the reviewer for the valuable feedback. We address the concerns raised. All the changes in the main manuscript are highlighted with blue color.

A comparison of pruning on FFN1 vs FFN2 layers. Table 1 (Table 19 of the updated paper) highlights the zero-shot performance of the LLaMA-2-7B model compressed with PruneNet with policy learned on different FFN matrices. In most cases, we observe marginal differences in the result ($< 1$%) when the policy is learned with FFN2 instead of FFN1. The observations emphasize that PruneNet is invariant to the choice of parameter used for learning the policy.

Table 1. A comparison of FFN1 vs FFN2 layers within the PruneNet framework for Llama-2-7B.

Usage of reward functions other than the KS distance. It is important to notice that compression reduces the cardinality of the spectrum. Therefore, the traditional pointwise distance measures (like Euclidean/Cosine/Frobenius) are not applicable for comparing the distance between the spectrum structure of uncompressed and compressed models. Therefore, we resort to probabilistic distance measures that can capture the shift in distribution structures of the compressed model. To further understand the effectiveness of PruneNet under different distance measures, we evaluate the LLaMA-2-7B model compressed using PruneNet with non-parametric Anderson–Darling measure of agreement. Table 2 (Table 18 of the updated paper) highlights the effectiveness of PruneNet with both Kolmogorov-Smirnov (highlighted as KS) and Anderson–Darling (highlighted as AD) distance measures. Under both reward functions, we observe a similar performance of PruneNet for different compression ratios with the LLaMA-2-7B model. The results further emphasize the stability of our proposed compression method under different choice metrics.

Table 2. Zero-shot performance of Llama-2-7B compressed with PruneNet with different reward functions.

Motivation behind the usage of discounted rewards and singular values of deeper layers. Weight matrices of LLMs have been known to exhibit a variation in the distribution of singular values across layers [1]. To further verify our observations from Figure 1 of the main text, we conduct a similar study on the FFN1 matrices of the OPT and Phi-2 models. We plot the cumulative distribution of the singular values of these matrices and observe consistent behaviour across all three model architectures. Figures 5a and 5b in the updated paper highlight the spectrum of the FFN1 module at different layers for Phi-2 and OPT-2.7B models, respectively. This further strengthens our assumption of using discounted rewards in our policy learning approach.

Detailed results of the one-sided KS test. For LLaMA-2-7B and Phi-2 models, we calculate the recovered performance ($\frac{\text{Compressed model performance}}{\text{Dense model performance}}$) for different compression ratios with PruneNet and SliceGPT. We conducted a one-sided KS test with the null hypothesis as SliceGPT recovered performance higher than PruneNet recovered performance for different compression ratios. For a compression ratio of $20$%, we obtained the test statistics value of 0.5 and $p$-value of $0.043$, suggesting that we reject the null hypothesis. Similarly, for compression ratio $25$% and $30$% we obtain statistics $0.5$ ($p$-value $0.043$) and $0.6$ ($p$-value $0.026$), respectively. With these statistical results, we conclude that the performance recovered with PruneNet is higher than SliceGPT with statistical significance.

Why does RFT reduce performance of models like Phi2? Although interpreting the performance of a model compression method requires much detailed analysis, which is much out of scope for our current study, we can make certain educated guesses to understand the behaviours of different LLMs under compression. One possible reason behind the inferior performance of Phi-2 under RFT could be attributed to its pre-training objective. Phi-2 is pre-trained on a vast amount of synthetic texts [2], which allows it to perform reasonably well on common sense and logical reasoning tasks (the tasks used in our work) even without fine-tuning. Therefore, fine-tuning on additional datasets like WikiText2 or Alpaca could hurt the model's performance. Table 2 of the paper shows that the post-compression Phi-2 model demonstrates higher performance recovery than any other LLM, indicating its robustness on complex reasoning tasks, even after compression. Unlike SliceGPT, PruneNet preserves the original knowledge base of LLMs by preserving spectral structures. Therefore, additional fine-tuning can distort the reasoning abilities of the compressed LLM, adversely affecting its zero-shot performance.

A comparison of PruneNet with SliceGPT with RFT. We report the results with LLaMA-2-7B with RFT on Wikitext2 and Alpaca datasets in Table 3 (Table 11 of the updated paper) and Table 4 (Table 12 of the updated paper), respectively. With RFT on the Wikitext2 dataset, PruneNet achieves, on average, $>5$% accuracy than SliceGPT. However, SliceGPT outperforms PruneNet when fine-tuned on the Alpaca dataset, with an average margin of $3$%. As SliceGPT uses the same datasets for calibration and RFT, it typically has access to more instruction-tuning datasets than PruneNet, allowing it to do better when fine-tuned on the Alpaca dataset. However, it is worth noticing that the average standard deviation in the performance of SliceGPT after RFT is significantly high ($5.5$) compared to PruneNet ($1.5$). The low standard deviation highlights the robustness of PruneNet when fine-tuned on different datasets to recover information loss during compression.

Table 3. Results with RFT on WikiText2 on Llama-2-7B

Table 4. Results with RFT on Alpaca on Llama-2-7B

Dear reviewer onBH,

We thank you for reassessing our paper and increasing the soundness score. We would greatly appreciate your guidance on any additional concerns you may have about improving the overall rating of our paper.

We are eager to hear any additional recommendations for increasing the impact of our work.

Thanks,

We thank the reviewer for the valuable feedback. We address the concerns raised. All the changes in the main manuscript are highlighted with blue color.

Comparison of PruneNet with other compression methods with RFT enabled. We report the results with LLaMA-2-7B with RFT on Wikitext2 and Alpaca datasets in Table 1 (Table 11 of the updated paper) and Table 2 (Table 12 of the updated paper), respectively. With RFT on the Wikitext2 dataset, PruneNet achieves, on average, $>5$% accuracy than SliceGPT. However, SliceGPT outperforms PruneNet when fine-tuned on the Alpaca dataset, with an average margin of $3$%. As SliceGPT uses the same datasets for calibration and RFT, it typically has access to more instruction-tuning datasets than PruneNet, allowing it to do better when fine-tuned on the Alpaca dataset. However, it is worth noticing that the average standard deviation in the performance of SliceGPT after RFT is significantly high ($5.5$) compared to PruneNet ($1.5$). The low standard deviation highlights the robustness of PruneNet when fine-tuned on different datasets to recover information loss during compression.

Table 1. Results with RFT on WikiText2 on Llama-2-7B

Table 2. Results with RFT on Alpaca on Llama-2-7B

Efficiency of policy learning. The primary motivation behind policy learning is to decouple the model pruning process from the model itself. Therefore, with an appropriate policy learner model, we can learn which parameters can be compressed within an arbitrary model component just by looking at the model parameters without posing any additional constraint on the component architecture. This allows us to flexibly use the same policy learner model for compressing different layers and components of an LLM.

We highlight the results with random policy in Table 3 (Table 16 of the updated paper), where we select the pruned indices randomly for each model parameter. We refer to this as a 'random' selection. We observe an average $2$% drop with a random selection method, demonstrating the need for a learnable policy for effective model compression.

Table 3. Effect of learnable policy on compressed Llama-2-7B model.

To further understand the effectiveness of the learned policy empirically, we perform experiments where we learn deterministic policy based on the parameter importance calculated in Equation 3 of the paper. In the deterministic policy, we chose only the topk parameters based on this computed importance metric, making the selection deterministic. Table 4 (Table 17 of the updated paper) highlights the results of LLaMA-2-7B with deterministic and stochastic policies (policy learned with PruneNet defined in Equation 5 of the paper). We observe that deterministic policy often underperforms the stochastic policy with an average margin of $2$%. The results highlight that parameter importance alone cannot determine which parameters to compress. Preserving the spectral structure between the compressed and uncompressed models is critical to ensure minimal performance drop post-compression.

Table 4. Effect of stochastic policy on compressed Llama-2-7B model.

Computational budget. The policy learner model is typically $1000$ times smaller than the LLM to be compressed (number of learnable parameters $<7$M). As Section 1 of the paper highlights, most existing structured pruning methods learn the compressed model after being calibrated on some dataset. Therefore, the computational budget utilized by the policy learner model is far less than the existing baselines, making our compression method more scalable and efficient.

Pseudo code. We thank the reviewer for the suggestion. We have added the pseudo-code of the proposed compression method in Algorithm 1 of the updated paper.

A detailed motivation behind design choices. We perform an ablation study to understand the importance of different components of PruneNet.

We conduct experiments with a random selection process, where the pruned parameter indices are chosen randomly. We highlight the results with random policy in the above Table 3 (Table 16 of the updated paper). We observe an average $2$% drop with a random selection method, justifying the need to learn which parameters to compress for more effective model compression.

Above Table 4 (Table 17 of the updated paper) highlights the results with LLaMA-2-7B with deterministic and stochastic (policy learned with PruneNet in Equation 5 of the paper) policies. In the deterministic policy, we chose only the topk parameters based on the importance metric defined in Equation 3. We observe that deterministic policy often underperforms the stochastic policy with an average margin of $4$%. The results highlight that parameter importance alone cannot determine which parameters to compress. Preserving the spectral structure between the compressed and uncompressed models is critical to ensure minimal performance drop post-compression.

To further understand the effectiveness of PruneNet under different distance measures, we evaluate the LLaMA-2-7B model compressed using PruneNet with non-parametric Anderson–Darling measure of agreement. Table 5 (Table 18 of the updated paper) highlights the effectiveness of PruneNet with both Kolmogorov-Smirnov (highlighted as KS) and Anderson–Darling (highlighted as AD) distance measures. Under both reward functions, we observe a similar performance of PruneNet for different compression ratios with the LLaMA-2-7B model. The results further emphasize the stability of our proposed compression method under different choice metrics.

Table 5. Zero-shot performance of Llama-2-7B compressed with PruneNet with different reward functions.

Table 6 (Table 19 of the updated paper) highlights the zero-shot performance of the LLaMA-2-7B model compressed with PruneNet with policy learned on different FFN matrices. In most cases, we observe marginal differences in the result ($< 1$%) when the policy is learned with FFN2 instead of FFN1. The observations emphasize that PruneNet is invariant to the choice of parameter used for learning the policy.

Table 6. A comparison of FFN1 vs FFN2 layers within the PruneNet framework for Llama-2-7B.

Response to Question. The key motivation behind the design of PruneNet is to have a predictor network which can compute the importance scores of the rows of any weight matrix; this allows us to reuse the same learned predictor network to prune the FFN weight matrix of any layer in the model, which turns out to be extremely efficient compared to other SOTA pruning methods.

Some clarification about equation (3) of main text is in order. Note that the output of equation (3) is an $n$-dimensional vector (and not a $d$-dimensional vector), where $n$ is the number of rows in the weight matrix $W$. In essence, the predictor network is computing importance scores for each row of a weight matrix in two steps: in the first part of equation (3), the network computes the relative importance of rows amongst themselves. This is needed since the importance of a row is potentially correlated with the importance of other rows of the weight matrix. The second part of equation (3), particularly the matrix multiplication within $\sigma$, simply projects the $n\times n$ matrix of relative importance scores into an output vector of importance scores of each row. In contrast, as suggested by the reviewer, directly optimizing a vector of importance scores conflicts with our design philosophy of reusability of the learned predictor network.

Regarding the use of a distributional distance: since pruning reduces the cardinality of the spectrum, the use of traditional pointwise distance measures (like Euclidean/Cosine/Frobenius norm) are not applicable for comparing the spectrum structure of uncompressed and compressed models. Due to this, we resort to distributional distance measures to capture the shift in the spectrum of model; a specific structural argument is given in Corollary 3.3. As mentioned above, we also evaluated the effectiveness of PruneNet under two different distance measures, namely the KS distance and the non-parametric Anderson-Darling measure of agreement.

We also explained above that directly optimizing the parameter importance often leads to poorer performance than the learned compression policy. We reuse the same policy learner for all the layers within the LLM. The policy learner is agnostic to the model architecture and learns a mapping between the compressed and uncompressed model parameters. The simplicity of the policy learner allows us to reuse the same model to compress an LLM at any given compression ratio, offering greater flexibility and efficiency.

Dear reviewer ooUi,

We thank you for your effort in reassessing our work and increasing the overall rating. We are committed to addressing all your concerns into the paper.

Thanks,

We thank the reviewer for the valuable feedback. We address the concerns raised. All the changes in the main manuscript are highlighted with blue color.

Structural interdependencies with subsequent layers. We argue that the interdependencies between subsequent layers are already captured during the pre-training phase of the LLM. Therefore, the spectral structure of the feedforward layers already preserves the information propagation between the subsequent layers. By preserving the spectral structure after compression, PruneNet preserves the information that needs to be maintained at later layers, therefore omitting additional effort to preserve the information flow between interconnected layers.

A comparison with simple random sampling. We highlight the results with random policy in Table 1 (Table 16 of the updated paper), where we select the pruned indices randomly for each model parameter. We refer to this as a 'random' selection. We observe an average $2$% drop with a random selection method, demonstrating the need for a learnable policy for effective model compression.

Table 1. Effect of learnable policy on compressed Llama-2-7B model.

Moreover, Table 2 (Table 17 of the updated paper) highlights the results with LLaMA-2-7B with deterministic and stochastic (policy learned with PruneNet) policies. In the deterministic policy, we chose only the topk (k depends on the compression ratio) parameters based on the importance metric defined in Equation 3. We observe that deterministic policy often underperforms the stochastic policy with an average margin of $4$%. The results highlight that learnable parameters alone are insufficient to determine which ones to compress. Preserving the spectral structure between the compressed and uncompressed models is also critical to ensure minimal performance drop post-compression.\newline

Table 2. Effect of stochastic policy on compressed Llama-2-7B model.

Efficacy of PruneNet on larger models. Table 3 (Table 14 in the updated paper) highlights the results with PruneNet and SliceGPT for different compression ratios for the LLaMA-2-13B model. The performance drops drastically for larger models like LLaMA-13B at a high compression ratio. However, the results in Table 4 (Table 15 in the updated paper) highlight that after recovery fine-tuning, the compressed models regain the performance quickly and can preserve up to $84$% of the original uncompressed model performance, even at a high compression rate of $30$%. PruneNet also outperforms SliceGPT with a margin of $2$%, showcasing a similar trend as the smaller LLMs used in our study.

Table 3. Llama-2-13B results without RFT.

Table 4. Llama-2-13B results with RFT on WikiText2.

Dear Reviewer ZGwT,

We have been trying to reach out to you with a request to check our responses to your comments. We have meticulously addressed all your concerns with new sets of experiments. We are sure our responses will address your comments. Since the discussion period will end tomorrow, it is our sincere request to you to check our responses at least once and reassess our paper.

Sincerely look forward to hearing from you.

Thanks

We thank the reviewer for the valuable feedback. We address the concerns raised. All the changes in the main manuscript are highlighted with blue color.

Comparison with SliceGPT with Alpaca calibration data. Table 1 (Table 10 of the updated paper) contains a comparison of the performance of the LLaMA-2-7B and Phi-2 pruned with SliceGPT using the Alpaca calibration dataset (as opposed to the WikiText2 dataset used in the main text of our paper). While LLaMA-2-7B pruned with SliceGPT exhibits better average performance for all compression ratios, Phi-2 exhibits an opposite trend, wherein PruneNet beats SliceGPT by a consistent margin of at least $2$% for all compression ratios. This trend reinforces a fundamental limitation of calibration data-based pruning techniques, namely their sensitivity to the choice and quality of the calibration dataset and the choice of model architectures. In conjunction with this, it should also be noted that the Alpaca dataset is an instruction-tuning dataset (as opposed to WikiText2). Therefore, it is not unreasonable for the performance of pruned models calibrated with the Alpaca dataset on generative tasks to be better than that of a calibration-free technique like PruneNet.

Table 1. Results without RFT, where SliceGPT uses Alpaca as calibration data

Motivation for calibration-free model compression techniques. While it is true that there is an abundance of calibration data for textual modalities, it is not the availability of calibration data alone that determines the quality of a pruned model. A detailed study by William et al., 2024 reveals a large degree of sensitivity in the performance on downstream tasks of pruned models based on the selected calibration data and model architectures, which questions the usability of calibration-based pruning methods in data-private settings; this has made the problem of sampling high-quality calibration datasets an active area of model pruning research. Another major limitation of calibration-based pruning techniques is that of reusability, wherein pruning multiple models necessitates running the forward/backward passes of pruned models on the calibration datasets, making the overall process highly inefficient. In contrast, PruneNet aims to be the first pruning technique to alleviate these difficulties and can also run right out of the box in low-resource/compute settings. Finally, PruneNet can be used in data modalities other than text, where sampling calibration data might be costly.

Justification of spectrum matching and comparison with other reward criterion. It is important to understand that pruning reduces the cardinality of the spectrum. Therefore, the traditional pointwise distance measures (like Euclidean/Cosine/Frobenius norm) are not applicable for comparing the spectrum structure of uncompressed and compressed models. Therefore, we resort to probabilistic distance measures that can capture the shift in the empirical distribution of spectrums of the compressed model. Corollary 3.3 suggests that compression reduces the spectrum range, and Figure 1 highlights that with more compression, the spectrum becomes more right-skewed. To further understand the effectiveness of PruneNet under different distance measures, we evaluate the LLaMA-2-7B model compressed using PruneNet with non-parametric Anderson–Darling measure of agreement.

Table 2 (Table 18 of the updated paper) highlights the effectiveness of PruneNet with both Kolmogorov-Smirnov (highlighted as KS) and Anderson–Darling (highlighted as AD) distance measures. Under both reward functions, we observe a similar performance of PruneNet for different compression ratios with the LLaMA-2-7B model. The results further emphasize the stability of our proposed compression method under different choice metrics.

Table 2. Zero-shot performance of Llama-2-7B compressed with PruneNet with different reward functions.

References

[1] Williams, Miles, and Nikolaos Aletras. "On the impact of calibration data in post-training quantization and pruning." In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10100-10118. 2024.

On spectrum matching. In Figure 1 of our paper, we highlighted how slicing impacts the spectrum structure of an LLM. In Corollary 3.3, we formalized this fact using the Poincare separation theorem. To understand the importance of using a learned policy instead of a static method such as random slicing of a weight matrix, we further carried out the following experiment: we randomly pick indices to slice off for each layer and report the standard deviation and mean of the obtained rewards. In Table 2, we observe a high standard deviation of rewards for smaller compression ratios, thereby requiring a mechanism that learns the indices to be pruned off.

Therefore, Corollary 3.3 and the above experiment suggest developing a learnable strategy to minimize the difference between uncompressed and compressed model spectrums for retaining information in the post-compressed LLM.

References

[1] Ashkboos, Saleh, Maximilian L. Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. "Slicegpt: Compress large language models by deleting rows and columns." arXiv preprint arXiv:2401.15024 (2024).

[2] Lin, Bin, Yang Ye, Bin Zhu, Jiaxi Cui, Munan Ning, Peng Jin, and Li Yuan. "Video-llava: Learning united visual representation by alignment before projection." arXiv preprint arXiv:2311.10122 (2023).

Dear reviewer bk7b,

We have addressed your recent concerns with new sets of experiments. Since the discussion period is ending soon, we request you to have a look at our recent responses and let us know if you have further concerns. Your feedback is valuable for us to address the potential weaknesses and improving the quality of our work.

If you feel our responses address your questions, please consider reassessing our paper.

Thanks

Dear reviewer bk7b,

We have addressed your recent concerns with new sets of experiments. We sincerely request you check our responses. If you feel our responses address your questions, please consider reassessing our paper. We are ready to address other concerns if any.

Thanks