LearnAdvanced Training & AdaptationKnowledge Distillation for LLMs

⚡HardFine-Tuning & Training

Knowledge Distillation for LLMs

Understand the main forms of knowledge distillation for LLMs, from logit matching and response-based supervision to on-policy KD. Learn when distillation helps, where student capacity becomes the bottleneck, and how to implement a correct teacher-student training loop.

29 min read

Learning path

Step 118 of 196 in the full curriculum

RLVR & Verifiable Rewards Model Merging and Weight Interpolation

Personalize this lesson

Adapt explanations and teaching visuals to your background and preferred voice.

The previous lesson used Reinforcement Learning from Verifiable Rewards (RLVR): train a policy against checked outcomes where a verifier exists. Distillation starts after that useful teacher exists. At 09:00, a 7B student language model, a smaller large language model (LLM), passes staging's familiar restore-window cases. At 09:15, it answers a new expired-admin case with 7 days; trusted policy says not allowed. The student is fast enough for production, but copying that regression isn't a deployment win.

Knowledge distillation trains a smaller student to imitate useful behavior from a larger or otherwise more capable teacher. In this example, a 70B-class teacher answers "How many days can a disabled production API key stay restorable?" well enough to pass your checks, but you can't afford that model on every request. A 7B-class student is viable only if it keeps enough of that behavior under a latency and serving-cost budget.

That transfer signal depends on what the teacher exposes. In white-box settings, the student can match softened token probabilities or internal features. In black-box settings, it fine-tunes on selected teacher-written answers, solution traces, or synthetic corpora. The student doesn't inherit the teacher's judgment. It copies the behavior you can measure, then earns release through the same restore-window checks, latency target, and cost budget.

Before computing a token-level loss, compare teacher and student on the same token vocabulary, prompt formatting, and response positions. Two arrays can share a shape while assigning column 17 to different tokens. The loss still runs, but it trains the student toward the wrong words.

Why soft labels teach more than hard labels

The staging regression points at a weakness in hard labels. Hinton et al. popularized a practical formulation in 2015: train a student to mimic a teacher's behavior rather than relying only on ground-truth labels. A 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}

Start with one response token. The trusted answer is 7d, so a hard label puts mass 1.00 on 7d and zero on every alternative. Predict what gets lost: the teacher's ranking among plausible wrong answers. Its output also contains dark knowledge. Here 3d is the strongest near-miss, with smaller mass on 30d and never. That extra signal is useful only if the teacher's ranking is itself useful. It isn't proof that the student learned a general access-policy rule.

Three bar charts of the same four restore-window answers. A hard label puts mass 1.00 on 7d and zero on 3d, 30d, and never. At temperature 1 the teacher keeps 0.818 on 7d and 0.111 on the near-miss 3d. At temperature 4 the ranking stays 7d then 3d then 30d then never, but the top gap shrinks from 0.708 to 0.156. — Hard labels hide alternatives. At T = 1 the teacher still ranks 7d above 3d by 0.708. Raising T to 4 keeps that order while shrinking the gap to 0.156, so near-misses contribute real gradient. A held-out check still has to decide whether the ranking is worth copying.

Those four bars are one simplified prediction. For a causal LLM, logit distillation applies the same idea at each predicted token position. Temperature $T$ divides the logits before softmax. $T = 1$ is the model's ordinary distribution. $T > 1$ flattens it so second-place tokens aren't rounded into zero. Push $T$ far enough and the ranking disappears.

Predict the top gap before running this check: it should shrink as $T$ rises, then nearly vanish. The code uses the same logits as the figure, $[4.0, 2.0, 1.0, 0.7]$ , and prints the four probabilities plus $p(7d) - p(3d)$ at $T = 1$ , $T = 4$ , and $T = 40$ .

temperature-softening.py

import math

logits = {"7d": 4.0, "3d": 2.0, "30d": 1.0, "never": 0.7}

def softened_probabilities(temperature: float) -> dict[str, float]:
    scaled = {token: math.exp(logit / temperature) for token, logit in logits.items()}
    total = sum(scaled.values())
    return {token: value / total for token, value in scaled.items()}

for temperature in (1.0, 4.0, 40.0):
    probs = softened_probabilities(temperature)
    rounded = {token: round(probability, 3) for token, probability in probs.items()}
    top_gap = probs["7d"] - probs["3d"]
    print(f"T={temperature:g} probabilities:", rounded, "top_gap:", round(top_gap, 3))

Output

T=1 probabilities: {'7d': 0.818, '3d': 0.111, '30d': 0.041, 'never': 0.03} top_gap: 0.708
T=4 probabilities: {'7d': 0.397, '3d': 0.241, '30d': 0.188, 'never': 0.174} top_gap: 0.156
T=40 probabilities: {'7d': 0.263, '3d': 0.25, '30d': 0.244, 'never': 0.242} top_gap: 0.013

Once alternatives carry signal, decide how much to trust that signal. A common training objective blends the teacher-matching loss with a ground-truth task loss:

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

$\mathcal{L}_{\text{distill}}$ matches teacher behavior, $\mathcal{L}_{\text{task}}$ matches trusted labels, and $\alpha \in [0, 1]$ sets the mix. Raising $\alpha$ copies more of the teacher, including its mistakes. That makes $\alpha$ a behavior-and-risk choice, not a knob to tune from training loss alone.

That trade-off is why distillation isn't automatically better than training from scratch. It's most useful when you already have a strong teacher, legal access to its signal, and a clear smaller deployment target. The Gemma 2 report gives a controlled example: the 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} That 7B teacher is an ablation sized to mimic their 27B-to-9B gap, not a claim about the production teacher.

Gemma 2 writes the distillation objective as a token-level negative log-likelihood under the teacher distribution, which matches forward KL up to the teacher's entropy, a constant with respect to the student. Different recipes still expose different 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} These examples motivate careful data and signal selection. They don't establish one universally best recipe.

Soft labels explain why a teacher distribution can beat a one-hot target. Next, put that target on the causal language-model path and see what can go wrong.

Matching the teacher's probabilities: logit distillation

Logit KD needs the teacher's logits: raw, unnormalized scores from the final layer, before softmax. Training minimizes KL divergence (Kullback-Leibler divergence) so the student's token distribution tracks the teacher's.

Send the same restore-window prompt through both models. Predict the update: the teacher's output is a target, while only the student's weights move.

Diagram showing Shared prompt tokens, Frozen teacher, Trainable student, and Teacher probs q at T. — Shared prompt tokens, Frozen teacher, Trainable student, and Teacher probs q at T.

Start by applying temperature scaling, dividing both models' logits by $T > 1$ before softmax:

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)}

$z_t$ are the teacher's logits, $z_s$ are the student's logits, and $T$ is temperature. The KL term is then

$\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)
$p_i$ is the student's softened probability for token $i$ (the current prediction)
$T$ is temperature (higher $T$ is softer and closer to uniform)
The $T^2$ factor compensates for gradient scaling. Hinton et al. note that gradients from soft targets scale as about $1/T^2$ , so multiplying by $T^2$ keeps the soft term's relative weight stable while you tune temperature.^{[1]Reference 1Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531}
The hard task loss, when you use one, is ordinary cross-entropy at $T = 1$
KL measures extra surprisal from coding samples from $q$ with a code built for $p$

Two implementation details can turn a plausible loss into bad supervision. First, next-token training needs 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 aren't enough: token id 42 must denote the same token in both models. If the output spaces differ, plain token-level KL no longer lines up and you usually fall back to response distillation or design an explicit mapping.

⚠️ Common mistake: Running logit distillation without verifying tokenizer alignment. Two models can share a vocabulary size and still disagree on token-id mappings. Compare the complete output mapping, not vocab_size alone, before training.

Use this stdlib toy loss to check both mechanics without a neural net. It shifts next-token labels, skips ignored positions, computes $T^2$ KL plus hard cross-entropy, and fails fast on a reordered vocabulary.

causal-kd-loss.py

import math

VOCAB = ("7d", "3d", "30d", "never")

def softmax(logits: list[float], temperature: float) -> list[float]:
    scaled = [logit / temperature for logit in logits]
    offset = max(scaled)
    exps = [math.exp(logit - offset) for logit in scaled]
    total = sum(exps)
    return [value / total for value in exps]

def log_softmax(logits: list[float], temperature: float) -> list[float]:
    return [math.log(prob) for prob in softmax(logits, temperature)]

def kl_divergence(teacher_probs: list[float], student_log_probs: list[float]) -> float:
    return sum(
        teacher_prob * (math.log(teacher_prob) - student_log_prob)
        for teacher_prob, student_log_prob in zip(teacher_probs, student_log_probs)
    )

def cross_entropy(logits: list[float], label: int) -> float:
    return -log_softmax(logits, 1.0)[label]

def distillation_loss(
    student_logits: list[list[float]],
    teacher_logits: list[list[float]],
    labels: list[int],
    student_vocabulary: tuple[str, ...],
    teacher_vocabulary: tuple[str, ...],
    temperature: float = 3.0,
    alpha: float = 0.5,
    ignore_index: int = -100,
) -> float:
    if student_vocabulary != teacher_vocabulary:
        raise ValueError(
            "Logit KD requires identical token-to-id mappings. "
            "Use response KD or design an explicit mapping when output spaces differ."
        )
    width = len(student_vocabulary)
    if any(len(row) != width for row in student_logits + teacher_logits):
        raise ValueError("Logit rows and vocabulary dimensions must agree.")

    # Causal LMs predict token t+1 from positions up to t.
    shift_student = student_logits[:-1]
    shift_teacher = teacher_logits[:-1]
    shift_labels = labels[1:]

    soft_terms: list[float] = []
    hard_terms: list[float] = []
    for student_row, teacher_row, label in zip(shift_student, shift_teacher, shift_labels):
        if label == ignore_index:
            continue
        teacher_probs = softmax(teacher_row, temperature)
        student_log_probs = log_softmax(student_row, temperature)
        soft_terms.append(kl_divergence(teacher_probs, student_log_probs) * temperature**2)
        hard_terms.append(cross_entropy(student_row, label))

    if not soft_terms:
        return 0.0
    soft_loss = sum(soft_terms) / len(soft_terms)
    hard_loss = sum(hard_terms) / len(hard_terms)
    return alpha * soft_loss + (1 - alpha) * hard_loss

student_logits = [
    [0.20, 0.10, 0.00, -0.10],
    [1.50, 0.40, 0.10, -0.20],
    [0.30, 2.00, 0.20, 0.00],
    [0.10, 0.20, 0.30, 0.40],
]
teacher_logits = [
    [2.00, 0.50, 0.10, 0.00],
    [3.00, 1.00, 0.20, 0.10],
    [0.20, 2.50, 0.40, 0.10],
    [0.00, 0.10, 0.20, 1.50],
]
labels = [0, 0, 1, -100]

loss = distillation_loss(
    student_logits,
    teacher_logits,
    labels,
    VOCAB,
    VOCAB,
    temperature=3.0,
    alpha=0.6,
)
mismatch_failed = False
try:
    distillation_loss(
        student_logits,
        teacher_logits,
        labels,
        VOCAB,
        tuple(reversed(VOCAB)),
    )
except ValueError as exc:
    mismatch_failed = "token-to-id mappings" in str(exc)

print("loss:", round(loss, 4))
print("loss_is_finite:", math.isfinite(loss))
print("mismatch failed:", mismatch_failed)

Output

loss: 0.72
loss_is_finite: True
mismatch failed: True

White-box logits are the richest token-level channel. Many teachers expose text only, so the supervision contract changes.

When you only have text: response distillation

Without weights or logits, response distillation is the available KD channel: the teacher writes answers, and the student fine-tunes on selected (prompt, response) pairs.

This is Supervised Fine-Tuning (SFT) on teacher-generated targets, not ground truth. An API commonly provides text rather than full token probabilities, so the channel is thinner than direct logit access. A teacher can provide worked solutions, decomposed subproblems, critiques, or multiple candidates, but those outputs should pass task-specific checks where you have them. For the restore-window assistant, ask which row you would train on before looking at the teacher's prose: only an answer that matches a trusted rule belongs in the target set. That is the same verifier boundary used in the RLVR chapter.

Student / recipe	Teacher	Signal transferred	Why it matters
Alpaca 7B^{[5]Reference 5Stanford Alpaca: An Instruction-following LLaMA Model.https://github.com/tatsu-lab/stanford_alpaca}	`text-davinci-003`	52K generated instruction-response examples	The repository reports preliminary instruction-following evaluation and clear non-commercial dataset terms.
Orca 13B^{[3]Reference 3Orca: Progressive Learning from Complex Explanation Traces of GPT-4.https://arxiv.org/abs/2306.02707}	GPT-4 + ChatGPT	Explanation traces and task instructions	Evaluates a richer generated-trace recipe, rather than logit KD.
phi-1.5^{[4]Reference 4Textbooks Are All You Need II: phi-1.5 technical report.https://arxiv.org/abs/2309.05463}	Existing LLMs + curated synthetic data	Textbook-like synthetic corpora	Adjacent synthetic-data recipe, not a teacher-distribution KD comparison.
Distilling Step-by-Step^{[6]Reference 6Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data.https://arxiv.org/abs/2305.02301}	Generated rationales plus labels	Rationales as an extra supervised target	Tests whether smaller students improve on the studied tasks with rationale supervision.
Gemma 2 2B / 9B^{[2]Reference 2Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}	Larger Gemma teachers	Token-probability distillation during pretraining	Reports a controlled 2B distilled-versus-from-scratch ablation.
DeepSeek-R1-Distill^{[7]Reference 7DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learninghttps://arxiv.org/abs/2501.12948}	DeepSeek-R1	About 800K selected SFT examples	Fine-tunes Qwen2.5 and Llama bases (1.5B, 7B, 8B, 14B, 32B, 70B) with SFT only, no RL stage on those distilled checkpoints. The paper reports 55.5% on AIME 2024 for Distill-Qwen-7B and 72.6% for Distill-Qwen-32B.

Response distillation and synthetic-data training overlap when a stronger model generates the selected targets. The label "distillation" doesn't remove the need for evaluation. Generated traces can be wrong, stylistically misleading, contaminated, or out of scope for the intended student.

Use the restore-window gate as a small prediction test. The teacher is trusted on disabled prod key restore and wrong on expired admin restore, so predict which prompt the filter rejects.

teacher-output-gate.py

generated = [
    {"prompt": "disabled prod key restore", "teacher": "7 days", "verified": "7 days"},
    {"prompt": "expired admin restore", "teacher": "7 days", "verified": "not allowed"},
    {"prompt": "rotate staging key window", "teacher": "not restorable", "verified": "not restorable"},
]

accepted = [
    example for example in generated
    if example["teacher"] == example["verified"]
]
rejected = [
    example["prompt"] for example in generated
    if example["teacher"] != example["verified"]
]

print("generated:", len(generated))
print("accepted:", len(accepted))
print("rejected prompts:", rejected)
print("teacher text is trusted label:", len(rejected) == 0)

Output

generated: 3
accepted: 2
rejected prompts: ['expired admin restore']
teacher text is trusted label: False

Text and logits supervise the output. With white-box access you can also ask hidden states to match, which introduces a projection problem that output losses don't have.

Aligning internal layers: feature distillation

Logit distillation matches output distributions. White-box access also lets a training objective 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$

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

Feature matching adds three decisions before it adds value: which layers correspond, how to train the projection, and whether its extra compute improves held-out outcomes. Hidden-state access is a richer interface, not a guarantee of a better student.

Method	Teacher signal	Main advantage	Main constraint
Response KD	Selected text outputs	Works without white-box access	Teacher errors become SFT targets unless filtered
Logit KD	Token probabilities	Preserves distribution information	Requires aligned output space or an explicit mapping
Feature KD	Selected hidden states	Exposes intermediate representations	Needs layer/projection design and more storage or compute
On-policy KD	Teacher scores on student samples	Visits prefixes the student produces	Requires online sampling and teacher evaluation

So far KL was written as $D_{\mathrm{KL}}(q \| p)$ , with the teacher on the left. For generation, that direction is a choice, not a law.

Forward versus reverse KL

When minimizing KL 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, that coverage pressure can be costly. Predict the consequence: a capacity-limited student may spread probability across many teacher modes instead of committing to a usable answer.

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 reports that the best divergence depends on task and the diversity-performance tradeoff.^{[9]Reference 9On-Policy Distillation for Language Models.https://arxiv.org/abs/2306.13649}

Direction	Formula	Behavior	Common fit
Forward KL	$D_{KL}(P_{teacher} \\| P_{student})$	Mean-seeking, covers more of the teacher distribution	Classic KD when broad coverage matters
Reverse KL	$D_{KL}(P_{student} \\| P_{teacher})$	Penalizes student mass in teacher-low-probability regions	Candidate objective to evaluate for generation

Stay inside the restore-window story. At a later token the teacher still prefers 7d (0.58) over 3d (0.40) and almost never predicts never (0.02). Predict which student should trigger the larger reverse-KL penalty: one covers the ranking but leaks 0.10 onto never; another dumps 0.25 onto never. The latter is the stronger out-of-support mistake, and reverse KL more than doubles while forward KL rises less.

Two bar charts of KL on a shared 0 to 0.50 axis. Teacher mass is 0.58 on 7d, 0.40 on 3d, and 0.02 on never. When the student leaks 0.10 onto never, forward KL is 0.051 and reverse KL is 0.084. When never-mass rises to 0.25, forward KL is 0.218 and reverse KL is 0.436, so reverse KL reacts more strongly to student mass where the teacher is near zero. — Both charts use the same 0 to 0.50 KL axis. Reverse KL jumps from 0.084 to 0.436 when the student piles mass on never, the token the teacher almost never predicts. That makes reverse KL a relevant candidate, not an automatic winner.

No divergence wins every task. Measure task quality, diversity, calibration, and failure rates under the decoding setup you will ship.

kl-direction-diagnostic.py

import math

teacher = {"7d": 0.58, "3d": 0.40, "never": 0.02}
students = {
    "covers_teacher": {"7d": 0.54, "3d": 0.36, "never": 0.10},
    "dumps_on_never": {"7d": 0.40, "3d": 0.35, "never": 0.25},
}

def kl(left: dict[str, float], right: dict[str, float]) -> float:
    return sum(prob * math.log(prob / right[token]) for token, prob in left.items())

for name, student in students.items():
    forward = kl(teacher, student)
    reverse = kl(student, teacher)
    print(name, "forward:", round(forward, 3), "reverse:", round(reverse, 3))

print("choose objective from evaluation, not slogan")

Output

covers_teacher forward: 0.051 reverse: 0.084
dumps_on_never forward: 0.218 reverse: 0.436
choose objective from evaluation, not slogan

KL direction decides how the student fits a distribution on a given prefix. It can't help with prefixes the distillation set never contains.

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. At inference the student generates its own tokens, so its prefix distribution shifts. Predict the failure: once one early token changes, later prefixes can leave the fixed training set, and errors compound. That's exposure bias.

On-policy methods such as 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. This is useful when fixed teacher data misses prefixes that the deployed student commonly enters, but the benefit has to be measured per task.

Gemma 2's post-training recipe applies the same idea after SFT on predominantly synthetic teacher responses: the authors also distill from the teacher on the student's distribution, citing GKD and MiniLLM.^{[2]Reference 2Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf}^{[9]Reference 9On-Policy Distillation for Language Models.https://arxiv.org/abs/2306.13649}^{[8]Reference 8MiniLLM: On-Policy Distillation of Large Language Models.https://arxiv.org/abs/2306.08543}

on-policy-prefix-coverage.py

fixed_teacher_prefixes = {
    "restore disabled prod key",
    "rotate staging key",
}
student_generated_prefixes = {
    "restore disabled prod key",
    "restore expired admin key",
    "restore key without approval",
}

unseen_in_fixed_data = student_generated_prefixes - fixed_teacher_prefixes
teacher_scored_prefixes = fixed_teacher_prefixes | student_generated_prefixes

print("fixed prefixes:", len(fixed_teacher_prefixes))
print("student prefixes needing new teacher scores:", sorted(unseen_in_fixed_data))
print("scored after on-policy collection:", len(teacher_scored_prefixes))

Output

fixed prefixes: 2
student prefixes needing new teacher scores: ['restore expired admin key', 'restore key without approval']
scored after on-policy collection: 4

The losses and sampling rules matter only if a training step updates the student and not the teacher. The next loop makes that invariant visible.

A practical distillation training loop

A typical loop freezes the teacher, sends the same input through both models, and updates only the student with the combined loss. The local example uses tiny positionwise PyTorch stand-ins rather than Transformers, so you can test the mechanics without downloading a real teacher.

tiny-kd-training-step.py

import torch
from torch import nn
import torch.nn.functional as F

class TinyLM(nn.Module):
    def __init__(self, vocab_size: int, hidden_size: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.output = nn.Linear(hidden_size, vocab_size)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.output(self.embedding(input_ids))

def kd_loss(student_logits: torch.Tensor, teacher_logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    temperature = 3.0
    alpha = 0.7

    shift_student = student_logits[:, :-1, :]
    shift_teacher = teacher_logits[:, :-1, :].detach()
    shift_labels = labels[:, 1:]

    student_flat = shift_student.reshape(-1, shift_student.size(-1))
    teacher_flat = shift_teacher.reshape(-1, shift_teacher.size(-1))
    labels_flat = shift_labels.reshape(-1)

    soft_teacher = F.softmax(teacher_flat / temperature, dim=-1)
    soft_student = F.log_softmax(student_flat / temperature, dim=-1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * temperature**2
    hard_loss = F.cross_entropy(student_flat, labels_flat)
    return alpha * soft_loss + (1 - alpha) * hard_loss

torch.manual_seed(0)
vocab_size = 12
teacher = TinyLM(vocab_size=vocab_size, hidden_size=16)
student = TinyLM(vocab_size=vocab_size, hidden_size=6)
teacher.requires_grad_(False)
teacher.eval()

input_ids = torch.tensor([
    [1, 2, 3, 4, 5],
    [1, 3, 5, 7, 9],
])
labels = input_ids.clone()
optimizer = torch.optim.AdamW(student.parameters(), lr=0.05)

with torch.no_grad():
    teacher_logits = teacher(input_ids)

before = kd_loss(student(input_ids), teacher_logits, labels)
optimizer.zero_grad()
before.backward()
has_grad = any(parameter.grad is not None for parameter in student.parameters())
optimizer.step()

after = kd_loss(student(input_ids), teacher_logits, labels)

print("before:", round(float(before), 4))
print("after:", round(float(after), 4))
print("has_grad:", has_grad)
print("after_is_finite:", bool(torch.isfinite(after)))
print("improved:", bool(after < before))

Output

before: 1.1513
after: 0.9793
has_grad: True
after_is_finite: True
improved: True

A deployed pipeline usually puts this pattern inside a framework trainer, with aligned-vocabulary models such as a larger Gemma teacher and smaller Gemma student when you need direct token-probability distillation. 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.

logit-cache-budget.py

tokens = 50_000_000
vocab_size = 32_000
bytes_per_logit = 2  # bf16
top_k = 64
bytes_per_topk_item = 2 + 4  # bf16 value plus int32 token id

dense_bytes = tokens * vocab_size * bytes_per_logit
topk_bytes = tokens * top_k * bytes_per_topk_item
gib = 1024 ** 3

print("dense cache GiB:", round(dense_bytes / gib, 1))
print("top-k cache GiB:", round(topk_bytes / gib, 1))
print("storage reduction:", round(dense_bytes / topk_bytes, 1), "x")
print("quality must still be evaluated:", True)

Output

dense cache GiB: 2980.2
top-k cache GiB: 17.9
storage reduction: 166.7 x
quality must still be evaluated: True

If an online-distillation pilot is feasible, compare it with an offline baseline before committing to large-scale data generation. That comparison can show whether fresh teacher scoring earns enough held-out quality to justify its compute cost for this task.

The training step assumes a legal, filtered corpus. For response KD, corpus design is the main lever you control, alongside student capacity and the evaluation split.

Building a distillation dataset

When using response-based distillation, selected data, student capacity, and training budget set the ceiling. Build a generation and selection pipeline that can reject incorrect, duplicate, contaminated, or irrelevant examples before they become targets.

Seed-Expand-Filter pipeline

A Seed-Expand-Filter pipeline gives teacher-data generation a sequence of decisions. It doesn't prove quality by itself. Each filter needs a measurable contract and a separate evaluation split.

Seed: Start with a small set of high-quality, human-written prompts (for example, 100 restore-window and key-rotation questions).
Expand: Ask the teacher to generate new, diverse variations of those prompts.
Generate: Have the teacher answer the new prompts, often with rationales or decomposed steps when richer supervision helps.
Filter: Use checks, deduplication, safety screening, or reviewed scoring rules to reject unsuitable generations.

Use the three-row fixture from the teacher-output gate to see those filters fire:

Prompt	Teacher text	Gate	Decision
`disabled prod key restore`	`7 days`	matches the trusted answer	keep
`Disabled prod key restore`	(not generated)	same normalized prompt	skip as duplicate
`expired admin restore`	`7 days`	trusted answer is `not allowed`	reject

This follows the broad shape of Alpaca's Self-Instruct-style pipeline and Orca's richer explanation-trace generation, while 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}

Two checks are easy to skip and expensive to recover from: deduplicate prompts before paying for generation, and verify teacher answers before they become student targets. The wrapper encodes both.

select-teacher-responses.py

from collections.abc import Callable

class DistillationDataGenerator:
    def __init__(
        self,
        teacher_generate: Callable[[str], str],
        verify_response: Callable[[str, str], bool],
    ):
        self.teacher_generate = teacher_generate
        self.verify_response = verify_response

    def generate_dataset(self, prompts: list[str]) -> list[dict[str, str]]:
        selected: list[dict[str, str]] = []
        seen: set[str] = set()
        for prompt in prompts:
            normalized = " ".join(prompt.lower().split())
            if normalized in seen:
                continue
            seen.add(normalized)
            response = self.teacher_generate(normalized).strip()
            if self.verify_response(normalized, response):
                selected.append({"prompt": normalized, "response": response})
        return selected

def fake_teacher(prompt: str) -> str:
    return "7 days"

trusted_answers = {
    "disabled prod key restore": "7 days",
    "expired admin restore": "not allowed",
}

def verify_response(prompt: str, response: str) -> bool:
    return trusted_answers[prompt] == response

generator = DistillationDataGenerator(fake_teacher, verify_response)
examples = generator.generate_dataset([
    "disabled prod key restore",
    "  Disabled prod key  restore  ",
    "expired admin restore",
])

print("selected prompts:", [example["prompt"] for example in examples])
print("selected count:", len(examples))
print("bad response retained:", any("expired admin" in example["prompt"] for example in examples))

Output

selected prompts: ['disabled prod key restore']
selected count: 1
bad response retained: False

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. Block exact normalized overlap before training. For real releases, extend the gate with near-duplicate and reference-solution checks.

held-out-contamination-gate.py

def normalize(prompt: str) -> str:
    return " ".join(prompt.lower().replace("?", "").split())

candidate_training_prompts = [
    "Compute retry budget for a failed restore",
    "Summarize failed restore logs",
    "How many days can a disabled production API key stay restorable?",
]
held_out_prompts = [
    "how many days can a disabled production API key stay restorable",
    "Estimate GPU queue delay for a remote region",
]

held_out_keys = {normalize(prompt) for prompt in held_out_prompts}
accepted = [
    prompt for prompt in candidate_training_prompts
    if normalize(prompt) not in held_out_keys
]
blocked = [
    prompt for prompt in candidate_training_prompts
    if normalize(prompt) in held_out_keys
]

print("accepted training prompts:", len(accepted))
print("blocked overlap:", blocked)
print("held-out exact overlap after gate:", any(normalize(p) in held_out_keys for p in accepted))

Output

accepted training prompts: 2
blocked overlap: ['How many days can a disabled production API key stay restorable?']
held-out exact overlap after gate: False

A clean corpus still can't make a small model remember behavior beyond its capacity or context window, and it can't settle license constraints.

Limitations and when not to distill

Distillation doesn't erase context-window or data-coverage limits. A student may beat its teacher on a narrow checked metric after filtering or task-specific training, while regressing on other behavior. Treat teacher and student as separate artifacts to evaluate.

Before investing in a distillation pipeline, write the behavior contract and its tests. If the student passes common restore cases but fails expired-admin or long-context cases, capacity, context length, prompt coverage, and teacher-target quality are competing explanations. Split those hypotheses with targeted slices instead of reading one aggregate score.

Behavior	Regression risk to test	Useful held-out gate
Domain answers	Generated targets can repeat teacher errors	Checked answer accuracy and abstention rate
Instruction following	Narrow traces can miss new constraints	Fresh constraint-following prompts
Multi-step solutions	Final answers can hide invalid steps	Step checks where available plus final-answer accuracy
Long-context use	Student architecture or context limit may differ	Retrieval and long-context slices at deployment length
Safety and policy behavior	Filtered corpus may omit refusals or edge cases	Safety-policy evaluation separate from task benchmark

Legal and ethical considerations

Because the student can mirror teacher outputs closely, the origins of that training data matter.

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

Choose a distillation method from measured quality and economics. Richer teacher access enables different losses; it doesn't rank final models without evaluation.

Approach	Required access	Main training cost	Release gate
Use teacher directly	Teacher inference	No student training	Baseline quality, latency, and cost
Response KD	Generated outputs and permitted use	Generation plus SFT	Output filtering and held-out task quality
Logit KD	Aligned teacher token probabilities	Teacher scoring or cache storage	Task quality plus cache/online cost
Feature KD	Hidden states and layer mapping	Extra projections and state transfer	Ablation against simpler KD baseline

Distillation loss is an optimization signal, not a release decision. 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.

deployment-gate.py

teacher = {"checked_accuracy": 0.94, "policy_error_rate": 0.01, "latency_ms": 180, "monthly_cost": 42000}
student = {"checked_accuracy": 0.92, "policy_error_rate": 0.04, "latency_ms": 42, "monthly_cost": 18000}
requirements = {
    "checked_accuracy": 0.90,
    "max_policy_error_rate": 0.02,
    "max_latency_ms": 60,
    "max_monthly_cost": 20000,
}

checks = {
    "quality": student["checked_accuracy"] >= requirements["checked_accuracy"],
    "policy": student["policy_error_rate"] <= requirements["max_policy_error_rate"],
    "latency": student["latency_ms"] <= requirements["max_latency_ms"],
    "cost": student["monthly_cost"] <= requirements["max_monthly_cost"],
}

print("student faster:", student["latency_ms"] < teacher["latency_ms"])
print("student cheaper:", student["monthly_cost"] < teacher["monthly_cost"])
print("release checks:", checks)
print("deploy student:", all(checks.values()))

Output

student faster: True
student cheaper: True
release checks: {'quality': True, 'policy': False, 'latency': True, 'cost': True}
deploy student: False

The gate makes that trade-off concrete: this student is faster and cheaper, and its checked accuracy clears the threshold, but its policy error rate doesn't. Keep the teacher or revise the student until every release check passes.

When distillation breaks

Symptom	Cause	Fix
Validation loss barely changes as you raise temperature.	The softened teacher distribution may be too flat or the soft-loss weight may be ineffective.	Inspect teacher probabilities and tune temperature and loss weight on held-out tasks.
Student looks strong on training prompts but weak on held-out tasks.	Distillation corpus is too narrow, repetitive, or too close to evaluation data.	Broaden prompt coverage, filter duplicates, and keep a separate held-out evaluation slice.
Student predicts current token instead of next token during logit KD.	Causal LM loss forgot the one-token shift.	Shift logits at position $t$ against labels at position $t+1$ before KL or cross-entropy.
Student copies teacher hallucinations and policy mistakes.	Distillation blindly transferred bad teacher outputs.	Filter teacher generations, add task loss, and evaluate against trusted labels or reward checks.
KL loss runs but student quality stays random.	Teacher and student token-to-id mappings don't align, even if sizes match.	Compare mappings exactly, use response distillation, or design an explicit output-space mapping.
Tiny student misses required checked behaviors.	Student capacity, context, or data coverage is insufficient for this release target.	Narrow task scope, increase student size, or revise training and evaluation design.
Offline metrics look great but production quality collapses.	Distillation and evaluation data leaked into each other.	Split generation, tuning, and evaluation sets cleanly before training starts.

Next Step

Continue to Model Merging and Weight Interpolation

Distillation trains a new student from teacher signal. Model merging asks whether compatible checkpoints can be combined into one candidate without another gradient-training run.

PreviousRLVR & Verifiable Rewards

Share this article