Compact Language Models via Pruning and Knowledge Distillation
Authors: Saurav Muralidharan
Sharath Turuvekere Sreenivas
Raviraj Bhuminand Joshi
Marcin Chochowski
Mostofa Patwary
Mohammad Shoeybi
Bryan Catanzaro
Jan Kautz
Pavlo Molchanov
Venue: NeurIPS 2024 poster (2024)
Keywords: llm
pruning
distillation
compression
Primary area: natural_language_processing
TL;DR: A novel structured pruning and retraining approach for LLMs that produces strong models using a small fraction of original training data

Abstract: 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. On these tasks, we perform better than Nemotron-3 8B and LLaMa2 7B using **up to 40x** fewer training tokens}, on par with Mistral 7B and Gemma 7B using **up to 85x fewer tokens** and slightly worse than LLaMa3 8B using **up to 159x fewer tokens**. Our models also compare favorably to state-of-the-art compression techniques from the literature.

Compact Language Models via Pruning and Knowledge Distillation
Saurav Muralidharan"
Sharath Turuvekere Sreenivas"
Raviraj Joshi
Marcin Chochowski
Mostofa Patwary
Mohammad Shoeybi
Bryan Catanzaro
Jan Kautz Pavlo Molchanov
NVIDIA
f sharatht, raviraaj, ravirajj mcchocowkkch m seeatwarr,mshooeee.. bcatanzaro, pwl.utzz.oollllaoooooo chc.o
Abstract
Large language models (LLMs) targeting different deployment scales and sizes are currently produced by training each variant from scratch; this extremely compute-intensive. In this paper, investigate if pruning existing LLM and then re-training it with a fraction (<3%) the original training data can be a suitable alternative to repeated. full retraining. To this end, we develop 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 detailed empirical exploration pruning strategies for each axis, methods to combine axes, distillation strategies, and search techniques for arriving at optimal compressed architectures. We use this guide compress the Nemotron-4 family of LLMs by 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.8 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 other community models such as Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-arr compression techniques from the literature. We have open-sourced MINITRON model weights on Huggingface ! with corresponding supplementary material including example code available on GitHub
Introduction
Minitron Gemma78
Nemotron-4 156
LLaMi-3
65
Mistral 70
Large language models (LLMs) now dominate real- world natural language processing and have demon- strated excellent proficiency in understanding diffi- cult contexts 77
better
Minitron 4B
Score
35.
49 48). To aid users tar- geting different deployment sizes and scales, model providers often train an entire family of models from scratch, each with different size (number of param- eters). For instance, the LLaMa-2 model family [49 includes three different variants with 7, 13, and 70 bil-
40
52
Minitron fresulting models)
Pruning
Pruning Path
Gemma228
5.0
10.0
12.
15.0
(trillion takens)
Equul contribution.
Minitron Collection on HuggingFace httrp::://gichuub.mmm//aabaab hwt.///ttraaanninrrooooo
Figure 1: MINITRON results. Our approach greatly reduces training costs for additional models (40x) while improving accuracy.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

DEP
MLP
ATT
EMB
Distillation Loss
LM Val Loss
5.35 - 0.38
2,062
6.33 0.37
2,049
5.07 - 0.42
2.101
8.35 0.49
2.155
Train from scratch (random init)
12:27 - 2.34
3,953
Table 1: Demonstration of how various pruning strategies perform before and after lightweight retraining using 1.8B tokens. We prune the Nemotron-4 15B model down to the size of Nemotron-3 8B and report the change in distillation loss (KL divergence (28) on logits) and the final LM validation loss with retraining. We see that width (attention, MLP, embedding) pruning outperforms depth, but only after retraining. The last row shows change in loss for the Nemotron-3 8B model.
lion parameters, while the Pythia family 6 offers a selection of eight models with sizes ranging from 80 million to 12 billion parameters. However, training multiple multi--illion parameter models from scratch extremely time, data and resource-intensive In this paper, we ask the following question: can we train one big model, and obtain smaller, more accurate (w.r.. training from scratch) models from it through a combination of weight pruning and retraining, while only using a small fruction of the original training data? Achieving such goal would make producing LLMS targeting different deployment scales significantly cheaper. Weight pruning a powerful and well-known technique for reducing model size [5122] 21). In this paper, we focus structured pruning. where blocks of nonzero elements are removed at once from model weights; examples of structured pruning techniques include neuron, attention head, convolutional filter, and depth pruning 32
smallfruction
18 53
55 26 While the
literature is rich with numerous papers on structured pruming, to an end-user, it's not always clear which technique to use, when, and how to combine them to consistently obtain good pruned models. Pruning is also often accompanied by some amount of retraining for accuracy recovery [5T]; this phase is extremely expensive in modern L.Mss, often requiring access to large amounts of curated data. To the best of our knowledge, no existing work on structured pruning explores data-efficient retraining techniques such distillation minimize retraining cost.
In this paper, we perform a thorough empirical exploration of structured pruning and retraining across multiple axes: neurons in feed-forward layers, heads in multi-head attention layers, embedding channels, and model depth. Through our experiments, we gain valuable non-trivial insights on the metrics and hyper-parann paperr-prramters use for each axis and how to effectively combine axes for higher compression rates. For instance, we discover that pruning neurons and heads alone is initially superior to pruming neurons. heads and embedding channels; however, after few steps of retraining, this order flips. Similarly, we discover that width pruning works better than depth, but only after some retraining (see Table [ for concrete example). We also investigate in detail how pruned model can be efficiently retrained for optimal performance using minimal additional data. Based on our findings, we develop a practical list of LLM compression and retraining best practices. Finally, apply our findings to prune the Nemotron-4 15B model [3] and produce a family of smaller models, named MINITRON, that compare favorably to similarly-sized models. MINITRON 8B achieves better accuracy than Nemotron-3 8B (39] (using 40x fewer training tokens) and LLaMa-2 7B [49), and comparable accuracy Mistral- 7B
(25), Gemma 7B [8] and Llama-3 8B; likewise, MINITRON 4B outperforms the similarly-sizee Gemma2 model and compares favorably to the Phi-2 model.
This paper makes the following key contributions:
1. Provides the first thorough empirical exploration of structured pruning and retraining in LLMs across multiple axes. It offers valuable insights on metrics and hyper-parameter settings for pruning, order of pruning, effects of combining different axes, and retraining techniques focusing on data efficiency.
2. Presents a list of effective and practical LLM compression and retraining best practices grounded in extensive empirical evidence.
3. Introduces the MINITRON family of L.Mss, which are obtained through direct pruning of the N 15B model. Deriving MINITRON models from Nemotron-4 15B requires up to 40x fewer training tokens compared to training from scratch, while still (1) comparing favorably to various popular community LLMs of similar size, and (2) outperforming state-of-the-art depth and width-pruned models from the literature.

---

2 Pruning Methodology
As shown in Figure 2] start the pruning process by first computing the importance of each layer, neuron, head, and embedding dimension and then sorting these importance scores to compute a corresponding importance ranking. In this section, detail how rankings are computed for each axis and then subsequently used to obtain pruned model.
2.1 Background and Notation
We begin with some formal definitions. Multi-Layer Perceptron (MLP) layers have two linear layers with a non-linear activation in between: MLP(X) = 5 WF) W2 here, X de- notes the input, and W and W2 are the two associated weight matrices in the MLP layer. W 2 E Rduciten Xdmodsi, where dmodel and dhidden are the embedding and MLP hidden di- mensions, respectively. 8(-) refers the non-linear activation function.
We define the Multi-Head Attention (MHA) operation for an input X as follows: MHA(X) = Connca(bhcad,,...eet
W and head, = Attn XW K, XWV")) here, a WK, V E Rd Xdmodel and W E RR Iddead Xdmddel where disaa is the size of single attention head, and L is the total number of heads.
Finally, the Layer Normalization operation (LayerNorm) 5 on an input is defined as follows: (N X-# @ + 3, where ju and o² represent the mean and variance across the embedding Va dimensions, is small value for numerical stability, and y and B are learnable parameters.
2.2 Importance Analysis
Estimating the importance or sensitivity of individual neural network components such neurons, attention heads, and layers is a well-studied area 9 13 41 In the context of LLMs, recent work has highlighted the ineffectiveness of traditional metries such weight magnitude for estimating importance [33]; instead, recent work on structured pruning of LLMs has focused on metrics such as gradient/Taylor
33 cosine similarity [34], and perplexity on a calibration dataset (26).
Owing their enormous size, computing gradient information on modern LLMs prohibitively memory and compute-intensive and one of our primary goals is to avoid this expensive step when trying to obtain importance information. In this paper, we propose purely activation-buss tance estimation strategy that simultaneously computes sensitivity information for all the axes we consider (depth, neuron, head, and embedding channel) using small (1024 samples) calibration dataset and only forward propagation passes. We now describe how this strategy is implemented for each individual axis.
impor-
Trained LLM
2. Estimate importance
fianstomm Bookk
room
Attention
Layer nomm
MLLP
Lioper
(mbii)
CH
Layer
Feee
Distillation
Iterative
Trim
Rank
CMI
CIIT
CIR 4B
Emioa
Emiito
CHI4
EmSO
Figure 2: High-level overview of our proposed iterative pruning and distillation approach to train 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.
3

---

we compute the importance of each head, neuron and embedding channel by examining the activations produced by the MHA, MLP and LayerNorm layers. respectively. We use small calibration dataset D) for this purpose Formally, we compute activation-based importance scores for heads, neurons, and embedding channels as: Ir head
Width:
EB.s || Atta(XX Q. XW K. XWWV))ll 122
Free neurrnn = EB.s X(Wi)², and F(ù = EB. LN(X);, Here, Wi refers to the 7th row of the weight matrix W1. EB.s refers to aggregation along the batch and sequence dimensions. We observe from our experiments that performing a simple summation here not always optimal. To this end, we perform detailed evaluation of various aggregation functions along each of these dimensions and their corresponding performance in Table Specifically, for a sequence of scores a n - (S, S)). Layer-wise scores are then summed up to obtain network-wiie -wiork-wiieeaoo importance scores for each axis.
S, we try three functions: (1) mean(abs): 1 >i=1 |S| (hereafter referred to as just mean), (2) L2
norm: 2i=1 S. and (3) variance: 1
Depth (Layers): for depth pruning, we evaluate the importance of each layer using two metries: (I) perplexity (PPI.) (26) and (2) Block Importance (BI) 34 For PPL-baaed ranking, we simply remove a single layer and compute its effect on perplexity of this pruned model; this serves as the "importance"
or sensitivity of the layer [26]. BI 34
uses the cosine distance between the input and output of a layer
x X
to estimate layer sensitivity. The BI score of layer is computed as: BL, = 1 - Ex, [x,, X
where X, refers to the input to layer i, and X denotes the tth row of X, The BI of all layers can be computed in a single forward pass, giving a significant speed advantage over PPL-based importance. Additionally, following Gromov et al. [14], we can extend BI to estimate importance of several contiguous layers at the same time.
Iterative Importance: in this setting, we iteratively alternate between pruning and importance estimation for given axis or combination of axes. Formally, given number of iterations T and source and target dimensions (layers, heads, etc.) d. and d,, respectively, we iteratively compute importance
on dg i. ( dimensions and prune to dg - (i + 1) (de-d) dimensions; E [0, T - 1]. We
evaluate the effectiveness of iterative importance estimation Table 12
2.3 Obtaining a Pruned Model
Figure 2 provides an overview of how pruned models are obtained. For a given architecture configu- ration, 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.
When pruning attention heads, we add the residual info from the pruned heads back into the remaining heads, with the aim of preserving relevant knowledge from the pruned heads. This idea is an MHA analog of Layer Collapse (55] for depth pruning and provides a boost to model accuracy in our experiments. Formally, given L original attention heads head,, headz, head, being pruned to K heads, each new head will have the form (for the ith head): head, - (head, heads (-oddk-i+)) i) for i E [K- - (L- - K), K]. In case of grouped query attention we apply this strategy only to the query heads.
Lightweight Neural Architecture Search: Figure provides an overview of our search strategy for finding optimal architecture configurations. Given search space and parameter budget (left side of the figure), we enumerate all feasible architectures meeting the parameter budget. At this stage, while it's possible to further reduce the search space size using strategies such as genetic search and/or Bayesian optimization, we found that sticking to commonly used neuron, head and embedding dimensions, along with reasonably narrow target parameter range (less than billion) was sufficient to obtain tractable solution sets (less than 20 candidates). The feasible candidates then undergo lightweight retraining (~1.8B tokens in this work). We show in Figure P that this retraining stage stabilizes relative rankings and helps us find more accurate candidate to train further. We note that parameter-efficient fine-tuning techniques such as LoRA [33 can also be applied at this stage; we leave the exploration of such techniques to future work.
WWe provide additional details of the calibration dataset in Section 4

---

bevrch SpMc
PLP
Pus de Soluisee
Arrs
Parth
Ravet
Fro
IvB
S120
FUOP
frumente
PART
LiSetin the KT
Plal
3ZL
VDL
SBL
85
Figure 3: Overview of our neural architecture search algorithm. We perform search on multiple axes: number of layers, attention head count, MLP and embedding dimensions to arrive at set of feasible architectures meeting the parameter budget. refers to retraining.
Embrddiings
Block 3 T-1
Mult-Hood Actonion
Food Forword (MUP)
Block To
Logis
Embatcingg
Losa
Block
LightLos
Embeddingg
Block m. S-1
Mult--iedd Attenton
Feed Forward (NLP)
Block St1. mn
Head
Legis
2 Stadint
Figure 4: Overview of Distillation. A studend model with N layers is distilled from a leacher model with M layers. The student learns by minimizing combination of embedding output loss, logit loss and transformer encoder specific losses mapped across student block S and teacher block T.
3 Retraining
We the term retraining refer the accuracy recovery process following pruning. In this paper, we explore two retraining strategies: (1) conventional training. leveraging ground truth labels, and (2) knowledge distillation using supervision from the unpruned model (teacher).
Retraining with Knowledge Distillation: Knowledge Distillation (KD) involves transfer of knowl- edge from a larger more complex model called the teacher to smaller/simpler model called the student [20]. 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 the uncompressed and pruned models correspond to the teacher and student, respectively.
The output probability distribution of an LLM for given token H; is computed as: p(z,, T) = exp(
where is the softmax temperature and [V] is the vocabulary size. Logit-based KD loss EIV exp
across the sequence of all output tokens is represented as: Liogis = + El- Loss (æ, T), T), (a, T)); here, / (z. 7) and pz(2;, (x, represent the teacher and student probability distributions on the kth token, respectively, and represents the sequence length.
For distillation, we explore various loss functions, and several combinations of intermediate states and mappings across the Transformer model as the loss components, along with their respective trade-off's This is illustrated in Figure The intermediate state-based KD loss across a sequence of
Transformer-specfffe hidden states represented as: Lis Exen E'= Loss (hfa) hih)), where hift and h represent the th teacher and student hidden state for the 7th token, respectively, and / represents the sequence length; H the set of chosen intermediate states. The mismatch in student and teacher hidden states is handled by learning shared linear transformation during distillation to upscale the student hidden state to the teacher hidden state dimension. The hidden states used are always post LayerNorm. We report our experimental results for retraining in Section 4.3
5

---

Models
Benchmark
Metric
Llama-3
Llama-2
Mistral
Gemna
Nemotron-4
Nemotron-3
MINITRON
Parameters
8B
6.7B
7,3B
8.5B
15.6B
8.5B
8.3B
Non-Emb. Params
5,98
6.4B
7.7B
12.5B
6.2B
Training Tokens
>15T
2T
8T
6T
8T
3.8T
94B
Knowlelge Loggk
winogrande
acc
83.6
arc challenge (25)
acc_ norm
57.8
53
60.3
61
58.8
52.8
52.6
hellaswag(10)
acc_norm
82.1
79
83.2
82
84.6
78.5
80.7
gsm8k(5)
acc
nuthfulqaa()
mn2
43.9
39
42.6
45
40.7
36.5
42.6
XI.Sum (20)3)
rougel.
30.9
31
4.80
17
32
30.9
31.2
Coding
MRPP(O)
pass@e
42.4
20
38.8
39
38
27.04
35.2
(==20)()
Pass@P1
Table 2: Performance of our pruned MINITRON 8B model compared multiple baselines: the original Nemotron-4 15B, previous generation Nemotron-3 8B, and multiple community models. MINITRON 8B uses x fewer tokens than Nemotron-3 8B. All evaluations by us, except for entries marked with *, which we report from the corresponding papers.
Models
Benchmark
Metric
Phi-2
Gemma
Gemma2
Qwen2
MiniCPM
MINITRON
Parameters Non-Emb. Params Training Tokens
2.7B
2.5B
2.6B
1.5B
2.7B
4,2B
2.5B
2B
1.3B
2.4B
26B
1.4T
3T
7T
1.IT
94B
winugrande are challengee (25)
acc norm
61
48
55.4
43.9
50.9
MMLU(5)
acc
57.5
42
56.5
Knowledge, Logic
hellaswag(10)
acc_ acc_norm
75.2
72
73.0
66.6
68.3
75.0
gsm8k(5)
acc
23.9
24.1
truthfulqu(()
mc2
44
33
45.9
42.9
XLSum en (20)())
rougeL
MBPP(0)
@
29
29.6
37.4
28.2
Coding
humaneval (a=20k()
pass@
50
24
17.7
31.1
23.3
Table Performance MINITRON model compared to similarly sized community models. All evaluations run by us, except for entries marked with * which we report from the corresponding papers. We only compare base models without SFT and DPO, therefore Phi-3 is excluded.
The total loss L is computed as L LCLM Llogits Lis; where LCLM is the student cross- entropy loss against the ground truth labels, and is a weighting coefficient. As the magnitudes of Llogits and Li, differ significantly, we found that computing dynamically as Lagiis achieves better results compared to using a constant.
Lia
4 Experiments and Results
We evaluate our pruning strategy on the N family models [33 specifically, we compress the Nemotron-4 15B model with 15.6 billion parameters down two target parameter ranges: (1) 8 billion, and (2) billion. We use the NVIDIA Megatron-LMM framework [7 to implement our pruning and distillation algorithms for compression and retraining.
Data and Training Hyperparameters: we use the Nemotron-4 curated trillion token (8T) base pretraining dataset and the continued training dataset (CT) 42
43 We use the 8T training blend for all our ablations and use a combination of both data blends to retrain our final models. Unless otherwise specified, dataset D used for importance estimation consists 1024 samples drawn randomly from the full dataset. We use the same optimizer settings and data split as [43 with cosine LR decay schedule from 2- to 4.5-7
1.8 billion tokens (400 steps) for lightweight retraining, The calibration
Downstream Tasks: following Touvron et al. 49
we evaluate our models of similar size on
series of downstream tasks, including MMLU 19 HumanEyal for Python code generation, several question-answering datasets for common-sense reasoning: Arc-C [10), HellaSwag (56). TruthfulQA [29] and WinoGrande (45) and XL-Sum English [07] for summarization. We report the 5-shot performance on MMLU, 5-shot 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@I scores we use temperature of 0.2 and nucleus sampling [22] with top-p- = 0.95. For instruction-tuned models, we use MT-Bench [57]. Cestructioo-Folloviii Eval (IFEval) [59].
ChatRAG-Bench 30, and Berkeley Function Calling Leaderboard (BFCI.) [54]
4.1 Main Pruning Results
We start by introducing the following list of structured compression best practices:
1. To train a family of LLMs, train the largest one and prune+distill iteratively to smaller LLMs.
2. Use (batch-L2, seq=mean) importance estimation for width axes and PPL/BI for depth.
3. Use single-shot importance estimation; iterative provides no benefit.
4. Prefer width pruning over depth for the model scales we consider (S 15B).
5. Retrain exclusively with distillation loss using KLD instead of conventional training.
6. Use (logit+intermediate statetembedding) distillation when depth reduced significantly.
7. Use logit-only distillation when depth isn't 't reduced significantly.
8. Prune a model closest to the target size.
9. Perform lightweight retraining to stabilize the rankings of searched pruned candidates.
10. If the largest model is trained using a multi phalt-ppass training strategy, it is best to prune and retrain the model obtained from the final stage of training.
We arrive at this list through detailed set of ablations and experiments, and each point is backed by empirical evidence, as we demonstrate in the rest of this section and the Appendix. We use this list obtain MINITRON pruned and retrained models, whose performance is shown in Tables and 3 Here, we compare the performance of our pruned models to multiple baselines: (1) the original Nemotron-4 15B model, (2) the previous generation Nemotron-3 8B model, and (3) a set of similarly- community models, all trained from scratch with trillions of tokens. Evaluation performed on the downstream tasks described earlier in this Section. In both tables, we list the number of full and non-embedding parameters, along with the number of training tokens used to arrive at the model.
We further compare the MINITRON models to state-of-the-are art depth and width-pruned baselines in Table 4 namely, LLM-Pruner [ SliceGPT (4), LaCo (55), ShorGGP (34], and Sheared LLaMa (53) Table 10 (Appendix) lists the architecture details of the Nemotron and MINITRON
models shown in Tables and 3 In the following subsections, will go into more detail on how
arrived at the MINITRON pruned models.
From Table 2 notice that MINITRON 8B compares favorably to the latest community models of the same size. Specifically, we outperform Nemotron-3 8B and LLaMa-2 7B, and perform on par with Mistral 7B, Gemma 7B and LLaMa-3 8B, all while using significantly fewer training tokens. MINITRON 8B also significantly outperforms multiple depth-pruned models of larger size (~> 10B parameters) (Table 4 From Table we notice that smaller model, MINITRON 4B, retains model capabilities better compared to small specialized models that score highly only on some tasks, outperforms the Gemma2 model and is significantly superior to multiple depth and/or width pruned models shown in Table
Instruction Tuning: to better understand how MINITRON models perform after supervised fine- tuning (SFT), we perform SFT on MINITRON using instruction-tuning data used for Nemotron-4 340B (38 to create MINITRON 4B-instruct, and evaluate it on various tasks, including instruction- following and roleplay (IFEval and MT-Bench), RAG QA (ChatRAG-Benill
and function calling (BFCL). The results for this experiment are shown in Tables 5 to g Tables 7 that MINITRON 4B-instruct has strong esstucttoonffollwinng roleplay and RAG capabilities, beating similarly sized models across all tasks. On function calling (Table 8) MINITRON 4B-instruct outperforms Gemma-2B-1 and even Llama-3-8B ---a----r---n-tettt..
Best Practice #1: in summary, Tables 2 8 provide strong empirical evidence to support the claim that training one single big model, and obtaining smaller ones from through pruning + retraining achieves higher accuracy and is extremely cost/compute-efficienn when compared to training them from scratch. Further, our efficient retraining strategy also eliminates the need to curate trillions of tokens of data.
7

---

Models
Benchmark
Metric
LLMPruner
SliceGPT
LaCo
ShortGPT
Sheared LLaMa
MINITRON
8 Ballion
#
9.81
9,98
9.8B
9.8B
8.38
# Non-Er Mon-Emb. Params
9.5B
9.5B
9.5B
9.5B
6.2B
MMLU(() hellaswag(1O)
RIC
25.2
37.1
45.9
54.7
63.88
SOC Borm
67.8
55.7
64.4
66.6
80.7
Parameters
4,91
4.98
4.9B
2.718
4.28
4 Biilion
# Non- Emb. Params
4.5B
4.6B
4.6B
4.6B
2.5B
2.6B
winogrande (5)
64.2
74
are challenge (25)
acc_norm
41.2
50.9
MMLU(5)
23.33
28.92
26.45
43.96
26.4
58.6
hellaswag(10)
Boc_norm
56.46
50.27
55.69
53.02
70.8
75
gsm8k(5)
Rcc
23.96
24.1
Table 4: Performance MINITRON models w.r. recent state-of state-of-ther the------thhrr models obtained through depth/width pruning. Top and bottom halves show results for MINITRON 8B and 4B, respectively.
Model
Non- Emb. Params
Tokens
Total
Prompi- level Acc. (strict)
Prompi- level Acc. (loose)
Instruction- level Acc. (loose)
MINITRON 4B-instruct
2.6B
90B
6,46
MINITRON 4B-
68.76
73.01
81.29
Phi-2
2.5B
1.4T
4.29
Qwen 1.5 Chat
1.2B
NIA
5,29
instruct
2B-IT
2B
6T
5.19
Gemma- IT
Qwen2-1
29
28.70
40.50
StableLM 2 Chat
LGB
2T
5.42
TiayLiama Chat
Instruct
Table 6: Evaluation results on IFEval.
Table 5: Evaluation results on MT-Bench.
Cost Savings for Training Model Family: the FLOPs required per training step $ for the 15B, 8B, and 4B models in the Nemotron-4 model family are, respectively: 4.4e17, 2.5e17 and 1.2e17. With the assumption that each model in the family is trained with an equivalent token count, steps and batch size, we obtain the following FLOP count for training each model in the family from scratch: (4.4e17 + 2.5e17 + 1 1.2 17 x steps. As noted from Tables and our approach requires 40x fewer training tokens for each additional model, hence resulting in the following updated FLOP count for the family: (4.4e17 + 2.5e17 /40 + 1. 1.2117440) X steps; the corresponding cost savings for training the full Nemotron-4 family using our approach is thus 1.8x.
We now dive deeper into our empirical ablations that help us arrive the list of best practices. Unless otherwise specified, we run these ablations on the Nemotron-4 15B checkpoint prior to continued training with the CT data blend.
4.2 Obtaining the Best Pruned Model
Best Aggregation Metric (Best Practice #2): we start by exploring the best aggregation metric for use with our activation based pruning criteria (see Section 2.2 for more details). Table [T shows how zero-shot LM loss and Wikitext2 perplexity [35] vary w.r.t different intra-batch and sequence aggregation functions. Here, the Nemotron-4 15B model is pruned to the Nemotron-3 8B architecture with no retraining. We notice that there is significant variation in zero-shot performance based on the aggregation metric, indicating the importance of selecting the right one. Both (batch-L2, seq=mean) and (mean, mean) perform well; in the remainder of the paper, use (12, mean) primarily due to its slightly better performance on the ST dataset. To further evaluate if these relative rankings hold after retraining, we perform related experiment: we prune the same 15B model to 8B using: (1) the best (L2, mean) metric, and (2) a poorly performing (L2, L2) metric, and perform retraining on both for 400 steps 1~BB tokens). The results of this experiment shown in Figure 5 From the Figure, we conclude that these rankings continue to hold post-retraining.
Iterative Importance (Best Practice #3): we evaluate whether iterative importance estimation provides any benefit (described in Section 2.2) and report results in Table 12. Here, we take the Nemotron-4 15B model and prune the embedding dimension alone using number of iterations T=1, 2, and iterations to the target value of 4096. We then perform lightweight retraining of all candidates
Assume batch size of 1152.
8

---

Model
Avg
Model
Avg.
MINITRON 4n-instruct
41.11
MINITRON 4B instruct
53.09
Gemma-28-IT
33,31
Gemma-2B-IT
41.63
Table 7: Evaluation results on ChatRAG- Bench.
Llama-3-B--nnsruuc
50,51
Table 8: Evaluation results on BFCL v2.
for 1.8B tokens. From the Table, we observe that while the iterative approach appears to be better before retraining, all 3 candidates converge to the same loss value, indicating no benefit.
Combining Depth and Width (Best Practice #4): we perform a simple experiment to compare the efficacy of width vs. depth pruning. Using the PPL and BI metries defined in Section 2.2 remove the 16 least important layers from the Nemotron 15B model based on both metrics to arrive at two variants of depth pruned models. We also perform neuron, head and embedding channel pruning to target the Nemotron-3 8B model and arrive at the width pruned variant. Finally, we combine depth (remove 4 least important layers) and width pruning arrive at the fourth variant. We report the results of this experiment in Table 13 We notice that even though the depth-width pruned variant has lower loss post-pruning, we the results flip around 200 steps of retraining (0.8B tokens); Table T and Figure 6 further illustrate this point.
4.3 Retraining and Search
Distillation vs. Conventional Training (Best Practice #5): in this experiment, compare: (1) training a 4B model with random initialization (4B-Random-Inic (2) pruning 15B to 4B, followed by retraining with conventional training (4B-Pruned), and (3) pruning 15B to 4B, and then retraining with distillation using the 15B model as the teacher (4B-Pruned-Distte Distilll. Since distillation adds training overheads (additional forward pass on the teacher model), compare approaches under iso-compute settings. Table [4] shows the results. Here, we observe a significant improvement in MMLU for (3), while both (1) and (2) score randomly. On HellaSwag, (3) (2) (1). This clearly demonstrates the superiority of distillation over conventional training after pruning.
Choice of Loss Function (Best Practice #5): experiment with Kullback-Leible divergence (KLD), MSE, cosine similarity and reverse KLD (R-KLD) to compute Liogite. Recent work has to be better fit than KLD in the SFT/instruction-followinn setting, and Agarwal et al. 2 claim the choice of loss is deskndepe...nt. We observe from Table 15 and 16 that
shown R- R-L [5 27
KLD is the best choice for pruned base model training.
Choice of Losses (Best Practices #6 and #7): typically, a weighted combination of LCLM and Loogits is used for distillation. We find that using Liogits alone results in the best performance as shown in Table [5 For Lis- Lemb + Lutt: + La I.o we make several observations similar to Lu et al. [31] these are listed in Appendix A.6 and in Table 17 Most notably, we observe improvements from using Lie when retraining models that don't prune the depth axis significantly, such MINITRON 8B and MINITRON 4B and hence use Liogits alone in such cases (see Table (8)
One-shot vs Iterative Pruning and Distillation Across Model Sizes (Best Practice #8): compress- ing Nemotron-4 15B to MINITRON 4B requires an aggressive 73.3% reduction of original model weights. We hypothesize that aggressive one-shot pruning loses out on important capabilities of the base LLM. We thus explore a simple iterative two-step pruning and retraining strategy where we first prune and retrain Nemotron-4 15B to create MINITRON 8B (~466% reduction) and further prune and retrain the latter to MINITRON 4B (~50% reduction). Table (last two rows) shows the comparison between single-shot and iterative pruning. and demonstrates that iterative achieves a 12% improvement in the MMLU scores compared to the one-shot strategy. During the final retraining step, we observe that using Nemotron-4 15B as the teacher achieves superior results compared to using MINITRON 8B. We provide additional ablations one-shot vs. iterative pruning in Appendix A.7
Search with Retraining (Best Practice #9): for lightweight neural architecture search, we use the search spaces defined in 9 MINITRON 8B and 4B. We further specify target parameter range of 8 and 4 billion parameters for the respective models, with a tolerance of 5%. With these settings, we obtain 15 and 18 feasible candidates for the 8B and 4B parameter targets, respectively. The architecture configurations for these candidates are provided in Table 19 As described in
Section 2.3 we perform lightweight retraining of all feasible candidates. Figure illustrates how

---

validation loss changes for the 8B candidates as training progresses. We notice that relative rankings undergo significant changes up 300 steps, and then stabilize.
Target
Layers
Heads
MLP Exp. Factor
Embedding
MINITRON SB
[29-32]
(32,48)
(2. 32.5,3,3.5,47
14096,4680, 5120 5632, 5632,6144) 6144) (2560,3072, 3584, 4096,
MINITRON 4B
[29-32]
(2,,32,48)
(2.5,3,3.5,4
Table 9: MINITRON 8B and 4B search space.
Single vs Multi-Phase Retraining (Best Practice #10): Recent studies 1
43 have shown improved results with multi-phase pretraining routines. Initially, models are trained on web data, followed by a lightweight phase with cleaner data. We explored two compression techniques: (I) prune the phase checkpoint, retrain with portions of phase and 2 data, and (2) prune the phase 2 checkpoint, retrain with a portion of phase 2 data. Table 20 shows that (2) is sufficient to regain accuracy and surpasses (1). This strategy is used for our best models, also suggesting that for further aligned models, it may suffice to prune the aligned model and retrain with a portion of the alignment dataset.
5 Related Work
Structured LLM Pruning:
there have been a number of recent structured pruning papers specifically
targeting LLMs; we can broadly classify these works into two main categories: (1) ones that prune only depth (layers), (2) ones that prune width (attention heads, MLP intermediate dimension, etc.) and/or depth. Recent work in the first category (depth pruning) includes ShortGPT (34. LaCo (55], and Shortened LLLMM [26]: for pruning layers in MINITRON models, we reuse and extend the metrics proposed some of these works (eg: block importance from ShortGPT (44 A number of recent papers have also proposed new saliency metrics and pruning strategies targeting width dimensions: namely, embedding channels, attention heads, and MLP intermediate channels [
53 33).
Most work in this category uses learnable masks, combined with an Augmented Lagrangian loss At LLM scale, this strategy has multiple
formulation to arrive at optimal width masks 4 53 3 disadvantages: (1) it requires compute and memory-intensive gradient computations, and (2) it requires a considerable amount of data and fine-tuming to arrive at reasonable masks. The notable exception in this line of work is Dery et al. (II). which recognizes these limitations and proposes saliency metrics that can be computed with only forward passes. To the best of our knowledge, we provide the first pruning strategy that (1) simultaneously targets both width and depth dimensions, (2) works at L.LM scale (i.e., uses only forward passes for computing importance and uses a small fraction of pretraining data), and (3) achieves state-of-the-ar compression and accuracy.
recent work has leveraged either teacher model which is larger/better (2] 27 or teacher-generatee synthetic data [6 36
Post-pruning A Accurayy Recovery:
337 to improve the accuracy of an existing trained smaller base model in the Supervised Fine Tuning (SFT)iinttructio following setting. Compared to recent width and depth pruning work [26] 34 53), to the best of our knowledge.
we are the first to employ distillation from an uncompressed teacher to improve the retraining of structurally -puuned student models.
6 Conclusions
This paper has presented a thorough empirical exploration of structured pruning and retraining in LLMs, offering unique insights into pruning order, effects of combining pruning axes, and retraining techniques for minimal data use. We have developed set of compression and retraining best practices, backed by extensive empirical evidence, which we employ to prune the Nemotron-4 15B model by factor of 2-4x. Our compressed MINITRON models are significantly cheaper to obtain compared to training each model from scratch (requiring up to 40x fewer training tokens), while still performing favorably to a number of similarly-sized community models; MINITRON models also outperform multiple state-of-the-art depth and width pruned models from the literature. Limitations: one notable limitation of our work is that we currently apply our proposed techniques only on the Nemotron family of models; we plan to address this by pruning other model families in future work. Also, even though it is short, our method requires full model retraining.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Ameya Sunil Mahabaleshwarkar, Hayley Ross, Brandon Rowlett, Oluwatobi Olabiyi, Ao Tang, and Yoshi Suhara for help with producing the instruction-tunee versions of MINITRON; additionally, James Shen for TRT-LLM support, and Sanjeev Satheesh, Oleksii Kuchaiev, Shengyang Sun, Jiaqi Zeng, Zhilin Wang, Yi Dong, Zihan Liu, Rajarshi Roy, Wei Ping, and Makesh Narsimhan Sreedhar for help with datasets. We'd d also like to gratefully acknowledge the insightful discussion and feedback from Chenhan Yu and Daniel Korzekwa.
References
[1] Marah Abdin, Sam Ade Jacohs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, Alon Benhaim, Misha Bilenko, Johan Bjorck, Sébastien Bubeck, Martin Cai, Caio César Teodoro Mendes, Weizhu Chen, Vishrav Chaudhary. Parul Chopra, Allie Del Giorno, Gustavo de Rosa, Matthew Dixon, Ronen Eldan, Dan Iter. Amit Garg. Abhishek Goswami, Suriya Gunasekar, Emman Haider. Junheng Hao, Russell J. Hewett, Jamie Huynh, Mojan Javaheripi, Xin Jin, Picro Kauffmann, Nikos Karampatziakis, Dongwoo Kim, Mahoud Khademi, Lev Kurilenko, James R. Loe, Yin Tat Lce, Yuanzhi Li, Chen Liang, Weishung Liu, Eric Lin, Zeqi Lin. Piyush Madan, Arindam Mitra, Hardik Modi, Anh Nguyen, Brandon Norick, Barun Patra, Danicl Perez- Becker, Thomas Portet, Reid Pryzant, Heyang Qin, Marko Radmilac, Corby Rosset, Sambudha Roy, Olatunji Ruwase, Olli Saarikivi, Amin Saied, Adil Salim, Michael Santacroce, Shital Shah, Ning Shang, Hiteshi Sharma, Xia Song. Masahiro Tanaka, Xin Wang. Rachel Ward, Guanhua Wang, Philipp Witte, Michael Wyatt, Can Xu, Jiahang Xu, Sonali Yadav, Fan Yang. Ziyi Yang. Donghan Yu, Chengruidong Zhang, Cyril Zhang, Jianwen Zhang. Li Lyna Zhang, Yi Zhang. Yue Zhang, Yunan Zhang, and Xiren Zhou. Phi-3 technical report: A highly capable language model locally on your phone, 2024.
[2] Rishabh Agarwal, Nino Vieillard, Yongehao Zhou, Piotr Stanczyk, Sabela Ramos Garca, Matthicu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations, 2024.
[3] Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
[4] Saleh Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler. and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. In The Twelfth International Conference on Learning Representations, 2023.
[5] Jimmy Lei Ba. Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 2XXiv:6607006555 2016.
[6] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Fric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2434 PMLR, 2023.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are fow-shot leamers. Advances neural information processing systems, 33:1877- 2020.
[8] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harrison Fdwards, Yura Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlow, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, David W. Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- William H. Guss, Alex Nichol, Igor Babuschkin. Suchir Balaji. Shantanu Jain, Andrew Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew M. Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. ArXiv, abs/2107.0333 2021.
[9] Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. A survey deep neural network compression: Challenges, overview, and solutions. IEFF Access, 6:39136-391.50 2018.
[10] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the A12 reasoning challenge. ArXiv, abs/ 803.05573, 2018.
[11] Lucio Dery, Steven Kolawole, Jean-Francois Kagey, Virginia Smith, Graham Neubig, and Ameet Tal- walkar. Everybody prune now: Structured pruning of Ilms with only forward passes. arXiv preprint arXiv:2402.0544 05406 2024.
11

---

[12] Yao Fu, How do language models put attention weights over long context? Yao Fi's Notion, Mar 2024. [13] Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao. An survey of neural network compression. arXiv preprint arXiv:2006.03665 2020.
[14] Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A. Roberts. The unreasonable ineffectiveness of the deeper layers. 2024.
[15] Yuxian Gu, Li Dong, Furu Wei, and Minlic Huang. Minillm: Knowledge distillation of large language models, 2024.
[16] Suriya Gunasekar, Yi Zhang. Jyoti Aneja, Caio Cesar, Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi. Mojan Javaheripi. Piero Kauffimann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang. Sébastien Bubeck. Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need, June 2023.
[17] Tahmid Hasan, Abhik Bhattacharjee, Md Saiful Islam, Kazi Samin, Yuan-Fang Li, Yong-Bin Kang, Sohel Rahman, and Rifat Shahriyar. XI-sum; Large-scale multilingual abstractive summarization for 44 languages, 2021.
[18] Yang He, Guoliang Kang, Xuanyi Dong, Yanwei Fu, and Yi Yang. Soft filter pruning for accelerating deep convolutional neural networks. arXiv preprint 180. 06866, 2018.
[19] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Mcasuring massive multitask language understanding, In International Conference on Learning Representattons, 2021.
[20] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean, Distilling the knowledge in a neural network, 2015.
[21] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks. arXiv preprint arXiv:2102.0055 005544 2021.
[22] Ari Holtzman, Jan Buys. Li Du Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. ArXiv, abs/1994.00975 2019.
[23] Edward Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2021.
[24] Shengding Hu, Yuge Tu, Xu Han. Chaoqun He. Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang. Chongyi Wang. Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng. Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategics, 2024.
[25] Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux. Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
[26] Bo- vo-Kyyong Kim, Geonmin Kim, Tac-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung- Kyu Song. Shortened LLaMA: A simple depth pruning for large language models. In ICLR 2024 Workshop on Mathematicul and Empirical Understanding of Foundation Models, 2024.
[27] Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se- Young Yun. Distilim: Towards streamlined distillation for large language models, 2024.
[28] Solomon Kullback and Richard A. Leibler. On information and sufficiency. Annals of Mathematical Statistics, 22(1):79-86, 1951.
[29] Stephanic Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Mcasuring how models mimic human falschoods, 2022.
[30] Zihan Liu. Wei Ping, Rajarshi Roy, Peng Xu, Chankyu Lee, Mohammad Shoeybi. and Bryan Catanzaro. Chatqa: Surpassing gpt-4 on conversational qa and rag. arXiv preprint N.NV::2201..022:: 2024.
[31] Chengqiang Lu, Jianwei Zhang, Yunfei Chu, Zhengyu Chen, Jingren Zhou, Fei Wu, Haiqing Chen, and
Hongxia Yang. Knowledge distillation of bransformer-based language models revisited, 2022.
[32] Jian-Hao Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. In Proceedings of the IEEE international conference on computer vision, pages 5058-5066
2017.
[33] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LLM--runer: On the Structural Pruning of Large Language Models. Advances in neural information processing systems, 36:21702-21720 2023.
[34] Xin Men, Mingyu Xu. Qingyu Zhang, Bingning Wang, Hongyu Lin. Yaojie Lu, Xianpei Han, and Weipeng Chen. ShortGPT: Layers in Large Language Models are More Redundant Than You Expect, 2024.
12

---

[35] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016.
[36]
Arindam Mitra, Luciano Del Corro, Shweti Mahajan, Andres Codas, Clarisse Simoes, Sahaj Agarwal, Xuxi Chen, Anastasia Razdaibiedina, Erik Jones, Kriti Aggarwal, Hamid Palangi, Guoqing Zheng, Corby Rosset, Hamed Khanpour, and Ahmed Awadallah. Orca 2: Teaching small language models how to reason, 2023.
[37] Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar. Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah. Orca: Progressive learning from complex explanation traces of gpt-4, 2023.
[38] Nvidia, i, Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H. Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, Sirshak Das, Ayush Dattagupta, Olivier Delalleau, Lcon Derczynski, Yi Dong, Daniel Egert, Ellie Evans, Aleksander Ficek, Denys Fridman, Shaona Ghosh, Boris Ginsburg, Igor Gitman, Tomasz Grzegorzek, Robert Hero, Jining Huang, Vibhu Jawa, Joseph Jennings, Aastha Jhunjhunwala, John Kamalu, Sadaf Khan, Oleksii Kuchaiev, Patrick LeGresley. Hui Li, Jiwei Liu, Zihan Liu, Fileen Long, Ameya Sunil Mahahaleshwarkar, Somshubra Majumdar, James Maki, Miguel Martinez, Maer Rodrigues de Melo, Ivan Moshkov, Deepak Narayanan, Sean Narenthiran, Jesus Navarro, Phong Nguyen, Osvald Nitski, Vahid Noroozi, Guruprasad Nutheti, Christopher Parisien, Jupinder Parmar. Mostofa Patwary, Krzysztof Pawelec, Wei Ping. Shrimai Prabhumoye, Rajarshi Roy. Trisha Saar, Vasanth Rao Naik Sabavat, Sanjeev Satheesh, Jane Polak Scowcroft, Jason Sewall, Pavel Shamis, Gerald Shen, Mohammad Shocybi, Dave Sizer, Misha Smelyanskiy, Felipe Soares, Makesh Narsimhan Sreedhar, Dan Su, Sandeep Subramanian, Shengyang Sun, Shubham Toshniwal, Hao Wang. Zhilin Wang, Jiaxuan You, Jiaqi Zeng, Jimmy Zhang, Jing Zhang, Vivienne Zhang, Yian Zhang, and Chen Zhu. Nemotron-4 340b technical report, 2024.
[39] NVIDIA. N 8B Model, 2023. Blog post: http:://ervealpprrvviddaaa.....a
with pott-pprocttion-rray--11111
[40]
OpenAT, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin. Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mo Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bermadett-Shaproo, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichacl, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Fcoffet, Atty Eleti, Tyna Floundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao. Elie Georges. Christian Gibson, Vik Goel, Tarun Gogineni. Gabriel Goh. Rapha Lontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu. Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschole, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Łukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lec, Jan Leike, Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanic Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini. Sam Manning. Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine Mcl Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reiichiro Nakano, Rajecv Nayak, Arvind Neelakantan, Richard Ngo. Hyeonwoo Noh, Long Ouyang, Cullen O' Keefe, Jakub Pachocki. Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish. Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng. Adam Perelman. Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokomy, Michelle Pokrass, Vitchyr Pong. Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rac, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders. Shibani Santurkar. Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher, Felipe Petroski Such, Natalie Summers, 1lya Sutskever, Jie Tang, Nikolas Tezak, Madeleine Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Preston Tuggle. Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang, Alvin Wang. Ben Wang, Jonathan Ward, Jason Wei, CJ Weinmann, Akila
13

---

Welihinda, Peter Welinder. Jiayi Weng, Lilian Weng. Mait Wiethoff, Dave Willner, Clemens Winter. Samuel Wolrich, Hannah Wong. Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo. Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang. Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. GPT-4 Technical Report, 2023.
[41] Seungeheol Park, Jachyeon Choi, Sojin Lee, and U Kang- A comprehensive survey of compression algorithms for language models. arXiv preprint arXiv:2401. 15334, 2024.
[42] Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Bo Liu, Aastha Jhunjhunwala, Zhilin Wang. Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Data, data everywhere: A guide for pretraining dataset construction, 2024.
[43] Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Mostofa Patwary, Sandeep Subramanian, Dan Su, Chen Zhu, Deepak Narayanan, Aastha Jhunjhunwala, Ayush Dattagupta, Vibhu Jawa, Jiwei Liu, Ameya Mahabaleshwarkar. Osvald Nitski, Annika Brundyn, James Maki, Miguel Martinez. Jiaxuan You, John Kamalu, Patrick LeGresley, Denys Fridman, Jared Casper, Ashwath Aithal, Oleksii Kuchaiev, Mohammad Shoeybi, Jonathan Cohen, and Bryan Catanzaro. Nemotron-4 15b technical report, 2024.
[44] Jupinder Parmar, Sanjev Satheesh, Mostofa Patwary, Mohammad Shoeybi. and Bryan Catanzaro. Reuse, don't retrain: A recipe for continued pretraining language models, 2024.
[45] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. WinoGrande: An adversarial winograd schema challenge at scale. Commun. ACM, 64(9), 2021.
[46] Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching Ilama2 performance with 0.1m dollars, 2024.
[47] Mohammad Shocybi. Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi billioiiinn parameter language models using model parallelism, 2020.
[48] Gemma Team, Thomas Mesnard, Cassidy Hardin, Rohert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Taccheti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Choquette-Choo, Clément Crepy, Daniel Cer. Daphne Ippolito, David Reid, Elena Buchatskaya, Eric Ni. Eric Noland, Geng Yan, George Tucker, George-Christe Christian Muraru, Grigory Rozhdestvenskly, Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean- Saan-Bapiiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katic Millican, Lars Lowe Sjocsund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez. Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross Mellroy, Ruibo Liu, Ryan Mullins, Samuel Smith, Sebastian Borgeaud, Sertan Girgin, Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De. Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojeiech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff' Dean, Koray Kavukcuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelie Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kencaly. Gemma: Open models based on gemini rescarch and technology, 2024.
[49] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molyhog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor. Adina Williams, Jian Xiang Kuan, Puxin Xu. Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama Open foundation and fine-tuned chat models. ArXiv, abx/2307. 009288, 2023.
[50] Wenhui Wang, Hangbo Bao, Shaohan Huang, Li Dong, and Furu Wei. MvniLMv2: Multi-head self- attention relation distillation for compressing pretrained transformers. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli, editors, Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 2 2140-2151, Online, August 2021. Association for Computational Linguistics.
[51] Wenxiao Wang, Wei Chen, Yicong Luo, Yongliu Long, Zhengkai Lin, Liye Zhang, Binbin Lin, Deng Cai, and Xiaofei He. Model compression and efficient inference for large language models: A survey. arXiv preprint 202.:2442000744455 2024.
14

---

[52] Jason Wei, Xuezhi Wang. Dale Schourmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le. Denny Zhou, al. Chain-of-thoughu prompting elicits reasoning in large language models. Advances Neural Information Processing Systems, 35:24824-24833 2022.
[53] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen, Sheared Ilama: Accelerating language model pre- pre-training via structured pruning. In The Twelfth International Conference Learning Representations, 2023.
[54] Fanjia Yan, Huanzhi Mao, Charlic Cheng-Jie Ji, Tianjun Zhang, Shishir G. Patil, lon Stoica, and Joseph E. Gonzalez. Berkeley function calling leaderboard. https://goriillacc..brr/ wtwp:/aaa......eeee...eee..eee.e catneaaaallnnnneeeeeeeeeee
2024.
[55] Yifei Yang. Zouying Cao,, and Hai Zhao. Laco: Large language model pruning via layer collapse. arXiv preprint arXiv:2402. 1887, 2024.
[56]
Rowan 7ellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Florence, Italy, July 2019. Association for Computational Linguistics.
[57] Lianmin Zheng, Wei-Lin Chiang. Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E Gonzalez, and Ion Stoica. Judging llm-as-a- judge with mt bench and chatbot arena. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 46595 46623, Curran Associates, Inc., 2023.
[58] Helong Zhou, Liangchen Song, Jiajie Chen, Ye Zhou, Guoli Wang, Junsong Yuan, and Qian Zhang, Re- thinking soft labels for knowledge distillation: bias-variance tradcoff perspective. CoRR, abs/2102.006500 2021.
[59] Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. instruction-folloviii evaluation for large language models, arXiv preprint 1rX19.2311.07911, 2023.
15

---

A Appendix
A.1 Pruned Architecture Details
Model
Layers
Hidden Size
Att. Heads
Query Groups
MLP Hidden
Parameters
Nemotron-4 15B
32
6144
48
24576
15.6B
Nemotron-3 8B
32
4096
32
32
16384
8.5B
MINITRON 8B
32
4096
48
16384
8.27B
MINITRON 4B
32
3072
24
9216
4.19B
Table 10: Architecture details of the uncompressed Nemotron and pruned MINITRON models. Vocabulary size 256k for all models.
A.2 Width Pruning
Best Aggregation Metric for Width Pruning:
Results post-pruning (zero-shot) shown in
Table 11 and after lightweight retraining in Figure 5
Batch
Sequence
8T L.M Loss
WikiText2 LM Loss
L2
L2
8.73
8.37
L2
mean
7.18
7.23
L2
8.18
8.61
mcan L2
8.41
7.84
mean mean
7.21
6.89
mean var
7.94
8.29
var L2
9.01
9.30
8.34
8.72
var var
10.55
11.14
Table 11: Zero-shot performance of activation-basee importance with different batch and sequence aggregation metrics. LM loss is reported the validation set of the 8T and WikiText2 datasets.
(L2, L2) vs (L2, Mean)
(1.12 [12 mar))
LM Validation Loss
100
150
200
250
300
350
400
Training Steps
Figure 5: LM validation loss curve for retraining of two pruned candidates with (L2, L2) and (L2, Mean) metrics for (batch, sequence) aggregation strategies.
Iterations
Initial (Zero-Shot) Validation Loss Final Validation Loss
T=1
5.43
1.92
T-2
5.55
1.92
T=4
5.24
1.92
Table 12: Comparison of one-shot importance estimation and pruning vs iterative importance estimation and pruning embedding dimension from the original size the target size. L.M validation loss is reported before and after lightweight retraining.
16

---

A.3
Depth vs. Width Pruning
Model
Parameters
LM Loss
MINITRON 8B Depth (PPL) [26
9,39B
2.155
MINITRON 8B Depth (BI) 34
9.39B
2.177
MINITRON 8B Width
7.74B
2.049
MINITRON 8B Depth Width
7.91B
2.062
Table 13:
Comparison retraining LM loss across different pruning strategies post retraining with 1.8B tokens. We explore depth only, width only, and combination of both. Width only strategy though with the least parameter count outperforms the rest.
Width alone vs Combined Width-Depth pruning
281-mlp-att-emb
32l-mpp-ttc-emb
LM Validation Loss
2.2
350
Training Steps
Figure 6: Comparison of retraining LM validation loss curves across two pruning choices, width alone combined depth and width. We observe a flip in the ranking prior 200 steps retraining. showcasing the need for a lightweight retraining phase.
A.4 Distillation vs. Conventional Training
Model
Tokens
Hellaswag
MMLU
4B R-ndRmn-nmmmi
150B*
46.22
24.36
4B Ai-RRandom-ine Init
400B
48.23
26.24
4B -r-Prr (prune 4 15B)
150B*
50.85
24.57
4B-Pruned-Disthe (prune Nemotron-4 15B)
100B*
51.04
37.81
4B-Pruned-Distie (prune MINITRON 8B)
100B'
52.04
42.45
Table 14: Accuracy comparison across different strategies to train 4B model. Pruning the 15B model and distillation results in gain of 4.8% Hellaswag and 13.5% MMLU compared training from scratch with equivalent compute. Pruning an 8B model instead of 15B model results in additional gain of 1 % and 4.6% the benchmarks. Indicates settings with iso-compute.
A.5 Retraining with Distillation
Choice of loss function: In our experiments with the previous generation of Nemotron models in Table 15] we see that KLD consistently outperforms R-KLD, cosine and MSE. WSL-KD (58) also performs inferior to KLD. Hence, we do not repeat all these studies with the experiment setup in Section rather only a subset as shown in Table 16
Temperature:
We experiment with =0.1, 0.5, 1.0, 3.0 in the softmax computation. Literature shows vision (classification) models output a spikey logit distribution and softening logit distribution with temperature results in an improvement when using distillation. However, LLM logit distributions have higher entropy and hence the inspiration for temperature to reduce the noise. We observe best results when 7=1.0.
17

---

Loss
LM loss
WikiText PPL
LCLM + Llogita (MSE)
2.144
9.007
LCLM + Llogits (RKL.)
(RKLD)
2.140
9.008
LOLM + Liogita (Cosiie(Cosine)
(Cosine)
2.134
8,965
Loss Function
LM
LOLM + Liogits (KLD)
2.117
8.791
Liogit (RKLD)
2.665
LKLg)).(KLD)
2.107
8.720
(KLD) 2.155
Table 15: LM loss comparison for various loss functions and loss components on Nemotron-3 8B.
Table 16: Comparison of loss functions with MINITRON 2B-Depth-prunce LM loss is reported on the validation sets of the 8T.
The loss component Llogits alone with forward KLD loss outperforms the rest.
Inspired by the top-K/top-P sampling approach used in LLM inference, experimented with retaining only top-K teacher and the corresponding student logits prior to computing Llogits- This should essentially remove noise from the low probability logits/tokens. We observe that low value of top-K (<=100) results in a significant drop in accuracy. The drop is no longer observed when increasing top-K, but better than not using top-K. Hence, skip using top-K for further experiments.
Top-K:
A.6 Choice of Losses
1. Using loss L, based the output activations of encoder block provides boost.
2. The final 1-2 layers in a Transformer for LLM are highly specialized [12] and mapping hidden states across (last-2):(last22) layers for both the student and teacher achieves the best result [31].
3. Using word embeddings based loss Lemb) improves accuracy.
4. Computing loss Lass(atenntin relation loss [50) based on query, key and value states does not show any improvement.
5. Adding loss L, based on the input to MLP makes no difference.
6. We weren't it able to experiment with attention scores due Flash Attention abstractions.
7. Cosine similarity loss performs the best.
Results are shown in Table
A.7 One-shot vs. Iterative Pruning and Distillation
Loss components
LM loss
Liogus
2.155
Liogas Lo(29:)3)
2.145
Liogits Lo(15:15)) Lemb
2.240
Liogits Lo(23:15)+ Lemb
2,205
Liogits Ioo(9:155) Lemb
2.203
Llogit Lo(30:15) Lemb
2.188
Liogia Lo(31:15) Lemb
2.180
Ltogita La(28:12)+ Lemb
2.141
Togits + I.o(29:13) + T.emb
2.141
Togits I.o(29:14) Tremb
2.152
Rogits Lo(30:14)+ Lemb
2.150
Liogita (.9:133) Lemb + L;929:13)
2.141
Table 17: Ablation study on loss components for computing Lis and different (teacher:student) layer mapping for Lo and L. LM loss is reported on the validation set of the 8T. Note: Layer indices start from 0, teacher Nemotron-4 15B layers (0-31), student MINITRON 8B-Depthhprundh layers (0-15).
One-shot Iterative for Importance Estimation and Pruning: refer to Table 12
One-shot vs Iterative within Dimension: understand the best prune-retrain strategy considering single dimension that can be pruned across the model (depth), experiment with two different
18

---

Loss
Tokens
MMLU
HellaSwag
HvmanEval
Llogits Lis
18.98
58.0
73.6
26.8
Llogis
18.9B
58.3
73.8
26.2
Liogits
94B
62.8
79.7
30.4
Table 18: Ablation study for MINITRON 8B with and without the loss component Lis, and increased
retraining token count with Ligits Adding Lis performs par with using Llogits alone.
approaches for depth pruning and retraining in order to arrive at the MINITRON SB-Depth-prunce model mentioned above.
As first step, we rank layer importance with the procedure mentioned in 2.2 borrowed from 33 Then we:
1. Iteratively prune and distill: Remove the least important layer. distill using 1.8B tokens and repeat the procedure 16 times. See iterative 167 Figure
2. One-shot prune and distill: Remove 16 least important layers, distill using 1.8B x 16(30.2B) tokens. See shot pruning 167 in Figure 7
In order to mitigate the sharp drop in accuracy and to prevent further catastrophic collapse of the model, we increase the compute budget from 1.8B x 1.8B tokens starting with pruning the 26 layer model which amounts to 86.4B tokens in total. See iterative x 4 16/ in Figure 7 Increasing the training budget improves accuracy, but still performs worse than the one-shot prune and distill strategy that uses 30.2B tokens.
With the iterative strategy, we can see in Figure 7 accuracy on:
Hellaswag and PIQA is retained up to 31 layers and start dropping gradually with further removal of layers. We MMLU score retained up to 26 layers and start dropping gradually with further removal of layers. We see sharp drop when the model is reduced to 20 layers.
sharper drop when the model is reduced to 25 layers.
This shows that a few layers can be removed from a pretrained model in a lossless manner with minimal retraining, As for the retraining strategy, it is best to follow the one-shot method. Our results agree with [26].
0.7
hellaswag baseline 15B (BT takens)
mmiu baseline 15B tokens)
0.6
humaneval baseline 15B (8T tokens)
Reeaaive x1 161 (+30 28 tokens)
1-shet pruning 16l (+30.20 takens)
iterative (+:30
iteratvve (+86. tokens)
1-shat pruning 16l (+30.20 tokens)
0.3
Reeaaive (+-30.26 takens)
iterative x4 161 (+886, 8B tokens)
0.2
1-shat pruning (+30.28 takens)
32
30
28
26
24
22
20
18
16
number of layers
Figure 7: Accuracy on MMLU, HellaSwag and HumanEval benchmarks for iterative vs one-shot depth pruning and retraining strategy. One shot pruning and retraining outperforms the iterative approach.
One-shot vs Iterative across Dimensions: we experiment with iterative EMB-+MLP-ATI pruning with retraining after both iterations and one-shot EMB-MLP-ATT pruning and retraining with equivalent token count as the former. As shown Figure shot achieves better results than the iterative approach.
19

---

One-shot mlp-att-emb vs Iterative emb-2mlp-att Pruning
One Shot ripattimmm
frmmire enb-irnl-47
LM Validation Loss
2.3
2.2
2.1
2.0
1.9
100
150
200
250
300
350
400
Training Steps
Figure 8: Comparison of LM validation loss curves for one-shot pruning of embeddings, MLP, attention and retraining for 400 steps vs the iterative approach; pruning the embeddings first and retraining for 250 steps, followed by pruning MLP, attention and retraining for additional 150 steps.
A.8 Search
All the feasible 8B candidates produced by search are shown in Table 19
ID
Layers
Heads
MLP Exp. Factor
Embedding
32
32
12800
5120
32
32
13824
4608
32
48
11520
4608
32
48
16384
4096
31
32
12800
5120
31
32
16128
4608
31
48
13824
4608
31
48
16384
4096
30
32
12800
5120
10
30
32
16128
4608
11
30
48
13824
4608
30
48
16384
4096
29
32
12800
5120
14
29
32
16128
4608
15
29
48
13824
4608
Table 19: MINITRON 8B feasible candidates produced by search.
A.9
Single vs. Multi-Phase Training
Table 20 compares the accuracy of single vs. multi-phase training.
Strategy
Tokens
MMLU
HellaSwag
PIQA
HumanEval
Phasel + Phase2
113B
54.7
80.3
77.2
25.6
Phase2 only
94B
61.9
80.1
76.7
30.5
Table 20: Accuracy comparison single vs multi-phase training approach with MINITRON 8B- Width-pruned. Note: This is not the searched 8B model in Table
A.10 Compute Resources
All experiments were performed on 16x NVIDIA DGX A.00 nodes (8x A.00 80GB) for short tumarruun times.
20

---

Recraining of Search candidates 121
M15
2.40
2.5
Validation
2.20
2:00
Training Scaps
Figure 9: Retraining of searched candidates for 8B target with 1.8B training tokens.
21

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
Justification: Experimental results provided in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims
made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: Limitations listed in Section Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: No theoretical proofs or results. Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
22

---

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide relevant information for reproducing our results in Section 4 and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibilityy the case of closed-source models, may be that access to the model limited in some way (c.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Model weights are public. We intend to release the code pending internal review.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc/
public/guides: Sobllis//uuides//ddeeebbmsssssss (ow.eSuaiiiiioooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not
23

---

reproducible, they including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.cc. cc.s..c/pp wwtpublic///cccccccccccccccmcllt wwwic//guiieseee Sutee/ooeeSuuiicccce ss (ooeeSubmsssioooooo.......
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why. At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Provided in Section 4 Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Training LLMs is extremely expensive, we are therefore unable to perform multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma eтror bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Detailed in Section A.10
24

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes)
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Generic algorithm for optimizing LLM training. Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: No such release artifacts. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring
25

---

that users adhere to usage guidelines or restrictions to access the model or implementing
safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors
should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Assets are properly credited. Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, ppapesswithhoee com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Does not involve crowdsourcing or research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
26

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Does not involve crowdsourcing or research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
27


Reviews:
[{"review_id": "iDcWcUyPtK", "reviewer": "Reviewer_qFD8", "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.", "questions": "see weakness", "limitations": "na", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper conducts extensive experiments, comparing the latest baselines, and the authors summarize extensive tricks for pruning the model.\n\nThe pruned 8B model performs well, surpassing the model trained from scratch.\n\nThe 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.\n\nWhen 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their encouraging comments and insightful feedback. Please find our responses below:\n\n> **Will the authors open-source the code? If the code and data are open-sourced, I would raise my score.**\n\n**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.\n\n**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.  \n\nModels 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.\n\n> **When pruning a 70B model to 7B, would the method in the paper still work?**\n\nWe 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. \n\nPruning 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. \n\n**References:**\n1. Parmar, Jupinder, et al. \"Data, Data Everywhere: A Guide for Pretraining Dataset Construction.\" arXiv preprint arXiv:2407.06380 (2024).\n2. Parmar, Jupinder, et al. \"Reuse, Don't Retrain: A Recipe for Continued Pretraining of Language Models.\" arXiv preprint arXiv:2407.07263 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OoqiDhazdQ", "reviewer": "Reviewer_enpQ", "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.", "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. \n\n2. 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.\n\n3. Check weakness 2", "limitations": "The authors have mentioned limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is very well written and provide excellent explanations for all of their modeling choices.\n\n2. 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.\n\n3. I also enjoyed reading the comprehensive details that are written in the appendix.\n\nOverall 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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for pointing out the strengths of the paper and providing valuable feedback. \n\n> **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 ...**\n\nThe 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:\n\n| Minitron  4B  | Tied | Other Model (Gemma-2 / Phi-2) |\n| -------- | ------- | ------- |\n| 43.43%  | 25.71%    | 30.86% |\n\nHere, 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.\n\nAdditionally, 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.\n\n> **In Table 2 is the nemotron-8B baseline also trained on the 94B token retraining dataset?**\n\nIn Table 2, the nemotron-8B baseline is not trained on the exact 94B retraining dataset.\n\n> **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.**\n\nWe 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. \n\nGoing 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:\n\n|    | Minitron 8B | Nemotron 8B |\n| -------- | ------- | ------- |\n|  Tokens |  **94.4B** | 3.5T |\n|  MMLU |   **0.521** | 0.485 |\n|  PIQA |   0.794 |  **0.7971** |\n|  HellaSwag |  **0.763** | 0.7587 |\n|  HumanEval |  **0.207** | 0.189 |\n\nMinitron 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.\n\nAs 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.\n\nIn the Figure:\n1. Gray: Train nemotron-8B from scratch on the retraining dataset\n2. Pink: Minitron 8B \n3. Orange: Further vanilla training of nemotron-8B on the retraining dataset\n4. Blue: Further vanilla training on minitron 8B on the retraining dataset\n\n2 being significantly better than 1, and 4 being better than 3 showcases the efficacy of our technique. \n\nThe 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.\n\nWith 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.\n\n> **The figures and tables need much much more detailed captions**\n\nWe 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):\n\n* 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.\n\n* Figure 5: LM validation loss curve for retraining of two pruned candidates with (L2, L2) and (L2, Mean) metrics for (batch, sequence) aggregation strategies.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X58oyoNpUm", "reviewer": "Reviewer_W3pm", "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.", "questions": "N/A", "limitations": "Suggestions:\nThe number of significant figures should remain consistent in Table 2 and Table 3.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written, and the best practices are easy to follow, which is useful in practical LMs.\n2. 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].\n2. 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.\n\n\n[1] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning. 2024. In ICLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their encouraging comments and insightful feedback. Please find our responses below:\n\n>  **... it is better to highlight the difference between the final choice in this paper and the approaches in previous work like Sheared LLaMa.**\n\nTo help with this comparison, we have created the following table to highlight the differences between our approach and Sheared LLaMa [1]:\n\n| Criteria    | Sheared LLaMa |  Minitron | Summary |\n| -------- | ------- | ------- | ------- |\n|  Pruning importance estimation | Learnable mask on  embedding, MLP and attention heads. Mask learning uses 0.4B tokens and is 5x slower than standard LM training [1]. | Forward pass on a tiny number of samples (1024 in paper) to compute importance for embedding, MLP and attention heads at once. This adds negligible overhead. | An expensive gradient-based importance estimation strategy (Sheared-LLaMa) is not required to achieve strong accuracy for pruned models. A simpler, forward pass only approach (Minitron) works well. |\n| Pruning dimensions | Depth, embedding, MLP, attention | Depth, embedding, MLP, attention | Both approaches support multiple width and depth axes. |\n| Retraining | Uses 50B tokens with conventional finetuning | Uses 94B tokens with knowledge distillation | Both papers find that retraining is required to recover accuracy loss on benchmarks. We showcase the superiority of knowledge distillation over conventional training and recommend the former as the retraining approach. |\n| Search | No search process. Matches a previously-defined fixed architecture configuration. | Ability to search for optimal architectures using a constrained brute-force algorithm. We also observe that lightweight retraining of searched candidates is essential to stabilize relative rankings. | Minitron can find better architectures through search. For eg: search finds that keeping attention heads intact is better for Minitron-8B. The paper also provides new insights on how zero-shot accuracy post-pruning isn’t reflective of final accuracy. |\n| Multiple compressed models | Requires repeating the 5x slower mask learning process N times for producing N compressed models. Each of the N models must also be finetuned. | Single importance estimation pass (negligible overhead) is sufficient for all N compressed models. Each of the N models must be distilled. | Minitron approach is significantly less costly when multiple compression targets are specified. |\n\n> **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.**\n\nImportance 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.\n\nRegarding 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.\n\nWe 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:\n1. `4B-Random-Init` is a 4B model trained from scratch using conventional ground truth CE loss.\n2. `4B-Pruned (prune Nemotron-4 15B)` is a 4B model pruned from a 15B model, using conventional ground truth CE loss for retraining.\n3. `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.\n\n| Model    | Tokens |  Hellaswag | MMLU |\n| -------- | ------- | ------- | ------- |\n| 4B-Random-Init  | 150B*    | 46.22 | 24.36 |\n| 4B-Random-Init |  400B     | 48.23 | 26.24 | \n| 4B-Pruned (Prune Nemotron-4 15B)   | 150B*  | 50.85 | 24.57 |\n| 4B-Pruned-Distill (Prune Nemotron-4 15B)   | 100B*   | 51.04 | 37.81 |\n| **4B-Pruned-Distill (Prune Minitron 8B)**   | **100B***    | **52.04** | **42.45** |\n\nWe 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. \n\n> **The number of significant figures should remain consistent in Table 2 and Table 3.**\n\nThank you for pointing this out, we will fix this in the final version.\n\n**References:**\n1. Xia, Mengzhou, et al. \"Sheared LLaMa: Accelerating language model pre-training via structured pruning.\" arXiv preprint arXiv:2310.06694 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their insightful feedback and comments. We have posted individual responses for each review, making every effort to provide additional results to support our responses within this limited rebuttal period. We hope our rebuttal addresses all reviewer concerns.

Meta-review: The work explores combing model pruning and knowledge distillation for LLMs. The reviewers increased the scores after discussion and agreed to accept the submission. The engineering findings would be useful for practitioners. The first pruning stage seems critical in the whole pipeline. However, it would restrict applying the proposed method for cross-architecture distillation, i.e., the teacher model and student model use different model architectures. Overall, this is a good submission with a lot of details.
Decision: Accept (poster)