LLM Pruning and Distillation in Practice

We would like to thank the reviewer for their encouraging comments and insightful feedback. Please find our responses below:

W1: ICLR 2025 style guidelines

We apologize for the confusion - we have followed the ICLR 2025 style guidelines to the best of our knowledge. We will re-verify the paper contents to ensure it meets all guidelines.

W2: I believe the writing could be improved, and some details in the paper also need enhancement, including ...

Thank you for pointing this out. We will incorporate the suggested changes in the final version.

W3: I recommend that the authors provide more analysis or experiments demonstrating the effectiveness of teacher correction across different models or datasets.

This is a relevant question. We validate the effectiveness of teacher correction by evaluating it on two distinct model families (Llama 3.1 and Mistral-NeMo, trained on different pretraining datasets). Figure 3 in the paper shows the effects of applying teacher correction on Mistral-NeMo 12B. In our internal experiments, we noticed a similar trend for the Llama 3.1 family as well.

W4: Comparison to ShearedLLaMa.

Thank you for pointing this out. To help with this comparison, we have created the following table to highlight the differences between our approach and Sheared LLaMa. We will add the appropriate citations to ShearedLLaMa in the final version.

Q1: Comparisons to SLEB, Shortened LLaMA, and ShortGPT

The table below compares our approach to SLEB, Shortened LLaMa, SliceGPT and ShortGPT:

Q2: I recommend that the authors provide quantitative evidence of the token distribution differences between datasets and demonstrate how these differences correlate with the effectiveness of teacher correction.

This is a great suggestion. Unfortunately, we have no access to the original datasets (Llama or Mistral-NeMo) to perform this comparative study across dataset distributions.

We would like to thank the reviewer for their encouraging comments and insightful feedback. Please find our responses below:

W1: Comparison to MobileLLM-350M

We thank the reviewer for pointing out the drop in accuracy. Though small models like MobileLLM-350M may shine on particular tasks, they still typically underperform when evaluated across multiple downstream tasks. In particular, MobileLLM-350M beats LLaMA-v2 7B on API function calling (a specific downstream task) but not other downstream tasks. On this particular task, we observe a similar trend where our Llama3.1-compressed-4B outperforms the original 8B model (~2x larger model, BFCL v2 with 59.6% vs 44.3%). From Tables 1 and 2, we observe that despite an aggressive 50% compression ratio, our compressed models retain competitive scores across the full suite of downstream tasks w.r.t models trained from scratch with an order of magnitude more training tokens.

W2: The authors should have compared their approach against a broader spectrum of existing compression techniques, including various pruning approaches, quantization methods.

Table 4 in the Minitron work [1] compares Minitron models to previous state-of-the-art compression techniques. Our compressed 8B and 4B models outperform all these baselines (including Minitron) as shown in Tables 1 and 2. We will add a detailed comparison to related work, including LLMPruner, SliceGPT, ShortGPT, etc. in the final version.

W3: Sub-1B parameter models

This is a great suggestion. We plan to apply our strategy to further compress existing models into the sub-1B parameter range in the near future.

Q1: Why Winogrande was chosen over other potential downstream tasks?

Winogrande is a good representative of commonsense reasoning tasks, especially for smaller models. We draw this conclusion based on our internal evaluations that indicate Winogrande doesn’t saturate at these model sizes and continues to scale. It also has the advantage that it can be computed quickly.

Q2: How sensitive is the method to the choice of downstream task for depth pruning?

We observed that different types of tasks utilize different heads of the model, though the general pattern is similar. Also we noticed that generation tasks like GSM8K are more sensitive to depth pruning than common knowledge tasks. The selection of candidates for further distillation may depend on the downstream task; we leave this topic for further research that targets domain-dependent pruning approaches.

References

1. Compact Language Models via Pruning and Knowledge Distillation, arXiv:2407.14679v2, 2024

We would like to thank the reviewer for their encouraging comments and insightful feedback. Please find our responses below:

W1: Domain adaptation vs. teacher correction

We agree with the reviewer that teacher correction is similar in spirit to domain adaptation. We wanted to highlight that in this case, finetuning on the distillation dataset is necessary even though we aren’t explicitly changing domains.

W2: dropping last contiguous layers (except the last layer) [3] vs. our approach

The main difference in our work is that we use drop in Winogrande accuracy as the metric to choose which contiguous layers have to be dropped. This results in a more accurate pruned model compared to using LM loss as the metric (as in [3]). The Table below compares the two approaches.

W3: I think the paper would benefit a lot from adding comparisons to similar small sized models, trained from scratch.

Table 1 in the paper compares our models to various other similarly-sized community models trained from scratch (eg: Llama 3.1, Gemma/Gemma2, Mistral, etc.).

W4: Could the authors provide an algorithm describing the final/best performing choices for pruning+distillation?

Our final approach for obtaining pruned+distilled models can be summarized as follows:

• If no access to the original dataset, perform teacher correction using ~100B tokens from the distillation dataset.
• Compute importance rankings for depth and width axes on corrected teacher. For width, use L2-norm and mean aggregation across the batch and sequence dimensions, respectively (as evaluated in [1]); for depth, find contiguous layers that least affect Winogrande accuracy.
• Use the computed rankings to prune the corrected teacher along width and/or depth axes.
• Distill knowledge from corrected teacher to pruned student model using logit only loss and forward KLDiv loss.

W5: Could you elaborate on the following points which are ambiguous/unclear ...

Lines 167-168 Based on the original Minitron paper, we focused on retaining network depth (except for Llama3.1-Depth) and attention heads. The remaining axes were pruned in proportion to the parameter budget with an emphasis on retaining hidden dimension size.

Lines 133-134 The 1024 calibration data samples are drawn randomly from the full distillation data corpus. We did not experiment with calibration data drawn from different domains, but we hypothesize that this will change the importance rankings proportional to how different the two data domains are.

Lines 312-313 Based on our observations, for depth based pruning, we prune a set of contiguous layers that results in the least drop in Winogrande accuracy.

Continuously corrected teacher

• In this setting, the (1) teacher correction and (2) distillation from the corrected teacher, are run in parallel.
• The teacher model/checkpoint used in (2) is getting updated periodically through (1).
• Continuous teacher correction is described in lines 268-273, with results presented in Figure 4.

W6: Full fine-tuning vs. PEFT

Full fine-tuning uses the same resources as model pretraining and has been shown to scale to models well over the sizes we consider in this work. We agree that parameter-efficient fine-tuning would be a more resource-efficient way to perform teacher correction. We leave this exploration to future work. We also note that correction is only necessary when access to original training data is restricted.

We agree that using PEFT for teacher correction (and potentially even distillation) would be an interesting avenue for future research.

W7: how do the observations translate to public datasets eg: openwebtext...

We have not evaluated our method on public datasets; however, we expect the general trends to remain the same. This would be another interesting avenue for future work.

W8: Code release

We plan to release the code and experimental pipeline in the coming weeks, pending internal legal review. Model weights will also be released on HF with a permissive license. If required, we can also upload our checkpoints immediately to HF under an anonymous ID.

Responses to remaining questions provided in the follow-up comment.