Knowledge Distillation › Distillation for LLMs

Scaling distillation to large language models introduces both new opportunities and new challenges. The core objective remains the same, but the scale, the nature of the outputs, and the relationship between teacher and student all change significantly.

Distilling a 70B-parameter model into a 7B model involves training the student while repeatedly querying the teacher for its output distribution over a vocabulary of tens of thousands of tokens. The teacher must be kept in memory alongside the student during training, which roughly doubles the GPU memory requirement. For very large teachers (100B+ parameters), this becomes practically infeasible unless the teacher is served on separate hardware and queried via API, with logits streamed back to the student's training process.

An alternative is offline distillation: pre-compute the teacher's soft labels on the entire training corpus and store them as training data. The student then trains on these pre-computed distributions without needing the teacher loaded during training. This trades memory for storage. Storing soft labels over a vocabulary of 50,000 tokens for billions of training tokens requires enormous disk space.

In practice, the teacher's distribution is often truncated to store only the top-$k$ probabilities (e.g., $k = 64$), discarding near-zero mass. This significantly reduces storage requirements while retaining the most informative part of the distribution. The bottom $49,936$ probabilities in a 50,000-token vocabulary are typically negligible, and storing them wastes space without providing meaningful signal.

There are two levels at which distillation can operate for language models.

Token-level distillation applies the standard distillation objective at every token position in the sequence. At position $t$, the student is trained to match the teacher's distribution over the next token given the prefix. This is the most direct application of the Hinton et al. framework and is used in models like DistilBERT and TinyBERT. It provides the richest training signal because the student receives a full distribution at every step, but it requires access to the teacher's logits, which may not be available for closed-source models.

Sequence-level distillation treats the problem differently. Rather than matching the teacher's token-by-token distribution, the student is trained on text generated by the teacher. The teacher generates sequences by sampling from its distribution (possibly with temperature), and the student uses those generated sequences as training data for standard language modeling. This approach, studied by Kim and Rush (2016), avoids the need to store or stream the teacher's logits and works well when high-quality teacher-generated data is available.

The choice between these approaches involves real tradeoffs. Token-level distillation provides a richer training signal but requires accessing the teacher's internal distributions. Sequence-level distillation can be performed from any model that produces text, making it applicable even when distilling from closed-source models like GPT-4. However, sequence-level distillation loses the detailed probability information that temperature-scaled soft labels provide. The student learns what the teacher says, but not how uncertain the teacher was when it said it.

For practical LLM compression, token-level distillation typically produces better results when the teacher's logits are accessible. Sequence-level distillation is the only option when they are not, and in those cases it still provides substantial benefit over training without teacher guidance.

DistilBERT, introduced by Sanh et al. (2019) at Hugging Face, demonstrates token-level distillation at scale. The goal was to compress BERT-base (110M parameters, 12 layers) into a model with 40% fewer parameters (66M parameters, 6 layers) while retaining 97% of BERT's performance on the GLUE benchmark.

The distillation approach combined three components:

$$\mathcal{L}_{\text{DistilBERT}} = \alpha \cdot \tau^2 \cdot \mathcal{L}_{\text{distill}} + \beta \cdot \mathcal{L}_{\text{CE}} + \gamma \cdot \mathcal{L}_{\text{cos}}$$

where:

• $\mathcal{L}_{\text{distill}}$: the standard distillation loss matching the teacher's masked language model predictions
• $\mathcal{L}_{\text{CE}}$: supervised cross-entropy loss against ground-truth masked tokens
• $\mathcal{L}_{\text{cos}}$: cosine embedding loss between teacher and student hidden state vectors, encouraging alignment of internal representations

The third term, $\mathcal{L}_{\text{cos}}$, represents a form of feature distillation: rather than only matching the teacher's output probabilities, the student is also trained to produce similar hidden representations layer by layer. We will explore this further in the next chapter on distillation variants. For now, note that it provides an additional training signal that helps narrow the representation gap between teacher and student beyond what output-level supervision alone can achieve.

DistilBERT's training also used weight initialization from the teacher: rather than initializing the 6-layer student randomly, the authors initialized its layers by taking every other layer from the 12-layer BERT teacher. This dramatically speeds up convergence because the student starts with weights already in a good region of parameter space. The distillation training then refines these weights to account for the fact that the student has fewer layers between them.

The reported results were substantial. DistilBERT achieved 97% of BERT's performance on GLUE, ran 60% faster at inference, and reduced memory consumption by 40%. It became one of the most widely used compressed language models and demonstrated that careful distillation can retain nearly all of a model's capability at a substantial cost reduction.

DistilBERT's success established several practical lessons for LLM distillation:

• Weight initialization from the teacher substantially accelerates convergence and should be used whenever architectures are compatible.
• The distillation objective is more effective than standard training alone even when the same data is used, meaning the soft labels add information beyond what the raw text provides.
• A 40% parameter reduction is achievable with minimal performance loss; larger compressions require more sophisticated techniques.
• The cosine representation loss is a useful complement to the output distillation loss, helping align internal representations not just final outputs.

TinyBERT, introduced by Jiao et al. (2020), pushed the compression ratio further by distilling both intermediate representations and the attention matrices, in addition to the final output. The student is trained to match the teacher's attention patterns at corresponding layers in addition to matching the output distribution.

This multi-layer distillation approach requires the student to have the same number of attention heads as the teacher (or a compatible mapping between them), but it transfers much richer information. The teacher's attention patterns encode how it distributes focus across the input sequence when processing each token, and matching these patterns constrains the student to develop similar internal processing strategies.

TinyBERT achieved competitive performance with only 14.5M parameters (about 13% of BERT-base's 110M) on GLUE tasks. The tradeoff is complexity: the distillation process requires careful layer mapping between teacher and student, and the multi-objective loss is harder to tune. But the compression ratio is dramatically better than DistilBERT, showing that when you extract richer signal from the teacher, you can get away with a much smaller student.

A different flavor of distillation has become prominent with instruction-following LLMs. Models like Alpaca, Vicuna, and Orca train a smaller model to mimic the instruction-following behavior of a larger model like GPT-4. The "teacher" in these cases is GPT-4, but its internal logits are not accessible. The distillation is purely sequence-level: GPT-4 generates responses to prompts, and the student is fine-tuned to produce similar responses on the same prompts.

This approach raises an important distinction between two uses of distillation:

• Capability distillation: transferring the teacher's factual knowledge, reasoning ability, and language model quality to the student
• Behavioral distillation: teaching the student to follow instructions, adhere to formats, and exhibit particular conversational styles

Both are valuable, but they require different teacher signals. Capability distillation benefits from token-level soft labels with temperature scaling, because the goal is to transfer the precise probabilistic structure of the teacher's language model. Behavioral distillation can work with generated sequences because the target behavior is the output format and conversational style, not the precise token probability distribution.

Alpaca and Vicuna demonstrated that even relatively modest fine-tuning on high-quality teacher-generated responses could substantially improve a smaller model's instruction-following capabilities. A LLaMA-7B model fine-tuned on 52,000 GPT-4 completions behaves far more helpfully and coherently than the base model. This is sequence-level distillation working as intended: the student learns the distribution of helpful responses by imitating a larger model's outputs.

Orca (Microsoft, 2023) took this further by using "explanation traces" from GPT-4 rather than final responses alone. The teacher was prompted to explain its reasoning step by step before giving an answer, and the student was trained on both the reasoning trace and the final answer. This richer supervision signal helped the student learn what to say and how to reason, resulting in substantially better performance on reasoning benchmarks than models trained on final answers alone.

A significant practical challenge in LLM distillation is distribution shift: the student generates different text than the teacher would, and the student must learn on text from the training corpus or generated by the teacher, rather than on its own outputs. This mismatch creates a situation where the student encounters teacher logits in contexts it would never reach when generating on its own.

To understand why this matters, consider what happens at inference time. The student generates token by token, and at each step its generation depends on all previous tokens it has produced. During training, the student has been trained on contexts from the training corpus or from the teacher's completions. If the student's own completions drift away from these contexts (which they do, because the student is different from both the corpus and the teacher), the student encounters contexts that are "off the training distribution." Its behavior in these off-distribution contexts is unpredictable.

On-policy distillation addresses this by having the student generate sequences during training and computing the teacher's logit distribution at each position of the student's own outputs. The student is trained to improve its distribution on the contexts it generates, rather than on fixed training data contexts. This is more expensive, requiring forward passes through both models on student-generated text, but it produces better alignment between the student's behavior and the teacher's guidance, especially for generation-heavy tasks.

On-policy distillation is related to techniques like RLHF (reinforcement learning from human feedback) where the model is also trained on its own outputs. The key difference is the reward signal: RLHF uses a scalar reward from a reward model, while on-policy distillation uses the teacher's full distribution as a dense, token-level reward. The dense signal is much more informative than a scalar, because it tells the student whether the output was good or bad and what a better output would have looked like at each token position.

When storing teacher distributions for offline distillation, the standard approach is to truncate to the top-$k$ logits. This raises an important question: how large should $k$ be?

If $k$ is too small, the stored distribution loses dark knowledge. With $k = 10$, only the ten most probable tokens are represented, and the relational structure among less likely tokens is discarded. With $k = 64$, a much richer slice of the distribution is preserved. With $k = 1000$, the storage cost grows but the distribution is nearly complete for most tokens.

In practice, $k = 64$ is a common choice that balances storage cost against information quality. For a 50,000-token vocabulary, the top 64 tokens typically account for more than 99% of the probability mass for most positions in natural language sequences. The remaining 49,936 tokens split a negligible fraction of probability between them. Their relative ordering is information the student could learn, but it is so faint that it rarely makes a measurable difference.

The choice of $k$ should also be matched to the temperature. At high temperature, more probability mass is spread across non-top tokens, making a larger $k$ more important. At low temperature, the distribution is sharp and a small $k$ captures almost everything.