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 an access-policy assistant: 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. Don't assume the student inherits the teacher; transfer useful behavior and verify the resulting quality-cost tradeoff.

White-box logit KD matches teacher token probabilities at each step. Response distillation uses selected teacher text instead of the teacher-probability branch.

Why soft labels teach more than hard labels

Hinton et al. popularized a practical formulation of distillation in 2015: train a student model to mimic a teacher model's behavior rather than relying only on ground-truth labels. The teacher's full probability distribution can carry more information than a one-hot label:^{[1]Reference 1Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531}

A concrete example

Suppose an operator asks, "How many days can a disabled production API key stay restorable?" The teacher's output distribution contains dark knowledge: relationships between classes that hard labels don't capture.

The hard label only says "7 days" is correct. The soft label also shows that "3 days" is the strongest near-miss, with smaller mass on longer windows and "never." That extra signal is useful only if the teacher's ranking is itself useful; it's not proof that the student learned a general access-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.

The soft-label example exposes one teacher ranking. A common training objective blends that teacher-matching loss with a ground-truth task loss:

$\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.

What happens when $\alpha$ moves closer to 1?

Answer

The student listens more to the teacher distribution and less to the hard task labels. That can transfer richer teacher behavior, but it also makes teacher mistakes more influential.

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 doesn't 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]Reference 2Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}

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]Reference 3Orca: Progressive Learning from Complex Explanation Traces of GPT-4.https://arxiv.org/abs/2306.02707[4]Reference 4Textbooks Are All You Need II: phi-1.5 technical report.https://arxiv.org/abs/2309.05463[2]Reference 2Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf} They motivate careful data and signal selection; they don't 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, so the softened distribution still contributes meaningful gradient signal.
• KL divergence measures how much information is lost when using the student's distribution $p$ to approximate the teacher's distribution $q$

Why use temperature $T > 1$ during logit distillation?

Answer

Higher temperature softens the teacher distribution, exposing near-miss classes instead of only the top answer. The student learns the teacher's ranking over alternatives rather than a hard label alone.

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 vocab_size alone, 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 deployment-planning assistant, a generated plan is useful training data only after checks for dependency order, rollback safety, and policy constraints.

Examples and adjacent synthetic-data recipes

• Alpaca^{[5]Reference 5Stanford Alpaca: An Instruction-following LLaMA Model.https://github.com/tatsu-lab/stanford_alpaca}: Fine-tuned LLaMA 7B on 52K instruction-following examples generated by text-davinci-003.
• Orca^{[3]Reference 3Orca: Progressive Learning from Complex Explanation Traces of GPT-4.https://arxiv.org/abs/2306.02707}: Learned from GPT-4 explanation traces plus guidance from ChatGPT, moving beyond shallow answer imitation.
• phi-1.5^{[4]Reference 4Textbooks Are All You Need II: phi-1.5 technical report.https://arxiv.org/abs/2309.05463}: Not classic KD, but an adjacent synthetic-data recipe built from textbook-like generated data; it doesn't compare KD loss choices.
• Distilling Step-by-Step^{[6]Reference 6Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data.https://arxiv.org/abs/2305.02301}: Uses generated rationales as an additional supervised target and evaluates whether smaller students improve on the studied tasks.
• DeepSeek-R1-Distill^{[7]Reference 7DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learninghttps://arxiv.org/abs/2501.12948}: 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 doesn't 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" shouldn't 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.

Why does feature distillation need the projection $g(\cdot)$?

Answer

The student and teacher may have different hidden sizes. The projection maps student features into the teacher's feature space so the layer representations can be compared.

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 can't 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]Reference 8MiniLLM: On-Policy Distillation of Large Language Models.https://arxiv.org/abs/2306.08543} GKD evaluates multiple divergences and explicitly reports that the best divergence depends on the task and diversity-performance tradeoff.^{[9]Reference 9On-Policy Distillation for Language Models.https://arxiv.org/abs/2306.13649}

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 a fixed set of prefixes, which may come from ground-truth sequences or teacher-generated outputs. 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 doesn't require a natural-language critique.^{[9]Reference 9On-Policy Distillation for Language Models.https://arxiv.org/abs/2306.13649} The cost is computational: both student sampling and teacher scoring run during training. It's useful when fixed teacher data misses prefixes that the deployed student commonly enters, but the benefit must be measured per task.

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 positionwise PyTorch stand-ins rather than Transformers so you can test the loss 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 doesn't prove 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 deployment planning).
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]Reference 5Stanford Alpaca: An Instruction-following LLaMA Model.https://github.com/tatsu-lab/stanford_alpaca[3]Reference 3Orca: Progressive Learning from Complex Explanation Traces of GPT-4.https://arxiv.org/abs/2306.02707}

This logic can be implemented by creating a small wrapper around a teacher client. The wrapper handles 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 doesn't 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]Reference 5Stanford Alpaca: An Instruction-following LLaMA Model.https://github.com/tatsu-lab/stanford_alpaca}
• 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. Don't 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 doesn't 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.

When distillation breaks