Compact Language Models via Pruning and Knowledge Distillation

Large language models (LLMs) targeting different deployment scales and sizes are currently produced by training each variant from scratch; this is extremely compute-intensive. In this paper, we investigate if pruning an existing LLM and then re-training it with a fraction (<3%) of the original training data can be a suitable alternative to repeated, full retraining. To this end, we develop a set of practical and effective compression best practices for LLMs that combine depth, width, attention and MLP pruning with knowledge distillation-based retraining; we arrive at these best practices through a detailed empirical exploration of pruning strategies for each axis, methods to combine axes, distillation strategies, and search techniques for arriving at optimal compressed architectures. We use this guide to compress the Nemotron-4 family of LLMs by a factor of 2-4x, and compare their performance to similarly-sized models on a variety of language modeling tasks. Deriving 8B and 4B models from an already pretrained 15B model using our approach requires up to 40x fewer training tokens per model compared to training from scratch; this results in compute cost savings of 1.8x for training the full model family (15B, 8B, and 4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to training from scratch, perform comparably to other community models such as Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art compression techniques from the literature. We have open-sourced Minitron model weights on Huggingface, with corresponding supplementary material including example code available on GitHub.

Paper

References (50)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer W3pm7/10 · confidence 4/52024-07-07

Summary

This paper empirically explores compressing language models with pruning and knowledge distillation. It summarizes the best practices of pruning and distilling language models, which are supported by extensive experiments.

Strengths

1. This paper is well-written, and the best practices are easy to follow, which is useful in practical LMs. 2. The paper includes sufficient experiments to support the main results (best practices).

Weaknesses

1. The novelty of compressing LMs with pruning and knowledge distillation is limited because the method in this paper seems to be a simple combination of these two widely used techniques. Although the main contribution of this paper may be the best practices summarized from extensive experiments, it is better to highlight the difference between the final choice in this paper and the approaches in previous work like [1]. 2. Extra computational cost should be considered. It seems the pruning process and the online inference of knowledge distillation requires extra computation. Besides the trained tokens, it is better to additionally compare the FLOPs of model training in the main experiments, like in Figure 1 and Table 2. [1] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning. 2024. In ICLR.

Questions

N/A

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Suggestions: The number of significant figures should remain consistent in Table 2 and Table 3.

Reviewer enpQ7/10 · confidence 4/52024-07-12

Summary

In this paper the authors explore compression of LLMs via pruning and Knowledge Distillation. They try out a variety of approaches for pruning as well as the retraining step and provide a comprehensive analysis of best practices for getting compact LLMs from their larger counterparts. The authors explore pruning across width and depth using importance scores and then retraining on a relatively small dataset via Knowledge distillation to get a capable compact model. This paper provides super interesting insights into model pruning which can be used in practice.

Strengths

1. The paper is very well written and provide excellent explanations for all of their modeling choices. 2. I really enjoyed reading the key takeaways of the paper being structured as best practices. Each point in the best practices part of the paper provide some key insight into model compression and what works and what does not work. 3. I also enjoyed reading the comprehensive details that are written in the appendix. Overall this is a very well thought out research paper with proper explanation for each modeling decision that is taken.

Weaknesses

1. Small nitpick: The figures and tables need much much more detailed captions. I should have some idea about what the table wants to say just from the caption. 2. While results on benchmarks are appreciated, I always feel that these benchmarks do not reflect everything about LLMs. I would have liked some form of qualitative study on comparing the generations from the models. Some form of human evaluation on say 25-50 long form generation examples would be great given that there is not enough time during rebuttal for a larger scale human evaluation.

Questions

1. In table 2 is the nemotron-8B baseline also trained on the 94B token retraining dataset? Even if the retraining dataset is a subset of the pretraining dataset of nemotron models, I believe the apples to apples comparison would be to retrain all the models on the retraining dataset as well. I understand retraining all the baselines would be tough given the time constraints but it would be great if you could retrain only the nemotron-8B model and share the results. 2. How does this compare to taking the retraining dataset, generate logits from nemotron-15B model and then using vanilla KD to train the nemotron-8B model further? So basically I ask for 2 more baselines if possible. First is further vanilla training on nemotron-8B on the retraining dataset and second is further KD on the nemotron-8B model with nemotron-15B as the teacher. 3. Check weakness 2

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors have mentioned limitations.

Reviewer qFD87/10 · confidence 4/52024-07-13

Summary

This paper proposes Compact Language Models via Pruning and Knowledge Distillation, which combines various tricks and methods to compress a 14B model to 8B while achieving better performance than training from scratch.

Strengths

The paper conducts extensive experiments, comparing the latest baselines, and the authors summarize extensive tricks for pruning the model. The pruned 8B model performs well, surpassing the model trained from scratch. The paper is well-written, and the overall structure is good. Despite having many conclusions, it does not confuse the reader.

Weaknesses

Will the authors open-source the code? If the code and data are open-sourced, I would raise my score. When pruning a 70B model to 7B, would the method in the paper still work? Would this model perform better than pruning a 13B model to 7B?

Questions

see weakness

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

na

Reviewer enpQ2024-08-07

I appreciate the authors for adding the additional experiments and the qualitative study in such a short amount of time. According to me most weaknesses are addressed in the rebuttal. I am increasing my score to a 7.

Authorsrebuttal2024-08-08

We'd like to thank the reviewer for raising their score. Truly appreciate it! We are very happy to see that our rebuttal has addressed your questions and concerns.

Reviewer W3pm2024-08-09

Official Comment by Reviewer W3pm

I thank the authors for their response. After reading their reply, I will increase my score from 6 -> 7. I suggest the authors to include the additional experiment results in the final version of the paper.

Authorsrebuttal2024-08-09

Thank you raising the score. Really appreciate it! We will make sure to include the additional results in the final paper.

Area Chair 4nTD2024-08-12

Dear reviewers: as you are aware, the reviewer-author discussions phase ends on Aug 13. We request you to kindly make use of the remaining time to contribute productively to these discussions. If you have not read and/or responded to author rebuttal, please do it asap so that the authors get a chance to respond to you. If you have more questions to ask or want further clarification from the authors, please feel free to do it.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC