Knowledge Distillation for LLMs: How to Train Smaller Models from Big Teachers
Running a massive Large Language Model is expensive. We are talking about models with hundreds of billions of parameters that demand clusters of GPUs and burn through electricity like crazy. For most businesses, deploying a giant model directly into production isn't just costly-it's often impossible due to latency constraints and hardware limits. This is where Knowledge Distillation comes in.
Knowledge Distillation (KD) is a model compression technique where a smaller 'student' model learns to mimic the behavior of a larger, more powerful 'teacher' model. Instead of training the student on raw data alone, you train it to copy the teacher’s output probability distributions. The result? A compact model that runs fast and cheap but still delivers performance close to its heavyweight parent.
This guide breaks down how knowledge distillation works for LLMs, the different types of knowledge you can transfer, and practical steps to implement it using tools like NVIDIA NeMo or Hugging Face Transformers.
How Knowledge Distillation Works for LLMs
To understand KD, imagine a senior engineer (the teacher) mentoring a junior developer (the student). The junior doesn’t just memorize the final code; they learn *how* the senior thinks-the reasoning, the edge cases considered, and the alternatives rejected.
In traditional machine learning, a model is trained to predict the single correct label (e.g., "cat" vs. "dog"). In knowledge distillation, the teacher model outputs a probability distribution over all possible tokens. For example, when asked to complete the sentence "The sky is...", the teacher might assign 80% probability to "blue," 15% to "clear," and 5% to "bright." These low-probability options contain valuable information about the relationship between words.
The student model is then trained to match this soft distribution, not just the top choice. This process uses a loss function that combines two parts:
- Distillation Loss: Measures how well the student matches the teacher’s probabilities (often using Kullback-Leibler divergence).
- Task Loss: Measures how well the student predicts the ground-truth labels.
A key parameter here is temperature. By dividing the logits by a temperature value (usually T=2 to T=4), we soften the probability distribution, making the differences between high and low probabilities less extreme. This allows the student to learn the relative rankings of incorrect answers, which helps generalize better.
Types of Knowledge You Can Transfer
Not all distillation is created equal. Depending on what you want your student to learn, you can transfer different layers of knowledge. Here are the main categories identified in recent research (including the 2024 survey on LLM Distillation):
- Logit-Level Knowledge: The most common form. The student matches the teacher’s token-level probability distributions at each step. This captures the "dark knowledge"-the subtle relationships between tokens that aren’t obvious from the correct answer alone.
- Sequence-Level Knowledge (Data Distillation): Here, the teacher generates full responses (like summaries or code snippets) on a dataset. The student is then trained on these generated sequences as if they were ground truth. This is essentially synthetic data generation.
- Preference-Level Knowledge: Used for alignment. If the teacher has been fine-tuned with human feedback (RLHF), you can distill those preferences into the student so it behaves safely and helpfully without needing its own expensive alignment phase.
- Intermediate-Representation Knowledge: The student tries to match the internal hidden states or attention maps of the teacher. This forces the student to process information similarly to the teacher, though it requires architectural compatibility.
Practical Implementation: From 8B to 4B Parameters
Let’s look at a concrete example. Suppose you have a Meta Llama-3.1-8B model (the teacher) and want to create a faster 4B parameter version (the student). You can achieve this by combining pruning and distillation.
Step 1: Pruning the Student Architecture First, you reduce the size of the model structure. Using tools like NVIDIA NeMo, you can perform depth pruning (removing transformer layers) or width pruning (reducing hidden dimensions). For instance, dropping half the layers from a 32-layer model cuts the parameter count roughly in half. However, pruning alone causes significant accuracy drops.
Step 2: Distilling the Knowledge Next, you run the distillation process. The pruned 4B model acts as the student. You feed the same dataset through both the 8B teacher and the 4B student. The training loop minimizes the difference between their outputs. According to NVIDIA’s tutorials, this often involves running scripts like `megatron_gpt_distillation.py` on multiple GPUs to handle the compute load.
Step 3: Balancing the Losses You need to tune the weight given to the distillation loss versus the standard task loss. If you rely too much on the teacher, the student might inherit its biases or errors. If you rely too little, you lose the benefits of compression. A common starting point is a balanced mix, adjusted based on validation performance.
| Technique | How It Works | Impact on Accuracy | Best For |
|---|---|---|---|
| Knowledge Distillation | Trains a smaller model to mimic a larger one’s outputs. | Low loss if done correctly. | Reducing parameter count while keeping intelligence. |
| Quantization | Reduces numerical precision of weights (e.g., FP16 to INT8). | Minimal loss with modern methods. | Reducing memory footprint and speeding up inference. |
| Pruning | Removes unnecessary neurons or connections. | Moderate to high loss without recovery training. | Simplifying model architecture. |
Challenges and Limitations
While KD is powerful, it’s not free. The biggest hurdle is computational cost. Proper distillation requires running the teacher model forward for every token in the training set. Since the teacher is large, this doubles the training time compared to standard training. Some teams use approximations, like sampling only the top 256 tokens from the teacher’s distribution instead of the full vocabulary, to save memory and compute.
Another issue is bias propagation. If your teacher model has toxic or biased tendencies, the student will likely learn them too. You cannot expect the student to be smarter or safer than the teacher. In fact, if the teacher fails on certain domains, the student will fail there too, sometimes even worse because it lacks the capacity to recover.
Finally, there is the capacity gap. If you try to squeeze a 70-billion-parameter model into a 1-billion-parameter student, no amount of distillation will make it work. The student simply doesn’t have enough "brain space" to hold the knowledge. You need to find a sweet spot where the student is small enough to be efficient but large enough to capture the teacher’s core capabilities.
When to Use Knowledge Distillation
You should consider KD if:
- You need to deploy an LLM on edge devices (phones, IoT) with limited RAM.
- You want to reduce inference costs by switching from a proprietary API (like GPT-4) to an open-source model hosted in-house.
- You have a specialized domain (legal, medical) where a general-purpose model is too slow, and you’ve already fine-tuned a large expert model.
If you just need a quick speedup and don’t care about retaining complex reasoning skills, simple quantization might be easier. But if you need to maintain high-quality generation while cutting parameters, KD is the gold standard.
What is the difference between Knowledge Distillation and Quantization?
Knowledge Distillation reduces the number of parameters in a model by training a smaller network to mimic a larger one. Quantization reduces the precision of the numbers used to represent weights (e.g., from 16-bit floats to 8-bit integers) without changing the model structure. They are often used together: first distill to get a smaller model, then quantize to make it lighter and faster.
Can I distill from a proprietary model like GPT-4?
Yes, but with limitations. You can use the API to generate responses or probability distributions (if available) to train an open-source student model. However, you must comply with the provider’s terms of service regarding data usage and copyright. This approach is popular for creating private, domain-specific assistants that behave like GPT-4 but run locally.
How much compute does Knowledge Distillation require?
It is computationally intensive. Because you need to run the large teacher model during training, the cost is roughly double that of training the student alone. For very large teachers, this can be prohibitive. Techniques like sampled soft labels or offline data generation can help mitigate these costs.
What is the best student-to-teacher size ratio?
There is no fixed rule, but generally, the student should be within the same order of magnitude. Compressing a 70B model to a 7B student is challenging but possible with advanced techniques. Compressing 70B to 1B usually results in severe quality loss. A 2x to 4x reduction in parameters is a safe starting point for experimentation.
Does Knowledge Distillation improve the model’s safety?
Not inherently. The student learns the teacher’s behaviors, including any biases or unsafe responses. To improve safety, you need to include preference-level distillation or additional alignment steps (like RLHF or DPO) specifically targeting safety criteria after the initial distillation.
