Compact Language Models via Pruning and Knowledge Distillation

Published 11/5/2024 by Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, Pavlo Molchanov

Overview

Compact Language Models via Pruning and Knowledge Distillation is a research paper that explores methods for compressing large language models while maintaining their performance.
The key ideas include pruning model parameters and knowledge distillation, which transfer knowledge from a larger "teacher" model to a smaller "student" model.
The researchers tested their techniques on popular language models like BERT and GPT-2, achieving significant size reductions with minimal accuracy loss.

Minitron compression significantly reduces training costs and improves results.

1/4

Original caption: Figure 1: Results for Minitron. Compression results in significant reduction of training costs for additional models (40×40\times40 ×) while producing better results.

Iterative pruning and distillation trains smaller LLMs.

Original caption: 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.

Student model distillation from a teacher model, minimizing combined losses.

Original caption: Figure 4: Overview of Distillation. A s⁢t⁢u⁢d⁢e⁢n⁢t𝑠𝑡𝑢𝑑𝑒𝑛𝑡studentitalic_s italic_t italic_u italic_d italic_e italic_n italic_t model with N𝑁Nitalic_N layers is distilled from a t⁢e⁢a⁢c⁢h⁢e⁢r𝑡𝑒𝑎𝑐ℎ𝑒𝑟teacheritalic_t italic_e italic_a italic_c italic_h italic_e italic_r model with M𝑀Mitalic_M layers. The s⁢t⁢u⁢d⁢e⁢n⁢t𝑠𝑡𝑢𝑑𝑒𝑛𝑡studentitalic_s italic_t italic_u italic_d italic_e italic_n italic_t learns by minimizing a combination of embedding output loss, logit loss and transformer encoder specific losses mapped across s⁢t⁢u⁢d⁢e⁢n⁢t𝑠𝑡𝑢𝑑𝑒𝑛𝑡studentitalic_s italic_t italic_u italic_d italic_e italic_n italic_t block S𝑆Sitalic_S and t⁢e⁢a⁢c⁢h⁢e⁢r𝑡𝑒𝑎𝑐ℎ𝑒𝑟teacheritalic_t italic_e italic_a italic_c italic_h italic_e italic_r block T𝑇Titalic_T.

Pruning strategies' performance on a large language model before and after retraining.

1/2

DEP	MLP	ATT	EMB	Distillation Loss	LM Validation Loss
Yes	Yes	Yes	Yes	5.35 → 0.38	2.062
No	Yes	Yes	Yes	6.33 → 0.37	2.049
No	No	Yes	Yes	5.07 → 0.42	2.101
No	No	No	Yes	8.35 → 0.49	2.155
Train from Scratch (Random Initialization)					12.27 → 2.34
						3.953

Original caption: Table 1: Demonstration of how various pruning strategies perform before and after lightweight retraining using ∼similar-to\sim∼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.

		Models
		Benchmark	Metric	Llama-3	Llama-2	Mistral	Gemma	Nemotron-4	Nemotron-3	Minitron
		# Parameters		8B	6.7B	7.3B	8.5B	15.6B	8.5B	8.3B
		# Non-Emb. Params		5.9B	6.4B	7B	7.7B	12.5B	6.4B	6.2B
		# Training Tokens		>15T	2T	8T	6T	8T	3.8T	94B
Knowledge, Logic	winogrande (5)	acc	77.6	74	78.5	78	83.6	75.9	79.0
Knowledge, Logic	arc_challenge (25)	acc_norm	57.8	53	60.3	61	58.8	52.8	52.6
	MMLU(5)	acc	65.3	46	64.1	64	66.6	54.7	63.8
	hellaswag(10)	acc_norm	82.1	79	83.2	82	84.6	78.5	80.7
	gsm8k(5)	acc	50.3	14	37	50	48.5	24.0	51.3
	truthfulqa(0)	mc2	43.9	39	42.6	45	40.7	36.5	42.6
Coding	XLSum en (20)(3)	rougeL	30.9	31	4.80	17	32	30.9	31.2
Coding	MBPP(0)	pass@1	42.4	20	38.8	39	38	27.04	35.2
	humaneval (n=20)(0)	pass@1	28.1	12	28.7	32	35.4	20.7	31.6

Original caption: Table 2: Performance of our pruned Minitron 8B model compared to multiple baselines: the original Nemotron-4 15B, the previous generation Nemotron-3 8B, and multiple community models. Minitron 8B uses 40×\times× fewer tokens than Nemotron-3 8B. All evaluations run by us, except for entries marked with *, which we report from the corresponding papers.

Plain English Explanation

Large language models like BERT and GPT-2 have achieved impressive performance on various natural language tasks. However, these models can be very large, requiring substantial computational resources to run. This makes them challenging to deploy on resource-constrained devices like smartphones or edge computing systems.

The researchers in this paper explored two main techniques to compress these large models:

Pruning: This involves selectively removing model parameters (the numerical values that define the model's behavior) that are deemed less important. By carefully pruning away parts of the model, it can be made significantly smaller without losing too much accuracy.
Knowledge Distillation: This involves training a smaller "student" model to mimic the behavior of a larger "teacher" model. The student model learns to approximate the outputs of the teacher model, allowing it to achieve similar performance in a more compact form.

By combining these techniques, the researchers were able to greatly reduce the size of popular language models like BERT and GPT-2 while preserving a large portion of their original capabilities. This could enable these powerful models to be deployed on a wider range of hardware, from powerful servers to resource-constrained edge devices.

Technical Explanation

The researchers first explored pruning techniques to remove less important model parameters. They experimented with various pruning methods, such as magnitude-based pruning, which removes parameters with small absolute values, and iterative pruning, which prunes parameters in multiple rounds.

To further compress the models, the researchers then applied knowledge distillation. This involves training a smaller "student" model to mimic the behavior of a larger "teacher" model. The student model learns to predict the same outputs as the teacher model, allowing it to achieve similar performance in a more compact form.

The researchers tested their techniques on popular language models like BERT and GPT-2. They were able to achieve significant size reductions, such as compressing BERT from 110 million parameters to just 13 million parameters, while maintaining a large portion of the original model's accuracy.

Critical Analysis

The researchers thoroughly explored the trade-offs between model size and performance, providing valuable insights for practitioners looking to deploy large language models in resource-constrained environments. However, the paper does not address potential issues that could arise from aggressive pruning or knowledge distillation, such as potential loss of rare or important information, or the impact on downstream tasks beyond the ones tested.

Additionally, the researchers only evaluated their techniques on a limited set of language models and tasks. It would be valuable to see how these methods perform on a wider range of models and applications, including more specialized or domain-specific language models.

Conclusion

This research demonstrates that it is possible to significantly reduce the size of large language models through a combination of pruning and knowledge distillation, without sacrificing too much of their original capabilities. These techniques could enable the deployment of powerful natural language processing models on a wider range of hardware, from powerful servers to edge devices. As AI systems become more ubiquitous, efficient model compression will be an increasingly important area of research.

Full paper

Loading PDF viewer...

Read original: arXiv:2407.14679

This summary was produced with help from an AI and may contain inaccuracies - check out the links to read the original source documents!

Follow @aimodelsfyi on 𝕏 →