Lesson Plan: Knowledge Distillation for Large Language Models

Duration: ~15 minutes Audience: Practitioners and students familiar with machine learning basics Goal: Understand what distillation is, why it matters, and how it works in the context of LLMs


1. Motivation — Why Distillation? (3 min)

Large language models like GPT-4 or Claude can have hundreds of billions of parameters. Running them is expensive: they require specialized hardware, consume significant energy, and introduce latency that makes real-time applications difficult. Yet smaller models trained from scratch on the same data often can't match their quality.

Distillation bridges this gap. The core idea is simple: instead of training a small model to predict the right answers from raw data alone, you train it to mimic a larger, more capable model. The large model has already learned rich representations of language — patterns, nuances, and reasoning strategies — and distillation transfers that knowledge into a compact form.

Discussion prompt: Where might you need a smaller, faster model in practice? (Edge devices, latency-sensitive APIs, cost-constrained deployments.)

2. The Core Mechanism — Teacher and Student (4 min)

Distillation involves two models:

Hard labels vs. soft labels

In standard training, a model learns from "hard" labels — the single correct answer (e.g., the next token in a sequence). But the teacher model produces a full probability distribution over all possible tokens. These "soft" labels carry much more information. For example, if the teacher predicts the next word after "The cat sat on the" as mat (60%), floor (25%), rug (10%), and chair (5%), those relative probabilities encode the teacher's understanding that these words are all plausible but in a particular order. A hard label would only say "mat" and throw away all that nuance.

Temperature scaling

To make the soft labels even more informative, both teacher and student outputs are passed through a softmax function with an elevated temperature parameter. Higher temperature "softens" the distribution further, spreading probability mass more evenly and revealing the teacher's beliefs about less-likely options. During inference, the student uses normal temperature.

The loss function

The student is trained with a combined loss:

A weighting parameter (often called α) balances these two objectives. In practice, the distillation loss often dominates early in training, while the task loss helps the student stay grounded in correctness.

3. Distillation Strategies for LLMs (4 min)

The original distillation framework (Hinton et al., 2015) was designed for classification. Applying it to LLMs introduces additional considerations:

Output-level distillation is the most common approach. The student learns to match the teacher's next-token probability distributions across large text corpora. This is conceptually straightforward but requires running the teacher over the entire training set to generate soft targets, which can be computationally expensive.

Feature-level distillation goes deeper. Instead of only matching final outputs, the student also tries to match the teacher's internal representations — hidden states or attention patterns at intermediate layers. This can be especially helpful when the teacher and student have very different architectures, as it gives the student richer learning signals.

Data-augmented distillation takes a different angle. The teacher generates synthetic training data — reasoning chains, question-answer pairs, or paraphrases — and the student trains on this augmented dataset. This is how many open-source LLMs are trained: they learn from text generated by more capable proprietary models. It's simpler to implement because you only need the teacher's text outputs, not its internal probabilities.

Reinforcement learning from AI feedback (RLAIF) is a related technique where the teacher acts as a reward model or evaluator, scoring the student's outputs and providing training signal through reinforcement learning rather than direct imitation.

4. Trade-offs and Practical Considerations (2 min)

Distillation isn't free. Some important considerations:

Capacity gap. If the student is too small relative to the teacher, it may not be able to absorb the teacher's knowledge effectively. Research suggests that progressive distillation — using a chain of increasingly smaller models — can help.

Distribution mismatch. The teacher's training data and the distillation data should overlap meaningfully. A student distilled on out-of-domain text may not capture the teacher's strengths.

Evaluation. Distilled models tend to perform well on the same kinds of tasks the teacher excels at, but may be more brittle on out-of-distribution inputs. Thorough evaluation across diverse benchmarks is essential.

Legal and policy questions. When the teacher is a proprietary model, using its outputs to train another model may raise intellectual property concerns. Several model providers explicitly restrict this in their terms of service.

5. Recap and Key Takeaways (2 min)

Further reading: Hinton, Vinyals & Dean, "Distilling the Knowledge in a Neural Network" (2015); Sanh et al., "DistilBERT" (2019); Gu et al., "MiniLLM" (2024).