Compact Language Models via Pruning and Knowledge Distillation

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

... it is better to highlight the difference between the final choice in this paper and the approaches in previous work like Sheared LLaMa.

To help with this comparison, we have created the following table to highlight the differences between our approach and Sheared LLaMa [1]:

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.

Importance estimation for pruning has negligible overhead, since we use a single forward pass on 1024 samples. This costs less than the forward pass on a single step for training the model, which uses 1152 samples.

Regarding the overhead for knowledge distillation, we present ablations using a fixed amount of compute (trained on 128 GPUs for a fixed time) for conventional training vs distillation in Appendix Table 11 (with a 3B model, due to time constraints). We believe that this is more representative of real compute cost than FLOPs. Results show that distillation provides an absolute 1.1% gain in accuracy on the standard LM evaluation harness and a 5+% gain on MMLU when compared to conventional training.

We also obtain the table below for a 4B model for consistency with the model sizes trained in the paper post submission. We plan to replace tables 11 and 12 in the paper with the one below. Rows marked with * below indicate settings with fixed compute/cost. In this table:

1. 4B-Random-Init is a 4B model trained from scratch using conventional ground truth CE loss.
2. 4B-Pruned (prune Nemotron-4 15B) is a 4B model pruned from a 15B model, using conventional ground truth CE loss for retraining.
3. 4B-Pruned-Distill (prune Nemotron-4 15B) is a 4B model pruned from a 15B model, using distillation loss w.r.t. to the 15B teacher for retraining.

We can see orthogonal improvements from 1 to 2, and 2 to 3 using a fixed amount of compute. This experiment demonstrates benefits of using knowledge distillation (+13% improvement in MMLU) with the additional computational cost for online inference factored in.

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

Thank you for pointing this out, we will fix this in the final version.

References:

1. Xia, Mengzhou, et al. "Sheared LLaMa: Accelerating language model pre-training via structured pruning." arXiv preprint arXiv:2310.06694 (2023).

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

Will the authors open-source the code? If the code and data are open-sourced, I would raise my score.

Data: more details on the composition of the Nemotron-4 dataset that we use were recently made public [1,2]. Parmar et al. [1] describe the distribution of domains, type of speech, toxicity, and quality of all our crawl data in Section 6, classifier information for all the above attributes in Appendix B, and information about the multilingual and programming language distribution on page 17. Parmar et al. [2] detail our continuous training recipe in Table 2, Figure 1, 2 and 4. Unfortunately, due to individual licenses for sub-data, we are unable to distribute the full dataset.

Code: the code is planned to be released as part of existing common libraries with a timeline of 1-2 months, with the first candidate being distillation with depth pruning, followed by width pruning (MLP, attention head, embedding). We are working on the modifications that will help apply the method to any existing models in Hugging face, and this requires refactoring.

Models 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.

When pruning a 70B model to 7B, would the method in the paper still work?

We observe that pruning should be performed iteratively in incremental steps. For example, as shown in Table 12 in the appendix, we observe an improvement of 4.6% points on MMLU by doing 15B->8B->4B instead of 15B->4B. Therefore one would expect better results by doing 70B->35B->15B etc.

Pruning models of size 70B and bigger is definitely in our interest. This requires supporting pipeline parallelization in addition to tensor parallelization. We are in the process of adding this functionality to the code-base and plan to release models as soon as they are available.

References:

1. Parmar, Jupinder, et al. "Data, Data Everywhere: A Guide for Pretraining Dataset Construction." arXiv preprint arXiv:2407.06380 (2024).
2. Parmar, Jupinder, et al. "Reuse, Don't Retrain: A Recipe for Continued Pretraining of Language Models." arXiv preprint arXiv:2407.07263 (2024).

We would like to thank the reviewer for pointing out the strengths of the paper and providing valuable feedback.

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 ...

The reviewer raises an important question of evaluation via human preference. We conduct a small-scale (due to time constraints) human evaluation experiment with 7 people and 25 random prompts drawn from the MT-Bench multi-term benchmark. We generate responses from 3 models (all instruction-tuned): Minitron-4B, Gemma-2 2B, and Phi-2 2.7B and ask users to rate responses based on correctness, relevance, conciseness, hallucination and personal preference. For each prompt, the user selects a better model among Minitron and a randomly chosen second model (one of Gemma-2 or Phi-2), or specifies a tie. Model names are anonymized to prevent bias. The results of this experiment are shown in the table below:

Here, we notice that Minitron is selected as a better model in 43.43% of cases. While this is a small-scale experiment at the moment, we will add a larger scale human evaluation to the final paper.

Additionally, to gain more insights into the qualitative performance of Minitron models, we perform supervised fine-tuning (SFT) on Minitron 4B using instruction-tuning data used for Nemotron-4 to create Minitron 4B-instruct, and evaluate it on various tasks, including instruction-following and roleplay (IFEval and MT-Bench), RAG QA (ChatRAG-Bench), and function calling (BFCL). The results for this experiment are shown in Tables 6 to 9 in the 1-page PDF attached to the global rebuttal. Tables 6 to 8 demonstrate that Minitron 4B-instruct has strong instruction-following, roleplay and RAG capabilities, beating similarly sized models across all tasks. On function calling (Table 9), Minitron 4B-instruct outperforms Gemma-2B-IT and even Llama-3-8B-instruct. We will add these results to the final paper.

In Table 2 is the nemotron-8B baseline also trained on the 94B token retraining dataset?

In Table 2, the nemotron-8B baseline is not trained on the exact 94B retraining dataset.

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.

We agree that comparing models trained on exactly the same data is better. However, note that Nemotron-8B model is already trained on 3.8T tokens. Hence, we believe further vanilla training on the retraining dataset would be an unfair comparison and propose to train Nemotron-8B from scratch on the retraining dataset, identical to the retraining routine for Minitron-8B. This would provide an apples to apples comparison for: 1) training from scratch vs 2) our proposed approach, when training a new 8B model.

Going the other way, we also have an ablation study where we train Minitron-8B on exactly the same Phase 1 dataset Nemotron-8B model was trained on:

Minitron outperforms Nemotron-8B on 3 out of 4 benchmarks above despite using a tiny random subset of the 3.8T dataset used by Nemotron 8B.

As requested, we also perform the experiment with the model being trained with 3.8T + 94B tokens. For fairness, we also perform the exact same further vanilla training on Minitron 8B resulting in a model trained with 94B + 94B tokens. The loss curves for the experiments are provided in Figure 1 in the 1-page PDF attached to the global rebuttal. Unfortunately, due to shortage of time, the new training jobs ran to 85% completion and we were unable to run downstream task evaluation.

In the Figure:

1. Gray: Train nemotron-8B from scratch on the retraining dataset
2. Pink: Minitron 8B
3. Orange: Further vanilla training of nemotron-8B on the retraining dataset
4. Blue: Further vanilla training on minitron 8B on the retraining dataset

2 being significantly better than 1, and 4 being better than 3 showcases the efficacy of our technique.

The experiment with KD on the Nemotron-8B model with Nemotron-15B as the teacher is not possible as these models use different tokenizers and therefore the logits will be misaligned between 2 models. Distillation of models with different tokenizers will be a great topic for future work.

With the Nemotron-8B model (trained on 3.8T tokens) as a starting point, we would expect similar results to Minitron (Minitron pruning technique provides significantly better weight initialization compared to random initialization). Using the additional compute required to train Nemotron-8B on 3.8T tokens towards Minitron will make the latter even better.

The figures and tables need much much more detailed captions

We absolutely agree that providing more details in the captions will help understand the underlying message. Some sample updated captions are provided below (unable to provide the full list of 12 changes due to space constraints, but will update them in the final paper):

• Figure 2: High-level overview of our proposed iterative pruning and distillation approach to train a family of smaller LLMs. On a pretrained LLM, we first evaluate importance of neurons, rank them, trim the least important neurons and distill the knowledge from the original LLM to the pruned model. The original model is replaced with the distilled model for the next iteration of compression.

• Figure 5: LM validation loss curve for retraining of two pruned candidates with (L2, L2) and (L2, Mean) metrics for (batch, sequence) aggregation strategies.