LLM Pruning and Distillation in Practice

## Abstract

Abstract Abstract: We present a comprehensive report on compressing the Llama 3.1 8B and Mistral NeMo 12B models to 4B and 8B parameters, respectively, using pruning and distillation [ 1 ] .
We explore two distinct pruning strategies: (1) depth pruning and (2) joint hidden/attention/MLP (width) pruning, and evaluate the results on common benchmarks from the LM Evaluation Harness [ 2 ] . The models are then aligned with NeMo Aligner and tested in instruct-tuned versions. This approach produces a compelling 4B model from Llama 3.1 8B and a state-of-the-art Mistral-NeMo-Minitron-8B (MN-Minitron-8B for brevity) model from Mistral NeMo 12B. We found that with no access to the original data, it is beneficial to slightly fine-tune teacher models on the distillation dataset.
We open-source our base model weights on Hugging Face with a permissive license. Models on Hugging Face: Mistral-NeMo-Minitron-8B-Base | Llama-3.1-Minitron-4B-Width-Base | Llama-3.1-Minitron-4B-Depth-Base

## 1 Introduction

LLM providers often train an entire family of models from scratch, each with a different size (number of parameters, e.g. Llama 3.1 8B, 70B, 405B); this is done to aid users targeting different deployment scales, sizes and compute budgets. However, training multiple multi-billion parameter models from scratch is extremely time-, data- and resource-intensive.

Recent work [] has demonstrated the effectiveness of combining weight pruning with knowledge distillation to significantly reduce the cost of training LLM model families. Here, only the biggest model in the family is trained from scratch; other models are obtained by successively pruning the bigger model(s) and then performing knowledge distillation [] to recover the accuracy of pruned models.

In this report, we successfully apply the Minitron compression strategy [] to two state-of-the-art models: Llama 3.1 8B [] and Mistral NeMo 12B [], compressing them down to 4B and 8B parameters, respectively.
Figure 1 provides a high-level overview of our approach.

While following the original paper [], we make a key modification: due to lack of access to the original training data, we fine-tune the teacher model on our own dataset before pruning and distillation. We refer to this step as teacher correction. Figure 4 shows that omitting teacher correction causes a data distribution mismatch, negatively impacting distillation.

Table 1 provides a summary of our results: our compression strategy yields a state-of-the-art 8B model (MN-Minitron-8B) which outperforms all similarly-sized models across the board on common language modeling benchmarks.
Our Llama-3.1-Minitron-4B models (both depth and width-pruned variants) also exhibit strong accuracy compared to the teacher
Llama 3.1 8B model and the previous-generation Minitron-4B model []; among the two variants, the width-pruned variant outperforms the depth-pruned one.
In terms of runtime inference performance measured using TensorRT-LLM, the Llama-3.1-Minitron-4B models provide an average speedup of 2.7$\times$ and 1.8$\times$ for the depth and width pruned variants, respectively, compared to the teacher Llama 3.1 8B model.

Figure: Figure 1: High-level overview of our proposed pruning and distillation approach. The total number of tokens used for each step is indicated in parentheses.
Refer to caption: x1.png

## 2 Methodology

A high-level overview of our approach is illustrated in Figure 1. Here, the teacher model is first lightly finetuned on the target dataset to be used for distillation - we refer to this step as teacher correction. Next, pruning is applied to compress the model, following which distillation is used to recover any lost model accuracy. We refer the reader to the Minitron paper [] for the full description of the pruning and distillation method.

### 2.1 Pruning

Weight pruning is a powerful and well-known technique for reducing model size. In this report, we focus on structured pruning, where blocks (or channels) of nonzero elements are removed at once from model weights; examples of structured pruning techniques include neuron, attention head, convolutional filter, and depth pruning []. In case of LLMs, as shown in Figure 2, we start the pruning process by first computing the importance of each layer, neuron, head, and embedding dimension. We then sort these importance scores to compute a corresponding importance ranking.

Figure: Figure 2: Pruning and distillation process outlined in the original paper []. We follow the same approach in this work.
Refer to caption: x2.png

Importance Estimation:
We use a purely activation-based importance estimation strategy that simultaneously computes sensitivity information for all the axes we consider (depth, neuron, head, and embedding channel) using a small calibration dataset and only forward propagation passes. We consider depth pruning as a special case and do not combine it with compressing other dimensions.

We compute the importance of each head, neuron and embedding channel by examining the activations produced by the multi-head attention (MHA), multi-layer perceptron (MLP) and LayerNorm layers, respectively. We use a small calibration dataset (1024 samples) for this purpose.

For depth pruning, we consider three distinct metrics for evaluating layer importance: (1) LM validation loss, (2) Block Importance (BI) [] and (3) accuracy on the downstream task. For loss-based ranking, we simply remove a single or a block of contiguous layers and compute its effect on LM loss; this serves as the “importance” or sensitivity of the layer. BI uses the cosine distance between the input and output of a layer or a block of layers. We notice that BI and LM loss metrics are highly correlated but do not produce the most accurate pruned model on downstream tasks as shown in Figures 8 and 9. We thus evaluate layer importance using the Winogrande benchmark [].

Model Trimming:
As shown in Figure 2, for a given architecture configuration, we first rank the elements of each axis according to the computed importance and perform trimming (reshaping) of the corresponding weight matrices directly. For neuron and head pruning, we trim MLP and MHA layer weights, respectively. In the case of embedding channels, we trim the embedding dimension of the weight matrices in MLP, MHA, and LayerNorm layers.
The original approach ([]) uses Neural Architecture Search (NAS) to find the best architecture; in this work, we skip this step and instead utilize the network architecture-related learnings from the original paper.

### 2.2 Retraining with Distillation

We use the term retraining to refer to the accuracy recovery process following pruning. In this work, we explore two retraining strategies: (1) conventional training, leveraging ground truth labels, and (2) knowledge distillation using supervision from the unpruned model (teacher).
Knowledge Distillation (KD) [] involves transfer of knowledge from a larger or more complex model called the teacher to a smaller/simpler model called the student. The knowledge transfer is achieved by having the student model mimic the output and/or the intermediate states of the teacher model. In our case, the uncompressed and pruned models correspond to the teacher and student, respectively. For distillation, we follow best practices from our previous work [] and use forward KL Divergence loss [] on the teacher and student logits only (following the []). This is illustrated in Figure 3.

Figure: Figure 3: Overview of Distillation: If the original training data is unavailable, a slight fine-tuning of the teacher model is recommended. Distillation is then performed by minimizing KL divergence on the logits, with the original model as the teacher and the pruned model as the student.
Refer to caption: x3.png

backgroundcolor=gray!10,
roundcorner=10pt

## 3 Training Details

### 3.1 Pre-training

Llama 3.1 8B [] and Mistral NeMo [] 12B are pretrained on different proprietary datasets, which we do not have access to. According to the Llama 3.1 tech report [], the 8B model is pretrained on 15T tokens. We start with the corresponding Base models that are openly available online on Hugging Face.

##### Dataset:

We use the Nemotron-4 curated continued training (CT) dataset []  [] for all our pruning and distillation experiments.

### 3.2 Pruning

Our simplified pruning recipe is based on the best practices outlined in the Minitron paper [] and is described in the Methodology section.
Specifically, for width pruning, we (1) use l2-norm and mean as the aggregation functions across the batch and sequence dimensions, respectively, and (2) perform single-shot pruning, avoiding iterative approaches.
For depth pruning, as described in the Methodology section, we follow the observations from Gromov et al. [] and drop a continuous subgroup of layers that results in the least accuracy drop on Winogrande [].
In this work, we skip the lightweight neural architecture search (NAS) phase, and go with a manual architecture configuration for both Llama-3.1-Minitron-4B and MN-Minitron-8B. The architectures we come up with are inspired by the Minitron-4B and Minitron-8B models, and are detailed in Table 3. We now describe the pruning recipes for each of our target compressed models:

**Table 3: Architecture details of our compressed models.**
|  | LLaMa-3.1-Minitron-4B | MN-Minitron |  |
| --- | --- | --- | --- |
|  | Width | Depth | 8B |
| Total params | 4.5B | 4.5B | 8.4B |
| Non-Emb params | 3.7B | 3.5B | 7.3B |
| Hidden size | 3072 | 4096 | 4096 |
| Vocabulary | 128256 | 128256 | 131072 |
| MLP hidden dim | 9216 | 14336 | 11520 |
| Depth | 32 | 16 | 40 |
| Attention groups | 8 | 8 | 8 |
| Query heads | 32 | 32 | 32 |
| Head dimension | 128 | 128 | 128 |

backgroundcolor=gray!10,
roundcorner=10pt

Llama-3.1-Minitron-4B-Width:

- •
Starting model: Llama 3.1 8B
- •
Hidden dimension: 4096 $\rightarrow$ 3072
- •
MLP hidden dimension: 14336 $\rightarrow$ 9216
- •
Attention heads: unchanged
- •
Depth: unchanged

Llama-3.1-Minitron-4B-Depth:

- •
Starting model: Llama 3.1 8B
- •
Hidden dimension: unchanged
- •
MLP hidden dimension: unchanged
- •
Attention heads: unchanged
- •
Depth: 32 $\rightarrow$ 16

MN-Minitron-8B:

- •
Starting model: Mistral NeMo 12B
- •
Hidden dimension: 5120 $\rightarrow$ 4096
- •
MLP hidden dimension: 14336 $\rightarrow$ 11520
- •
Attention heads: unchanged
- •
Depth: unchanged

### 3.3 Distillation

##### Teacher Correction:

Using the Mistral NeMo 12B model directly as a teacher performs sub-optimally on our dataset. This is due to the change in distribution of sub-word tokens across the original dataset the teacher model was trained on vs. the dataset being distilled on. To account for this, we first fine-tune the teacher on our dataset using $\sim$127B tokens. As shown in Figure 4, such a correction is essential if the original dataset is not available during distillation. We thus apply this technique on both the Mistral-NeMo and Llama-3.1 teacher models. The fine-tuning process has a minor effect on the teacher model’s accuracy on downstream tasks, with some tasks improving and some degrading as shown in Table  1. We hypothesize this to be an artifact of the dataset used for fine-tuning.

##### Retraining:

Following the learnings in the Minitron work [], we opt for logit-only distillation, minimizing the forward KL Divergence [] loss across the teacher and student probabilities, and ignore the LM cross-entropy loss altogether. Here, the unpruned and pruned models correspond to the teacher and student, respectively. We use the hyperparameters listed in Table 4 during distillation. We use 32 NVIDIA DGX H100 nodes for our training jobs.

**Table 4: Hyperparameters used during distillation-based retraining.**
|  | Llama-3.1- | MN-Minitron |
| --- | --- | --- |
|  | Minitron-4B | 8B |
| Peak learning rate | 1e-4 | 1e-4 |
| Min learning rate | 1e-5 | 4.5e-7 |
| Warm-up steps | 40 steps | 60 steps |
| LR decay schedule | Cosine | Cosine |
| Global batch size | 1152 | 768 |
| Context length | 8192 | 8192 |
| Total tokens | 94B | 380B |

Figure: Figure 4: Training convergence plot for the compressed 8B student model. We compare supervision from the original teacher and the corrected teacher.
Refer to caption: x4.png

### 3.4 Instruction Tuning

To evaluate the instruction-following capabilities of our distilled models, we perform supervised fine-tuning (SFT) on the Llama-3.1-Minitron 4B models using NeMo-Aligner [] with the instruction tuning dataset used for Nemotron-4 340B []. As shown in Table 2, we evaluate the aligned models for instruction-
following and roleplay (IFEval [] and MT-Bench []), RAG QA (ChatRAG-Bench []), and function-calling capabilities (BFCL []).

## 4 Analysis

We perform a series of ablation studies to better understand the compression characteristics of these newer models. We report our results in this section.

##### Width vs Depth Pruning:

Figure 5 shows the training curve of Llama-3.1-Minitron-4B pruned for width vs. depth. We notice that width pruning results in smaller initial loss and consistently outperforms the depth-pruned model, despite both variants having the same number of parameters.

##### Pruning and Distillation:

Figure  6 demonstrates orthogonal benefits of our proposed approach with pruning and distillation. We compare (1) random weight initialization and distillation, (2) random pruning and distillation, where components are pruned randomly ignoring the importance scores, (3) our proposed pruning with typical cross entropy based LM loss training and (4) our proposed pruning with distillation-based training. We notice that pruning results in a significantly better starting point compared to random initialization, and also that distillation-based training outperforms conventional training methods while requiring significantly fewer training tokens (up to $50\times$ in our case).

Figure: Figure 5: Convergence of width- and depth-pruned Llama 3.1 8B to 4B models. Width pruning consistently outperforms depth pruning for a given parameter budget.
Refer to caption: x5.png

Figure: Figure 6: Training convergence plot for Mistral Nemo 12B compressed model. We compare (a) random initialization with distillation, (b) randomly pruned weights with distillation, (c) pruning with standard LM loss, and (d) our pipeline with pruning and distillation.
Refer to caption: x6.png

##### Teacher Correction:

We compare two approaches for teacher correction: (1) pruning and distilling the corrected teacher, and (2) pruning the original teacher and distilling from a continuously corrected teacher. The results in Figure 7 suggest that teacher correction doesn’t affect the optimality of pruning, and that distillation from a corrected teacher is crucial. Teacher correction can be performed in parallel with distillation to bridge the gap.

Figure: Figure 7: Training convergence plot for Mistral Nemo 12B compressed model. We compare (1) pruning and distilling the corrected teacher with (2) pruning the original teacher and distilling from a continuously corrected teacher.
Refer to caption: x7.png

##### Depth Pruning Metrics:

when examining how LM validation loss increases as contiguous blocks of layers are removed (Figure  8), we observe that the layers at the beginning and end are the most important. Removing non-contiguous layers can result in even better LM validation loss (the dashed line). However, this observation does not necessarily hold when evaluating downstream task performance.
Figure  9 shows that dropping 16 layers selected based on per-layer importance ( [, ]) yields a random Winogrande accuracy of 0.5, while removing layers 16 to 31 continuously ( []) results in an accuracy of 0.595. The gap holds during distillation-based retraining and we opt for the latter approach.

Figure: Figure 8: LM loss value on validation set after removing 1, 2, 8 or 16 contiguous layers with Llama 3.1 8B. For example, the purple line at layer no. 16 indicates the LM loss if we dropped the first 16 layers. Layer no. 17 indicates the LM loss if we leave the first layer intact and drop layers 2 to 17. The dashed line corresponds to LM loss value when removing 16 non-contiguous layers least increasing the loss.
Refer to caption: x8.png

Figure: Figure 9: Accuracy on the Winogrande task when removing 16 contiguous layers with Llama 3.1 8B. Layer no. 17 indicates the LM loss if we leave the first layer intact and drop layers 2 to 17. The dashed line corresponds to the accuracy when removing 16 non-contiguous layers least increasing the loss.
Refer to caption: x9.png

## 5 Evaluation

##### Benchmarks

following Touvron et al. [], we evaluate our compressed models on a series of downstream tasks, including MMLU [], HumanEval [] for Python code generation, several question-answering datasets for common-sense reasoning: Arc-C [], HellaSwag [], TruthfulQA [] and WinoGrande [] and XL-Sum English [] for summarization. We report the 5-shot performance on MMLU, 5-shot on Winogrande, 25-shot on ARC-Challenge, 10-shot on HellaSwag, 0-shot on 20% of XL-Sum and average pass@1 scores for HumanEval and MBPP. For pass@1 scores we use a temperature of 0.2 and nucleus sampling [] with top-p $=$ 0.95.
For instruction-tuned models, we use MT-Bench [], Instruction-Following Eval (IFEval) [], ChatRAG-Bench [], and Berkeley Function Calling Leaderboard (BFCL) [].

### 5.1 Base Models

Base model evaluation results are shown in Table 1. Compared to similarly-sized models, MN-Minitron-8B demonstrates superior accuracy across the board, outperforming the recent Llama 3.1 8B model
using 40$\times$ fewer training tokens (380B vs. 15T).
Similarly, the Llama-3.1-Minitron 4B models perform favorably compared to the teacher Llama 3.1 8B model using $150\times$ fewer training tokens (94B vs. 15T); our pruned Llama models also outperform the previous generation Minitron 4B model. We note from Table 1 that the width-pruned variant outperforms the depth-pruned one.
These results clearly demonstrate the advantages of our methodology: state-of-the-art accuracy coupled with an order of magnitude improvement in training efficiency.

### 5.2 Instruct Models

The performance of the instruction-tuned Llama-3.1-Minitron 4B variants is shown in Table 2. We compare the Llama-3.1-Minitron 4B variants to other similarly-sized baselines and notice that our models demonstrate strong instruction-following and roleplay capabilities, only lagging behind Gemma2 in IFEval [] and MT-Bench []. On retrieval based question answering (ChatRAG-Bench []) and function-calling (BFCL []), Minitron models achieve state-of-the-art performance.

### 5.3 Runtime Performance Analysis

We optimized the Llama 3.1 8B and Llama-3.1-Minitron 4B variants with NVIDIA TensorRT-LLM, an open-source toolkit for optimized LLM inference.

Figure: Figure 10: TensorRT-LLM FP8 throughput comparison for the Llama-3.1-Minitron-4B models with the Llama 3.1 8B model w.r.t. increasing input and output sequence lengths.
Refer to caption: extracted/5814671/fig/techreport/fp8-throughput-llama-3-1-minitron-4b.png

Figure 10 shows the throughput in requests per second for the various models in FP8 precision obtained on a single H100 80 GB GPU. Different use cases are represented by increasing input sequence length/output sequence length (ISL/OSL) combinations, at a batch size of 32 and 64 for the 8B-12B models and the 4B models respectively. The smaller memory footprint of the 4B model allows for larger batches.
We notice that Llama-3.1-Minitron-4B-Depth is the fastest, achieving an average throughput improvement of $2.7\times$ over Llama 3.1 8B; the Llama-3.1-Minitron-4B-Width variant achieves an average throughput improvement of $1.8\times$ over Llama 3.1 8B.
Compared to BF16, we notice that FP8 delivers a performance boost of $1.4\times$.

## 6 Insights

In this Section, we summarize some interesting and surprising observations.

#### 6.0.1 General

- 1.
Teacher correction is crucial for distillation to work optimally on a new, unseen dataset. Fine-tuning the teacher with the dataset used for distillation in this manner yields over a 6% reduction in LM validation loss. Teacher correction doesn’t affect the optimality of pruning and can even be performed in parallel with distillation.
- 2.
In line with the Minitron paper’s observations, we require only 380B tokens to achieve state-of-the-art accuracy post pruning with distillation.
- 3.
For width pruning, we achieve stronger accuracy by retaining attention heads and pruning the other dimensions (MLP intermediate dimension, embedding channels).

#### 6.0.2 Mistral NeMo 12B to MN-Minitron-8B:

- 1.
Our compressed model outperforms the teacher on two benchmarks, GSM8k and HumanEval after pruning and distillation: GSM8k increases from 55.7% to 58.5% and HumanEval increases from 23.8% to 36.2%. This improvement is likely influenced by the dataset. However, retraining is performed using the distillation loss alone.

#### 6.0.3 Llama 3.1 8B to Llama-3.1-Minitron 4B:

- 1.
Width pruning delivers better accuracy with MMLU at 60.5%, while depth pruning yields 58.7%, for Llama-3.1 compression.
- 2.
Reasoning ability is impacted further significantly, with GSM8K accuracy at 41.24% for width and 16.8% for depth.
- 3.
Depth pruning boosts throughput, achieving $2.7\times$ speedup over Llama-3.1 8B, while width pruning provides $1.7\times$ speedup.
- 4.
For depth pruning, we observe that dropping contiguous layers from the model is more effective than using non-contiguous, importance-based pruning.

## 7 Acknowledgments

This work would not have been possible without contributions from many people at NVIDIA. To mention a few:

Foundational Model: Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Pavlo Molchanov, Mostofa Patwary, Daniel Korzekwa, Ashwath Aithal, Mohammad Shoeybi, Bryan Catanzaro and Jan Kautz

Alignment: Ameya Sunil Mahabaleshwarkar, Hayley Ross, Brandon Rowlett, Oluwatobi Olabiyi, Shizhe Diao and Yoshi Suhara

Datasets: Sanjeev Satheesh, Jupinder Parmar, Shengyang Sun, Jiaqi Zeng, Zhilin Wang, Yi Dong, Zihan Liu, Rajarshi Roy, Wei Ping, Makesh Narsimhan Sreedhar and Oleksii Kuchaiev

TensorRT-LLM: Bobby Chen, James Shen and Chenhan Yu

Hugging Face Support: Ao Tang, Yoshi Suhara and Greg Heinrich

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Given the newly released llama-3.2-1b and llama-3.2-3b and ministral-3b models, could the authors compare against these model families?
- I currently lack a concrete answer to the question : Does a (student) model pruned and distilled from a larger model outperform a (student) model of similar size trained from scratch on a pretraining dataset? Could the authors elaborate on their observations here? 
- Have the authors tried to replace full-fine-tuning with PEFT techniques like LoRA[1] and Galore[2]? If yes do the observations about distillation and importance sorting hold there?
- Could the authors provide on-device latency gains using the pruned models eg: on a A100/H100?

I am willing to raise my score if the authors adequately respond to my questions and the weaknesses pointed out above.

[1] Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.

[2] Zhao, J., Zhang, Z., Chen, B., Wang, Z., Anandkumar, A. and Tian, Y., 2024. Galore: Memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507.
soundness: 3
strengths: Experiments
- The paper evaluates the proposed teacher correction scheme on different models (llama-3.1-8B and mistral nemo 12B) for a more thorough study of effectiveness of distillation (pre-corrected or continuously corrected)
- Detailed evaluation of gains from different components random student initialisation, importance sorting and knowledge distillation
- Exhaustive evaluation on different tasks

Writing
- Paper is written in an easy to understand manner in most parts

Originality and significance 
- While most parts of the paper are similar to/derive from observations in [1], the empirical investigation into precise gains by each of the components is interesting and significant. 

[1] Muralidharan, S., Sreenivas, S.T., Joshi, R., Chochowski, M., Patwary, M., Shoeybi, M., Catanzaro, B., Kautz, J. and Molchanov, P., 2024. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407.14679.
summary: This paper studies the combination of structural pruning and knowledge distillation to obtain compressed and performant language models with higher throughput. The paper proposes a "teacher correction" to allow for lightweight finetuning of the teacher model (a llama-3.1-8B or mistral nemo 12B). The step is proposed with the goal of adapting the distribution of the teacher to the domain of the finetuning data before the distillation step to a smaller model. Furthermore, the paper uses different saliency metrics to compute importance of depth and intermediate mlp/attention/width dimension and sort the teacher in order of importance, before applying pruning. In addition, the paper also explores different saliency metrics for improved depth pruning. The pruned models are evaluated on a wide variety of commonsense, math and instruction tuning tasks.
weaknesses: Originality and Significance
- A lot of this work is similar to  [1], including importance sorting, architecture selection procedure, distillation loss
- Domain adaptation of a teacher is also studied in [2] for BERT like models and I don't find the proposed teacher correction to be significantly different from the idea of adapting a teacher to a specific domain on interest
- The idea of dropping last contiguous layers (except the last layer) have been studied in [3] and the scheme chosen here is similar to the scheme chosen in the paper. Is my understanding correct and could the authors present the key differences/observations here?

Experiments
- Check questions section, I think the paper would benefit a lot from adding comparisons to similar small sized models, trained from scratch. 

Clarity
- While the paper is written in an easy to understand manner in most parts, I am lacking a cohesive story of the overall approach of the paper. Could the authors provide an algorithm describing the final/best performing choices for pruning+distillation? 
- Could you elaborate on the following points which are ambiguous/unclear : 167-168, what are the architecture related learnings, how was the student designed?  Line 133-134, which calibration set is used, does the domain of the calibration set matter? "continuously corrected teacher", is not defined properly in the paper, is the teacher periodically updated? line 312-313, what is the general scheme used for pruning different teachers, which contiguous layer indices are to be dropped?

Scaling to larger models
- As per my understanding the paper studies only full fine-tuning (FFT) of the teacher model, for teacher correction. However, this does not scale to larger models (eg: llama-3.1-70B) due to memory/compute constraints and hence cannot exploit larger/better teacher models.
- While the authors do mention the possible use of Parameter-Efficient-Fine-Tuning (PEFT) schemes like LoRA[4] and Galore[5], I think given the focus of the paper on large language models, generalisability of the observations in the paper for PEFT schemes on even larger language models, should be studied in detail. 

Code and Reproducibility
- The models are finetuned and distilled on a proprietary dataset, how do the observations translate to public datasets eg: openwebtext or finetuning datasets like alpaca, commonsense170k, math10k?
- The paper does not release the code and experimental pipeline for their work and I encourage the authors to do this to increase the impact and adoptability of their work. 

[1] Muralidharan, S., Sreenivas, S.T., Joshi, R., Chochowski, M., Patwary, M., Shoeybi, M., Catanzaro, B., Kautz, J. and Molchanov, P., 2024. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407.14679.

[2] Yao, Y., Huang, S., Wang, W., Dong, L. and Wei, F., 2021. Adapt-and-Distill: Developing Small, Fast and Effective Pretrained Language Models for Domains. Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021.

[3] Gromov, A., Tirumala, K., Shapourian, H., Glorioso, P. and Roberts, D.A., 2024. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887.

[4] Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.

[5] Zhao, J., Zhang, Z., Chen, B., Wang, Z., Anandkumar, A. and Tian, Y., 2024. Galore: Memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - Why Winogrande was chosen over other potential downstream tasks?
- How sensitive is the method to the choice of downstream task for depth pruning?
soundness: 3
strengths: - The paper introduces teacher correction as an innovative solution for compressing models without original training data access - a common industry challenge with proprietary models. The approach is elegantly simple yet effective, requiring only ~100B tokens for adaptation and integrating smoothly into existing distillation pipelines. This practical focus makes the method immediately valuable for real-world applications.
- The results show impressive efficiency gains with the MN-COMPRESSED-8B achieving SOTA performance using 40× fewer training tokens, and LLAMA 3.1-COMPRESSED-4B showing competitive results with 150× fewer tokens. The method delivers significant speedups (2.7× for depth pruning, 1.8× for width pruning) while maintaining strong performance across diverse tasks including reasoning, coding, and instruction following.
summary: This paper addresses an increasingly important challenge in the field of large language models (LLMs): how to effectively compress models when access to the original training data is restricted. The authors present a practical approach combining structured pruning and knowledge distillation, introducing several notable innovations in the process. The key contributions include a novel "teacher correction" phase for adapting the teacher model to target dataset distribution, an improved depth pruning saliency metric based on downstream task performance, and empirical validation through compression of Mistral NeMo 12B to 8B and Llama 3.1 8B to 4B parameters.
weaknesses: - The performance degradation observed in the compressed models is substantial, especially considering the modest compression ratios achieved. When compressing Llama 3.1 8B to 4B parameters (only a 2x reduction), the MMLU performance drops notably from 65.3% to 60.5% with width pruning, or even lower to 58.7% with depth pruning. This performance drop is particularly concerning when viewed against recent developments in the field - for instance, MobileLLM-350M demonstrates that it's possible to achieve comparable performance to LLaMA-v2 7B in specific tasks with a model that's 20 times smaller. The fact that this paper shows significant performance degradation with just 2x compression, while requiring additional fine-tuning, makes the proposed approach less compelling for practical applications. 

- The evaluation of the method's effectiveness is hampered by insufficient baseline comparisons. The authors should have compared their approach against a broader spectrum of existing compression techniques, including various pruning approaches, quantization methods. 

- The paper's experimental scope is also notably limited. While the authors validate their approach on two recent and popular models (Mistral NeMo 12B and Llama 3.1 8B), they don't explore more aggressive compression scenarios such as creating sub-1B parameter models, which would be particularly valuable for resource-constrained deployments.

[MobileLLM] MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. There are several methods for measuring the importance of layers, including SLEB, Shortened LLaMA, and ShortGPT. For a stronger publication, I suggest that the authors include a comparative analysis section that directly compares their proposed importance measurement method with these existing methods on a common set of metrics or benchmarks.

2. Is there any empirical validation for the authors' argument: "We hypothesize this is due to the change in the distribution of sub-word tokens across the original dataset the teacher model was trained on versus the dataset being distilled on"? 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.
soundness: 2
strengths: 1. The authors provide their base model on Hugging Face, and based on this, the results seem promising.

2. Teacher correction can improve the effectiveness of distillation when the original data is unavailable, as demonstrated with Llama 3.1 8B and Mistral NeMo 12B.

3. The authors propose a new method for measuring importance in pruning, but there seems to be a lack of comparison with other methods for evaluating LLM importance.
summary: The paper explores the practical application of structured pruning and knowledge distillation to compress large language models (LLMs) efficiently. The main goal is to create smaller models with fewer parameters while maintaining performance, even when the original pretraining dataset is unavailable. To achieve this, authors proposed

**Teacher Correction**: A novel phase where the teacher model is fine-tuned with a small amount of new data to adapt to the specific distribution of the target dataset. This improves the distillation process when the original training data is inaccessible.

**Pruning Strategies**: Authors explored two pruning strategies, including depth pruning and width pruning with their proposed importance measurement for each component.
weaknesses: 1. This paper does not follow the ICLR 2025 style guidelines, which could be a serious reason for desk rejection.

2. The overall writing is somewhat unclear. I believe the writing could be improved, and some details in the paper also need enhancement, including the proper use of \citep and \citet, as well as capitalization at the beginning of sentences.

3. The reasoning behind the improvement from teacher correction is not very clear. Even though it is a crucial component of the work, it seems to only apply to a single case. I recommend that the authors provide more analysis or experiments demonstrating the effectiveness of teacher correction across different models or datasets. This would help establish whether the improvement is generalizable or limited to a specific case.

4. Overall, the main components of the paper—distillation and pruning—appear to overlap significantly with previous works, particularly Sheared LLaMA, which first proposed efficient training using distillation and pruning. However, the paper does not cite this work.

Meta-review: This paper tackles an increasingly relevant challenge in the field of large language models (LLMs): compressing models effectively when access to the original training data is restricted. The key contributions include a novel "teacher correction" phase to adapt the teacher model to the target dataset distribution, an enhanced depth pruning saliency metric based on downstream task performance, and empirical validation demonstrated through compressing Mistral NeMo 12B to 8B and Llama 3.1 8B to 4B parameters.

The strengths of the paper lie in its focus on an important research question and its demonstrated efficiency improvements. However, the work exhibits significant weaknesses, including limited technical novelty and non-compliance with the ICLR style guidelines.

Given these limitations, I recommend rejecting this submission.

Decision: Reject