Knowledge Distillation for LLMs

RLVR trained a policy against checked outcomes where a verifier exists. Knowledge distillation asks a different deployment question: after you have a useful large language model (LLM) teacher, how do you transfer selected behavior into a smaller model that's cheaper to serve?

Knowledge distillation trains a smaller student model to imitate useful behavior from a larger or otherwise more capable teacher. Imagine a return-policy chatbot: a large model performs well on an evaluated question set, but running it for every user costs too much. A smaller student is viable only if it retains enough answer quality under a latency and serving-cost budget.

The transfer signal depends on teacher access. In white-box settings, the student can match softened token probabilities or internal features. In black-box settings, it can fine-tune on selected teacher-written answers, solution traces, or synthetic corpora. Each channel carries different information and different failure modes. The goal isn't to assume the student inherits the teacher; it's to transfer useful behavior and verify the resulting quality-cost tradeoff.

The flow below shows the white-box logit KD case. Response distillation uses selected teacher text instead of the teacher-probability branch.

Why soft labels teach more than hard labels

In machine learning, the core idea (introduced by Hinton et al. in 2015) is to train a student model to mimic a teacher model's behavior, rather than only the ground truth labels. The student learns from the teacher's full probability distribution, which contains richer information than simple one-hot labels:^[1]

$\mathcal{L} = \alpha \cdot \mathcal{L}_{\text{distill}} + (1 - \alpha) \cdot \mathcal{L}_{\text{task}}$

Where $\mathcal{L}_{\text{distill}}$ matches teacher behavior, $\mathcal{L}_{\text{task}}$ matches ground-truth labels, and $\alpha \in [0, 1]$ controls the tradeoff. This combined loss makes sure the student learns from both the teacher's soft predictions and the true labels.

A concrete example

Suppose a customer asks, "How long do I have to return an electronic item?" The teacher's output distribution contains dark knowledge: relationships between classes that hard labels don't capture.

The hard label only says "30 days" is correct. In this simplified answer-choice example, the soft label also records that the teacher ranks "14 days" above "1 year." That extra signal is useful only if the teacher's ranking is itself useful; it is not proof that the student learned a general return-policy rule.

The probabilities above describe one simplified prediction decision. For a causal LLM, logit distillation applies this idea at each predicted token position. Temperature exposes more of the teacher's ranking, but it can also flatten the distribution until very little preference signal remains.

When distillation beats training from scratch

Distillation is most useful when you already have a strong teacher, legal access to its signal, and a clear smaller deployment target. It does not win automatically over training from scratch. The Gemma 2 report gives a controlled example: its authors train the 2B and 9B models with token-probability distillation, and a 2B ablation trained for 500B tokens scores 67.7 when distilled from a 7B teacher versus 60.3 from scratch on their three-benchmark average.^[2]

Different recipes expose different supervision channels: Orca trains on explanation traces, phi-1.5 uses curated synthetic textbook-like data, and Gemma 2 uses teacher token probabilities for small models.^[3][4][2] They motivate careful data and signal selection; they do not establish one universally best recipe.

Matching the teacher's probabilities: logit distillation

This method requires access to the teacher model's logits (the raw, unnormalized scores output by the final layer of the network before the softmax function). We minimize the KL Divergence (Kullback-Leibler divergence), a mathematical measure of how one probability distribution differs from another, to align the student's probability distribution with the teacher's.

First, we apply temperature scaling (dividing logits by a temperature $T > 1$ before softmax to soften the probability distribution) to both models' logits to get softened probability distributions:

$$q_i = \frac{\exp(z_{t,i} / T)}{\sum_j \exp(z_{t,j} / T)}, \quad p_i = \frac{\exp(z_{s,i} / T)}{\sum_j \exp(z_{s,j} / T)}$$

Where $z_t$ are the teacher's logits, $z_s$ are the student's logits, and $T$ is the temperature. Then we compute the KL divergence loss:

$\mathcal{L}_{\text{KL}} = T^2 \sum_i q_i \log \frac{q_i}{p_i}$

Reading the formula

• $q_i$ is the teacher's softened probability for token $i$ (the target we want the student to match)
• $p_i$ is the student's softened probability for token $i$ (what the student currently predicts)
• $T$ is the temperature (higher $T$ creates a softer, more uniform distribution)
• The $T^2$ scaling factor compensates for gradient scaling: as temperature increases, gradients from soft targets scale down by approximately $1/T^2$. Multiplying by $T^2$ keeps the relative contribution from soft targets roughly stable as you tune temperature. Think of it like turning up the volume on a quiet signal so it can compete with the loud one.
• KL divergence measures how much information is lost when using the student's distribution $p$ to approximate the teacher's distribution $q$

For causal LMs, two implementation details matter. First, next-token training requires a one-token shift: logits at position $t$ train against the label at position $t+1$. Second, direct logit KD assumes teacher and student use the same token-to-id output mapping. Equal vocabulary sizes alone are insufficient: token id 42 must denote the same token for both models. If output spaces differ, plain token-level KL no longer lines up cleanly and you usually fall back to response distillation or design an explicit mapping. The snippet below handles the shift, masks ignored positions, and fails fast on a reordered vocabulary.

Common mistake: Running logit distillation without verifying tokenizer alignment. Two models can have the same vocabulary size and different token-id mappings. Compare the complete output mapping, not only vocab_size, before training.

When you only have text: response distillation

When you don't have access to the teacher's weights or logits, but can obtain generated text, response distillation is the available KD channel. In this approach, the teacher generates text responses for prompts, and the student fine-tunes on selected (prompt, response) pairs.

This is Supervised Fine-Tuning (SFT) on teacher-generated targets, not ground truth. Because an API commonly provides text rather than full token probabilities, the training channel is less detailed than direct logit access. A teacher can provide worked solutions, decomposed subproblems, critiques, or multiple candidate answers, but these outputs should be filtered with task-specific checks where possible. For a warehouse-routing assistant, a generated route is useful training data only after checks for capacity, delivery-window, and policy constraints.

Examples and adjacent synthetic-data recipes

• Alpaca^[5]: Fine-tuned LLaMA 7B on 52K instruction-following examples generated by text-davinci-003.
• Orca^[3]: Learned from GPT-4 explanation traces plus guidance from ChatGPT, moving beyond shallow answer imitation.
• phi-1.5^[4]: Not classic KD, but an adjacent synthetic-data recipe built from textbook-like generated data; it does not compare KD loss choices.
• Distilling Step-by-Step^[6]: Uses generated rationales as an additional supervised target and evaluates whether smaller students improve on the studied tasks.
• DeepSeek-R1-Distill^[7]: Fine-tunes Qwen2.5- and Llama-based students (1.5B to 70B) for two to three epochs on the paper's roughly 800K-example SFT collection. The paper reports strong results for these distilled models on selected reasoning benchmarks; it does not prove that every teacher behavior transfers through text.

Response distillation and synthetic-data training overlap when a stronger model generates the selected targets. Calling a dataset "distillation" should not bypass evaluation: generated traces can be incorrect, stylistically misleading, contaminated, or out of scope for the intended student.

Aligning internal layers: feature distillation

Logit distillation matches output distributions. With white-box access, a training objective can also match selected student hidden states to selected teacher hidden states through a learned projection.

$\mathcal{L}_{\text{feature}} = \sum_l \|f_l^{\text{teacher}} - g(f_l^{\text{student}})\|^2$

Where $f_l^{\text{teacher}}$ and $f_l^{\text{student}}$ are hidden states at layer $l$, and $g(\cdot)$ projects student features into the teacher feature space before comparison.

Since the student commonly has different hidden dimensions, $g$ maps the student's representation into the teacher comparison space. Feature matching introduces extra decisions: which layers correspond, how the projection is trained, and whether its additional compute improves held-out outcomes. Hidden-state access is a richer interface, not a guarantee of a better student.

Forward versus reverse KL

When minimizing KL divergence for language generation, direction matters. Classical distillation commonly minimizes Forward KL (teacher || student), which penalizes a student for missing probability mass that the teacher assigns to continuations. When a small student cannot model the teacher distribution well, this pressure may be costly.

Reverse KL (student || teacher) places more pressure on probability mass the student assigns where the teacher assigns little. MiniLLM reports improvements over its studied standard-KD baselines using reverse KL with an on-policy optimization algorithm in instruction-following experiments.^[8] GKD evaluates multiple divergences and explicitly reports that the best divergence depends on the task and diversity-performance tradeoff.^[9]

No divergence is the default winner for every task. Measure task quality, diversity, calibration, and failure rates under the actual decoding setup.

Off-policy versus on-policy distillation

Off-policy (standard) distillation trains the student on teacher outputs for a static dataset. During inference, the student generates its own tokens, causing distribution shift. Errors compound as the student drifts from training data (exposure bias).

On-policy methods like Generalized Knowledge Distillation (GKD) sample sequences from the student, then compare student and teacher token distributions on the prefixes the student produced. GKD can mix fixed outputs and student-generated outputs through a student-data fraction $\lambda$; it does not require a natural-language critique.^[9] The tradeoff is computational: both student sampling and teacher scoring run during training. It is useful when fixed teacher data misses prefixes that the deployed student commonly enters, but the benefit must be measured per task.

Temperature: expose rankings without trusting them

The temperature parameter $T$ controls how much we smooth the teacher distribution used in a logit loss. It changes which relative token preferences are visible to the loss; it does not guarantee that those preferences are useful or that the student can represent them.

Mathematically, the softmax function converts raw logits into probabilities using an exponential function. When one logit is significantly larger than the rest, the exponential makes it dominate the entire probability mass, crushing the others to near-zero. By dividing all logits by a temperature $T > 1$ before applying the exponential, we reduce the relative difference between the largest logit and the smaller ones. This prevents the top choice from monopolizing the probability space and allows the relative scores of the "incorrect" choices to emerge.

A worked example

Imagine the teacher sees the prompt "Return policy for electronics?" and produces these raw logits for the next token:

At $T = 1$, the probabilities are sharp: 30 gets about 0.82, 14 gets 0.11, 90 gets 0.04, and 1 gets 0.03. The student sees only a weak signal that 90 is a more reasonable guess than 1.

At $T = 4$, after dividing each logit by 4 and applying softmax, the probabilities spread out: 30 gets about 0.40, 14 gets 0.24, 90 gets 0.19, and 1 gets 0.17. The distillation loss now exposes more of the teacher's ranking among alternatives.

Key insight: Temperature changes the training target, not the trusted answer. If the teacher ranks an invalid option highly, a softened loss makes that error easier to copy. Use labels or checks and held-out evaluation alongside KD.

• Low Temperature ($T \approx 1$): The distribution is sharp. The model is very confident in its top choice (30). The relationships between incorrect choices (14 vs 90) are hidden because their probabilities are near zero.
• High Temperature ($T > 1$): The distribution flattens. Probability mass spreads to alternatives, making the teacher's relative ranking more visible to the loss.
• Too High ($T \gg 10$): The distribution becomes nearly uniform. The ranking signal is lost, and the student learns little from the teacher distribution.

In practice, tune temperature rather than assuming a universal constant. The right value depends on the teacher distribution, soft-loss weight, task, and evaluation metrics. Inspect probability spreads as in the executable example above, then compare held-out quality.

A practical distillation training loop

A typical distillation training loop involves a frozen teacher model and a trainable student model. We forward the same input through both models and compute the combined loss. The local example below uses tiny PyTorch models so you can test the mechanics without downloading a real teacher.

In production, the same pattern usually lives inside a framework trainer, with aligned-vocabulary models such as a larger Gemma teacher and smaller Gemma student when direct token-probability distillation is required. Teams may pre-compute some teacher signal offline to avoid running the teacher inside every student update. Dense next-token logits across a long corpus are costly to store, so a design may consider top-k logits, teacher responses, or online scoring, then measure the quality effect of compression. The payload-only estimate below ignores metadata and storage-format overhead, so treat it as a lower-bound sizing exercise.

If an online-distillation pilot is feasible, compare it with an offline baseline before committing to large-scale data generation. The result can expose whether fresh teacher scoring is worth its compute cost for this task.

Distilled models in practice

When the teacher is only available through generated text, response distillation becomes the default. When you control both models and can inspect logits or hidden states, white-box distillation exposes a richer training signal. Recent systems use both patterns depending on what the teacher exposes.

What changes from system to system isn't the core idea, but the supervision channel: raw logits, hidden states, instruction-response pairs, rationale traces, or synthetic corpora generated by a stronger model.

Building a distillation dataset

When using response-based distillation, the selected dataset is one of the main controllable inputs, alongside student capacity and training budget. Build a generation and selection pipeline that can reject incorrect, duplicate, contaminated, or irrelevant examples.

Seed-Expand-Filter pipeline

A structured approach to generating teacher data involves a Seed-Expand-Filter pipeline. It does not ensure quality by itself; each filter needs a measurable contract and a separate evaluation split.

1. Seed: Start with a small set of high-quality, human-written prompts (e.g., 100 examples about warehouse operations).
2. Expand: Ask the teacher model to generate new, diverse variations of these prompts.
3. Generate: Have the teacher answer these new prompts, often with rationales or decomposed steps when richer supervision helps.
4. Filter: Use checks, deduplication, safety screening, or reviewed scoring rules to reject unsuitable generations.

This is close to the spirit of Alpaca's Self-Instruct-style pipeline and Orca's richer explanation-trace data generation, even though real systems add deduplication, safety filters, and task balancing.^[5][3]

This logic can be implemented by creating a small wrapper around a teacher client. The example below focuses on two easy-to-miss requirements: deduplicate prompts before paying for generation, and verify teacher answers before they become student targets.

Keep training data out of evaluation

Teacher generation can quietly contaminate a benchmark when prompts, reference solutions, or close rewrites enter the student training set. At minimum, block exact normalized overlap before training. For real releases, extend the gate with near-duplicate and reference-solution checks.

Limitations and when not to distill

Distillation does not make capacity, context-window, or data-coverage constraints disappear. A student may beat its teacher on a narrow checked metric after filtering or task-specific training, while regressing on other behavior. Treat the teacher and student as separate artifacts to evaluate.

Before investing in a distillation pipeline, define which behaviors matter and how they will be tested.

Legal and ethical considerations

Beyond technical constraints, distillation introduces unique licensing and evaluation challenges. Because the student model closely mirrors the teacher's outputs, the origins of that training data are important.

• Provider terms matter: the Stanford Alpaca release was research-only and non-commercial, and the repo points to both the underlying LLaMA restrictions and the dataset's CC BY-NC 4.0 terms.^[5]
• Restrictions must be reviewed: before generating a corpus or shipping a student, review the teacher access terms, base-student license, generated-data license, and permitted use of outputs. Do not infer permission from technical access.
• Imitation isn't capability proof: a student may reproduce style or familiar output patterns while failing new checked tasks. Held-out evaluation, not resemblance, establishes value.

Cost-quality tradeoff

Deciding whether to distill, and which method to use, comes down to measured quality and economics. Richer teacher access can enable different losses; it does not rank final models without evaluation.

Also, don't stop at distillation loss. A student can match teacher probabilities on training batches and still regress on held-out generation quality, long-context behavior, or latency targets. Measure task metrics, pairwise win rate, and real serving cost together.

Practice checkpoints

How does temperature affect distillation? Temperature controls how soft the teacher distribution becomes. A higher $T$ spreads probability mass across alternatives, exposing more of the teacher ranking to the loss. Too high becomes nearly uniform and loses ranking signal; any setting still needs held-out validation.

Can you distill from a proprietary API model? If its permitted interface and terms allow generated outputs for training, the available KD channel is response distillation. You generate candidate prompt-response pairs or checked solution traces and fine-tune the student on selected text. You don't get token probabilities, so filtering, use rights, and held-out evaluation are central.

When does distillation fail to transfer capability? It fails on a target behavior when student capacity, context length, architecture, or training coverage cannot reproduce that behavior at the required threshold. Define slices such as long-context tasks, policy edge cases, and checked multi-step solutions so the failure is visible.

How do you decide whether to deploy a distilled student? Set quality, policy, latency, and cost gates before training. Lower latency is not enough if a student fails a policy or accuracy gate, as the executable deployment example shows.

Can you directly distill between different tokenizers? Not with plain token-level KL unless token-to-id output mappings agree. Same vocabulary size is not enough. If output spaces differ, use response distillation or design and validate an explicit mapping.

Check your understanding

Before moving on, try to answer these questions without looking back at the article.

1. Why does a soft label teach more than a hard label?
   Hint: Think about what the student learns about the near-miss answers, not only the correct one.

2. Why might $T = 1$ provide less ranking signal than a higher temperature? Hint: Consider what a sharp teacher distribution exposes about alternatives.

3. Why do causal language models need a one-token shift when computing distillation loss?
   Hint: Remember that a causal LM predicts the next token from all previous tokens.

4. Why should forward versus reverse KL be chosen through evaluation? Hint: Think about mode coverage, low-teacher-probability mass, and task-dependent diversity requirements.

Click to see solution sketches

1. A hard label only tells the student which answer is selected. A soft label exposes the teacher's full ranking over alternatives. That extra signal can help, but a trusted check still has to catch bad teacher rankings.

2. If the teacher distribution is sharp at $T = 1$, the top token receives most of the probability mass and alternatives contribute little signal. Raising temperature can expose their relative ranking, until excessive smoothing removes useful separation.

3. In a causal LM, the logits produced at position $t$ predict the token at position $t+1$. If you don't shift the labels by one, you are training the model to predict the current token from itself, which is trivial and wrong.

4. Forward KL penalizes missing teacher mass; reverse KL penalizes student mass where the teacher is low. Their quality and diversity tradeoffs vary by task and decoding setup, so evaluate both where the objective choice matters.

What to remember

1. Core concept: Distillation transfers teacher signal, not guaranteed truth. Keep trusted labels, checks, and held-out evaluation in the design.
2. Methods: Response KD uses selected generated text. Logit KD uses aligned token probabilities. Feature KD adds hidden-state comparisons. On-policy KD scores prefixes produced by the student.
3. Temperature: Softening can expose more of a teacher ranking; too much removes separation. With the classical soft-loss formulation, scale KL by $T^2$ to preserve gradient magnitude.
4. Loss function: For causal-LM logit KD, shift positions correctly, ignore masked positions, and verify identical token-to-id output mappings. Divergence direction is an evaluated design choice.
5. On-policy versus off-policy: Static teacher targets may miss prefixes the student produces. GKD adds student-generated sequences and teacher token-distribution scoring at additional compute cost.
6. Data and release gates: Deduplicate and check generated targets, protect evaluation splits, and release only if quality, policy, latency, and cost gates pass.

Common pitfalls

What you should be able to defend

By the end of this chapter, you should be able to:

• Explain why soft labels carry more information than hard labels.
• Choose between response, logit, feature, and on-policy distillation based on available teacher signal and measurable costs.
• Implement causal-LM logit distillation with one-token label shifting and vocabulary checks.
• Explain why forward KL and reverse KL behave differently for generative models.
• Design a generated-data pipeline with checks, deduplication, held-out protection, and licensing clarity.
• Decide whether a distilled student passes quality, policy, latency, and cost release gates.