Source-linked AI summary
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
TL;DR
Training separate LLM sizes from scratch is costly, so this paper studies obtaining smaller models by pruning and retraining one large model with limited data. Its structured-pruning and distillation practices produce 2-4× smaller MINITRON models that require up to 40× fewer training tokens while performing favorably against similarly sized models and outperforming several pruned alternatives.
Problem
Training multiple multi-billion-parameter LLMs from scratch is extremely time-, data-, and resource-intensive, motivating the question of whether pruning and retraining one large model can produce smaller, accurate models with little additional data.
Method
The paper empirically evaluates structured pruning across neurons, attention heads, embedding channels, and depth, then develops compression best practices using pruning, retraining, and knowledge distillation.
Results
The resulting MINITRON models compress Nemotron-4 15B by 2-4×, use up to 40× fewer training tokens than training from scratch, perform favorably against similarly sized community models, and outperform multiple pruned models.
Takeaways & Limitations
Training one large model and obtaining smaller models through pruning plus retraining can be more accurate and substantially more compute-efficient than training each model from scratch.
Takeaways & Limitations
The reported empirical conclusions include an assumption of batch size 1152.
Abstract
from arXiv · showhide
Large language models (LLMs) targeting different deployment scales and sizes are currently produced by training each variant from scratch; this is extremely compute-intensive. In this paper, we investigate if pruning an existing LLM and then re-training it with a fraction (<3%) of the original training data can be a suitable alternative to repeated, full retraining. To this end, we develop a set of practical and effective compression best practices for LLMs that combine depth, width, attention and MLP pruning with knowledge distillation-based retraining; we arrive at these best practices through a detailed empirical exploration of pruning strategies for each axis, methods to combine axes, distillation strategies, and search techniques for arriving at optimal compressed architectures. We use this guide to compress the Nemotron-4 family of LLMs by a factor of 2-4x, and compare their performance to similarly-sized models on a variety of language modeling tasks. Deriving 8B and 4B models from an already pretrained 15B model using our approach requires up to 40x fewer training tokens per model compared to training from scratch; this results in compute cost savings of 1.8x for training the full model family (15B, 8B, and 4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to training from scratch, perform comparably to other community models such as Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art compression techniques from the literature. We have open-sourced Minitron model weights on Huggingface, with corresponding supplementary material including example code available on GitHub.
1 Introduction
Training separate LLM sizes from scratch is costly, motivating a study of pruning and retraining one large model into smaller variants. The paper empirically develops compression practices and applies them to produce smaller MINITRON models with favorable accuracy and lower training costs.
- Training multiple multi-billion-parameter models from scratch is extremely time-, data-, and resource-intensive.
- The paper asks whether pruning and retraining one large model can produce smaller, more accurate models using only a small fraction of the original training data.
- The study explores structured pruning and retraining across neurons, attention heads, embedding channels, and model depth.
- The authors develop practical compression and retraining best practices from empirical evidence on pruning metrics, hyperparameters, axis combinations, and data-efficient retraining.
- MINITRON models derived from Nemotron-4 15B require up to 40× fewer training tokens than training from scratch while comparing favorably with similarly sized models.
2 Pruning Methodology
The methodology ranks components across depth, width, attention, and embedding axes using activation- and sensitivity-based importance, then trims weights, searches feasible architectures, and supports iterative compression with distillation. It combines axis-specific pruning strategies with retraining-oriented design choices, including residual information reuse during attention-head pruning.
- Obtaining a pruned model: The overall process iteratively ranks and trims components, distills the pruned model, and uses the distilled model as the source for the next compression iteration.Distillation combines embedding-output, logit, and transformer-encoder-specific losses mapped between student and teacher blocks.
- Importance analysis: Activation-based importance estimates sensitivity across depth, neurons, heads, and embedding channels using forward passes on a 1024-sample calibration dataset.The strategy avoids gradient computation, which the paper describes as prohibitively memory- and compute-intensive for modern LLMs.
- Importance analysis: Aggregation functions for activation scores are evaluated across batch and sequence dimensions because simple summation is not always optimal.The evaluated functions include mean(abs), L2 norm, and variance, with layer-wise scores summed into network-wide axis scores.
- Importance analysis: Depth pruning ranks layers using perplexity or Block Importance, with Block Importance computable in one forward pass and extendable to contiguous layer groups.Block Importance uses cosine distance between layer inputs and outputs to estimate sensitivity.
- Obtaining a pruned model: For a target architecture, ranked elements are trimmed directly from MLP, MHA, embedding, and LayerNorm weight matrices, while architecture search enumerates feasible configurations within a parameter budget.The search varies layer count, attention heads, MLP dimensions, and embedding dimensions, using commonly used dimensions and a narrow target parameter range.
- Obtaining a pruned model: Attention-head pruning adds residual information from removed heads back into remaining heads to preserve relevant knowledge and improve accuracy.For grouped-query attention, the strategy is applied only to query heads.
3 Retraining
Retraining recovers accuracy after pruning by combining conventional ground-truth training with knowledge distillation from the uncompressed teacher. The paper evaluates output, hidden-state, embedding, and other intermediate-state losses, finding benefits from selected components and dynamic loss weighting.
- Retraining means recovering accuracy after pruning through conventional training or knowledge distillation from the uncompressed teacher.The pruned model is treated as the student and the uncompressed model as the teacher.
- Knowledge distillation transfers teacher information by matching student outputs and/or intermediate Transformer states.The paper explores combinations of intermediate states and mappings as distillation loss components.
- A shared linear transformation maps student hidden states to the teacher dimension before comparing post-LayerNorm states across selected intermediate layers.The loss is computed over chosen hidden states and sequence positions.
- Dynamic weighting of the intermediate-state loss performs better than using a constant coefficient.The total loss combines student cross-entropy, logit loss, and α-weighted intermediate-state loss.
4 Experiments and Results
The experiments evaluate pruning and retraining Nemotron-4 models, establish practical compression best practices, and compare MINITRON models with community and pruned baselines. Results show strong performance with substantially lower training cost, while ablations identify effective distillation, pruning, and search choices.
- Experimental Setup: The study compresses Nemotron-4 15B into 8B and 4B models and evaluates them on language modeling, reasoning, coding, and summarization tasks.The evaluation includes MMLU, HumanEval, Arc-C, HellaSwag, TruthfulQA, WinoGrande, and XL-Sum.
- Best Practices: The best practices recommend iterative prune-and-distill family construction, width-focused pruning, KLD distillation, and search choices tailored to pruning axes and model depth.The recommendations also distinguish importance estimation for width and depth axes and specify when intermediate-state and embedding losses help.
- Main Results: MINITRON 8B uses 40× fewer tokens than Nemotron-3 8B while outperforming Nemotron-3 8B and LLaMA-2 7B and matching Mistral 7B, Gemma 7B, and LLaMA-3 8B.These comparisons use the reported baseline evaluations in Table 2.
- Main Results: MINITRON 4B retains capabilities across tasks, outperforms Gemma2, and surpasses multiple depth- and width-pruned models.The paper contrasts this broad performance with smaller specialized models that excel only on selected tasks.
- Main Results: MINITRON 4B-instruct beats similarly sized models across instruction-following, roleplay, and retrieval-augmented generation tasks, and outperforms Gemma-2B-IT and Llama-3-8B-instruct on function calling.These results concern the instruction-tuned variant rather than the base model.
- Retraining and Search: Distillation improves pruned-model retraining: the distilled 4B model significantly improves MMLU, while HellaSwag ranks distilled above conventionally retrained and randomly initialized models.The comparison is performed under iso-compute settings despite the teacher’s additional forward pass.
- Retraining and Search: Iterative pruning improves MMLU by 12% over one-shot pruning when compressing Nemotron-4 15B to MINITRON 4B.The iterative route first produces MINITRON 8B and then prunes it to 4B; the 4B reduction is 73.3% of the original weights.
5 Related Work
Prior structured LLM pruning work mainly targets depth alone or width and depth together. This paper distinguishes its approach by using distillation from an uncompressed teacher during retraining of structurally pruned models.
- Structured LLM Pruning: Structured LLM pruning studies broadly divide into depth-only methods and methods pruning width, depth, or both.Examples of depth-pruning work include ShortGPT, LaCo, and Shortened LLaMA.
- Post-pruning Accuracy Recovery: The paper identifies its use of an uncompressed teacher for retraining structurally pruned models as distinct from prior width- and depth-pruning work.Related post-pruning recovery studies use larger teachers or teacher-generated synthetic data mainly for instruction-following settings.
6 Conclusions
The paper presents empirically supported pruning and retraining practices for producing smaller MINITRON models from Nemotron-4 15B. These models require far fewer training tokens while remaining competitive with similarly sized models and prior pruning methods.
- Conclusions: The method compresses Nemotron-4 15B by 2-4× using structured pruning and retraining best practices developed through extensive empirical exploration.The exploration covers pruning order, combinations of pruning axes, and data-efficient retraining.
- Conclusions: MINITRON models require up to 40× fewer training tokens than training each model from scratch while performing favorably against similarly sized community models.They also outperform multiple state-of-the-art depth- and width-pruned models from the literature.
A.1 Width Pruning
The width-pruning analysis compares activation-based importance aggregation choices before and after lightweight retraining. It also examines whether importance estimation and pruning should be performed in one shot or iteratively.
- Aggregation Metrics: Zero-shot width-pruning results compare activation-based importance metrics across batch and sequence aggregation functions using validation loss on 8T and WikiText2.The aggregation study is reported in Table 13.
- Retraining Curves: Figure 5 plots validation-loss curves during retraining for two pruned candidates using (L2, L2) and (L2, Mean) aggregation metrics.The figure focuses on retraining behavior rather than only post-pruning performance.
- One-shot versus Iterative Pruning: Table 14 compares one-shot and iterative importance estimation and pruning for reducing the embedding dimension to a target size.Validation loss is reported before and after lightweight retraining.
A.2 Depth vs. Width Pruning
Figure 6 compares validation-loss retraining curves for width-only pruning and combined depth-and-width pruning. Their ranking reverses before 200 retraining steps, so early curves do not determine the final ordering.
- Figure 6 compares LM validation-loss curves for width-only versus combined depth-and-width pruning.
- Before 200 retraining steps, the ranking between the two pruning choices flips.
A.3 Retraining with Distillation
The retraining experiments favor forward KLD, a temperature of τ=1.0, and avoiding top-K logit filtering. Validation-loss curves also show that pruning choices can change rank during early retraining.
- Forward KLD consistently outperforms R-KLD, cosine, MSE, and WSL-KD in the reported Nemotron experiments.
- Width-only and combined depth-and-width pruning reverse their validation-loss ranking before 200 retraining steps.
- Table 15 reports that the Llogits component with forward KLD outperforms the other evaluated loss configurations.
- τ=1.0 produces the best distillation results among the tested temperatures.
- Top-K values ≤100 significantly reduce accuracy, while larger values provide no benefit over using all logits.
A.4 Choice of Losses
The paper evaluates pruning and retraining choices across depth, dimensions, loss components, and architecture search. Its reported comparisons favor one-shot pruning, activation-based importance estimation, and targeted knowledge-preserving mechanisms.
- A.5 One-shot vs. Iterative Pruning and Distillation: One-shot pruning and distillation outperform iterative pruning despite the iterative strategy receiving up to 86.4B retraining tokens.
- A.5 One-shot vs. Iterative Pruning and Distillation: Iterative pruning retains HellaSwag and PIQA through 31 layers and MMLU through 26 layers before accuracy begins to decline.
- Architecture Search: The search enumerates feasible architectures within a parameter budget using common neuron, head, and embedding dimensions.
- Training Strategy: The reported experiments compare single-phase and multi-phase training, including a width-pruned MINITRON 8B model distinct from the searched 8B model.