Compressing Large Language Models (LLMs)

🌈 Abstract

The article discusses the challenges of working with large language models (LLMs) and how model compression techniques can help overcome these challenges. It covers three broad categories of compression techniques: Quantization, Pruning, and Knowledge Distillation. The article also provides a hands-on example of compressing a 100M parameter model using knowledge distillation and quantization.

🙋 Q&A

[01] Overview of Model Compression Techniques

1. What are the three broad categories of model compression techniques discussed in the article?

• Quantization: Lowering the precision of model parameters to reduce model size
• Pruning: Removing model components that have little impact on performance
• Knowledge Distillation: Transferring knowledge from a larger teacher model to a smaller student model

2. How do these compression techniques help address the challenges of working with large LLMs?

• Reduced computational requirements and model size, enabling wider accessibility of powerful ML models, lower-cost integration of AI into consumer products, and on-device inference for improved user privacy and security.
• The techniques can be combined for maximum compression, allowing for significant reductions in model size without sacrificing performance.

[02] Quantization

1. What is the difference between Post-training Quantization (PTQ) and Quantization-Aware Training (QAT)?

• PTQ compresses the model by replacing parameters with a lower-precision data type, which is a relatively easy way to reduce model costs but can lead to performance degradation.
• QAT trains models from scratch with lower-precision data types, which can lead to significantly smaller, well-performing models, but is more technically demanding.

2. What is Quantization-aware Fine-tuning, and how does it compare to PTQ and QAT?

• Quantization-aware Fine-tuning is an approach between PTQ and QAT, consisting of additional training of a pre-trained model after quantization.
• It aims to overcome the limitations of PTQ while being less technically demanding than QAT.

[03] Pruning

1. What is the difference between Unstructured and Structured Pruning?

• Unstructured pruning removes unimportant weights from the neural network, but the resulting sparse matrix operations require specialized hardware to be efficient.
• Structured pruning removes entire structures from the neural network (e.g., attention heads, neurons, layers), avoiding the sparse matrix operation problem.

2. How do structured pruning approaches identify which structures to remove?

• Structured pruning approaches seek to remove structures with the smallest impact on model performance, using various techniques to identify the least important structures.

[04] Knowledge Distillation

1. How does knowledge distillation work?

• Knowledge distillation transfers knowledge from a larger teacher model to a smaller student model, either by using the teacher's output logits to train the student or by learning from synthetic data generated from the teacher model.

2. What is the example of knowledge distillation discussed in the article?

• The article provides an example of using knowledge distillation to compress a 100M parameter model into a 50M parameter model, and then further compressing it using 4-bit quantization to a 7x smaller final model.

[05] Hands-on Example

1. What are the key steps in the example implementation?

• Load the dataset and teacher model
• Define a student model architecture with fewer layers and attention heads
• Tokenize the dataset and define an evaluation function
• Train the student model using a distillation loss function that combines the teacher's soft targets and the ground truth hard targets
• Further compress the student model using 4-bit quantization

2. What are the performance results of the compressed models compared to the original teacher model?

• The student model outperformed the teacher model on both the test and validation sets, demonstrating that the compression techniques can lead to improved performance.
• The final 4-bit quantized model maintained the performance improvements while being 7x smaller than the original teacher model.