Knowledge Distillation for LLMs

The previous lesson used Reinforcement Learning with Verifiable Rewards (RLVR) to train against checked outcomes. For a large language model (LLM), distillation asks a different question: can a smaller model learn useful behavior from an existing model? The teacher might come from RLVR, ordinary pre-training, or another training recipe.

Consider a hypothetical access-policy assistant. A 7B student passes familiar key-restoration cases but answers a new expired-admin case with 7 days when the fixture's trusted rule says not allowed. Lower latency doesn't compensate for that regression. The policy rules and deployment numbers in this lesson are invented examples, not a real provider's API-key policy or measured serving results.

Knowledge distillation trains a student using signals from a teacher, often to make a smaller model useful. The teacher needn't have more parameters; self-distillation can use a model's own improved targets. In our 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 four answer classes, each treated as a single token in a miniature vocabulary. For the disabled-key case, 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 alternatives. This relative probability information is often called dark knowledge. Here 3d gets more mass than 30d or never. That ranking isn't a policy severity scale, and a real tokenizer might split these strings into several tokens.

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 ordinary distribution. For finite $T > 1$, the distribution flattens but rank is preserved in exact arithmetic. Not every pairwise probability gap must shrink monotonically. As $T$ approaches infinity, the probabilities approach uniform; rounding can hide their differences much earlier.

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

import math

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

def softened_probabilities(temperature: float) -> dict[str, float]:
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be finite and positive.")
    offset = max(logits.values())
    scaled = {token: math.exp((logit - offset) / 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))

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

Token-distribution KD needs aligned teacher probabilities or logits, the raw scores before softmax. Full normalized log-probabilities are also enough to reconstruct a temperature-scaled distribution; a truncated top-k list generally isn't. Here 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.

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

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

• $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$ softens the distribution and flattens probability peaks)
• The $T^2$ factor compensates for approximate high-temperature gradient attenuation when mixing soft and hard losses; it doesn't make every finite-temperature gradient identical.^{[1]Reference 1Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531}
• The hard task loss, when included, uses ordinary cross-entropy at $T = 1$
• KL divergence measures the extra surprisal of coding samples from $q$ using a code optimized for $p$

Mathematical derivation of the T² gradient scaling factor

Why multiply soft KL by $T^2$? Hinton et al. motivate this convention through a high-temperature approximation.^{[1]Reference 1Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531} The exact derivative explains both its usefulness and its limits. Other loss weights or learning-rate choices can compensate differently; $T^2$ isn't required for a valid objective.

Consider a single token position across vocabulary $V$. Let $z_s$ and $z_t$ represent the raw student and teacher logits. With temperature $T$, the softened probability distributions are:

At a given prefix, the frozen evaluation-mode teacher supplies detached probabilities $q_i$, held fixed for this gradient calculation. New prefixes can supply different targets. Minimizing the KL divergence $D_{\mathrm{KL}}(q \parallel p) = \sum_i q_i \log q_i - \sum_i q_i \log p_i$ with respect to student parameters is identical to minimizing the cross-entropy loss under softened targets:

$\mathcal{L}_{\text{soft}} = - \sum_{i \in V} q_i \log p_i$

Let $v_i = z_{s,i} / T$ denote the scaled student logit. Differentiating $\mathcal{L}_{\text{soft}}$ with respect to $v_i$ yields the standard softmax cross-entropy gradient:

$\frac{\partial \mathcal{L}_{\text{soft}}}{\partial v_i} = p_i - q_i$

To find the gradient with respect to the actual trainable student logit $z_{s,i}$, apply the chain rule:

$\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_{s,i}} = \frac{\partial \mathcal{L}_{\text{soft}}}{\partial v_i} \cdot \frac{\partial v_i}{\partial z_{s,i}} = (p_i - q_i) \cdot \frac{1}{T}$

The chain rule supplies one factor of $1/T$. For the next approximation, temperature must be high relative to the logit spread. Softmax is invariant to a shared offset, so center each model's logits first. Using $\exp(x) \approx 1+x$ when the centered logits divided by $T$ are small:

where $\bar{z}_s = \frac{1}{|V|} \sum_j z_{s,j}$ is the mean student logit across vocabulary $V$. Applying the geometric expansion $(1 + \epsilon)^{-1} \approx 1 - \epsilon$:

The exact same expansion applies to the teacher's softened probabilities:

$q_i \approx \frac{1}{|V|} \left(1 + \frac{z_{t,i} - \bar{z}_t}{T}\right)$

Subtracting $q_i$ from $p_i$ reveals that the probability difference itself shrinks inversely with temperature:

$p_i - q_i \approx \frac{1}{|V| T} \left( (z_{s,i} - \bar{z}_s) - (z_{t,i} - \bar{z}_t) \right)$

Now substitute this probability difference back into the chain-rule gradient with respect to student logit $z_{s,i}$:

$\frac{\partial \mathcal{L}_{\text{soft}}}{\partial z_{s,i}} = \frac{1}{T} (p_i - q_i) \approx \frac{1}{|V| T^2} \left( (z_{s,i} - \bar{z}_s) - (z_{t,i} - \bar{z}_t) \right)$

For fixed finite logits in this limit, the unscaled gradient is approximately proportional to $1/T^2$. The hard-loss gradient, evaluated at $T=1$, doesn't change with the soft branch's temperature. Moving from $T=1$ to $T=4$ isn't necessarily a sixteenfold decrease: the low-temperature distribution may be far from the approximation. The mixing weight and optimizer also determine how much the teacher influences training.

Multiplying the soft distillation loss by $T^2$ cancels this attenuation:

$\mathcal{L}_{\text{KD}} = T^2 \cdot D_{\mathrm{KL}}(q \parallel p)$

Differentiating the scaled distillation loss yields:

$\frac{\partial \mathcal{L}_{\text{KD}}}{\partial z_{s,i}} = T^2 \cdot \frac{1}{T} (p_i - q_i) = T(p_i - q_i) \approx \frac{1}{|V|} \left( (z_{s,i} - \bar{z}_s) - (z_{t,i} - \bar{z}_t) \right)$

For fixed finite logits, the scaled loss approaches $\frac{1}{2|V|}\|\tilde z_s-\tilde z_t\|^2$, where each $\tilde z$ is centered by its own mean. This matches centered logits in the high-temperature limit. At finite temperatures the gradient still changes, and a teacher-matching loss can be exactly zero when the student already matches its target.

Predict two gradients before running the next exercise: does a one-hot target give alternatives zero gradient, and does $T^2$ make the soft gradient constant at every temperature? Neither. Hard-label cross-entropy has derivative $p_i-y_i$, so an incorrect class with positive student mass still receives a gradient. What it lacks is the teacher's ranking of those alternatives.

import torch
import torch.nn.functional as F

teacher = torch.tensor([4.0, 2.0, 1.0, 0.7], dtype=torch.float64)
base_student = torch.tensor([1.0, 0.4, 0.1, -0.2], dtype=torch.float64)
limit = ((base_student - base_student.mean()) - (teacher - teacher.mean())) / 4

for temperature in (1.0, 4.0, 40.0, 400.0):
    student = base_student.clone().requires_grad_()
    q_log = F.log_softmax(teacher / temperature, dim=0)
    p_log = F.log_softmax(student / temperature, dim=0)
    loss = F.kl_div(p_log, q_log, log_target=True, reduction="sum")
    raw_gradient, = torch.autograd.grad(loss, student)
    scaled_gradient = temperature**2 * raw_gradient
    exact_gradient = temperature * (p_log.exp() - q_log.exp())
    assert torch.allclose(scaled_gradient, exact_gradient, atol=1e-12)
    print(f"T={temperature:g}: raw_norm={raw_gradient.norm():.6f}; scaled_norm={scaled_gradient.norm():.6f}")
assert torch.allclose(scaled_gradient, limit, atol=0.001)

student = base_student.clone().requires_grad_()
F.cross_entropy(student.unsqueeze(0), torch.tensor([0])).backward()
assert bool((student.grad[1:] > 0).all())
print("hard-label alternatives have nonzero gradients=True")
print("scaled gradient approaches centered-logit limit=True")

T=1: raw_norm=0.434065; scaled_norm=0.434065
T=4: raw_norm=0.030553; scaled_norm=0.488844
T=40: raw_norm=0.000272; scaled_norm=0.435812
T=400: raw_norm=0.000003; scaled_norm=0.429383
hard-label alternatives have nonzero gradients=True
scaled gradient approaches centered-logit limit=True

Two implementation details can turn a plausible loss into bad supervision. First, hard next-token labels are shifted: logits at position $t$ predict the label at $t+1$. Teacher and student logits compare at the same position and prefix, not one position apart. Second, direct logit KD assumes the same token-to-id output mapping. Equal vocabulary sizes aren't enough. Verify normalization, special tokens, chat formatting, and the actual input token sequence too. Different tokenizers can segment identical text differently, leaving no one-to-one token-position alignment. Response KD avoids that particular constraint by retokenizing text for the student.

⚠️ 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 standard-library loss to check both mechanics without a neural net. It shifts hard labels, uses their mask for both losses, and rejects reordered vocabularies or mismatched sequence lengths. An all-ignored example raises an error rather than silently reporting zero training loss.

import math

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

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

def log_softmax(logits: list[float], temperature: float) -> list[float]:
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be finite and positive.")
    if not logits or not all(math.isfinite(value) for value in logits):
        raise ValueError("Logits must be nonempty and finite.")
    offset = max(logits)
    scaled = [(value - offset) / temperature for value in logits]
    normalizer = math.log(sum(math.exp(value) for value in scaled))
    return [value - normalizer for value in scaled]

def kl_divergence(teacher_probs: list[float], student_log_probs: list[float]) -> float:
    if len(teacher_probs) != len(student_log_probs):
        raise ValueError("Distribution lengths must match.")
    return sum(
        teacher_prob * (math.log(teacher_prob) - student_log_prob)
        for teacher_prob, student_log_prob in zip(teacher_probs, student_log_probs)
        if teacher_prob > 0  # The limiting contribution of zero mass is zero.
    )

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 not math.isfinite(alpha) or not 0 <= alpha <= 1:
        raise ValueError("alpha must be finite and in [0, 1].")
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be finite and positive.")
    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 width == 0 or len(set(student_vocabulary)) != width:
        raise ValueError("Vocabulary must contain distinct tokens.")
    if not len(student_logits) == len(teacher_logits) == len(labels) or len(labels) < 2:
        raise ValueError("Matching sequences with at least two positions are required.")
    if any(type(label) is not int or (label != ignore_index and not 0 <= label < width) for label in labels):
        raise ValueError("Labels must be token IDs or ignore_index.")
    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:
        raise ValueError("No supervised next-token positions remain.")
    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)

loss: 0.72
loss_is_finite: True
mismatch failed: True

The worked loss uses token distributions on aligned prefixes. Many teachers expose text only, so the supervision contract changes.

When you only have text: response distillation

Without weights or logits, response distillation (also called sequence-level distillation or black-box KD) becomes the primary transfer channel: the teacher generates full answers or reasoning traces, and the student fine-tunes on curated (prompt, response) pairs.

This is Supervised Fine-Tuning (SFT) on teacher-generated targets. Direct token matching needs aligned output meanings and positions, or an explicitly designed cross-tokenizer method. It doesn't inherently require a terabyte cache or materializing every logit at once: online scoring and chunked losses can reduce storage and activation memory. Text targets offer different practical options:

1. Tokenizer independence: Teacher and student can use completely different tokenizers, subword merges, and vocabulary sizes. The student tokenizes the teacher's raw text using its own tokenizer.
2. Decoupled compute: Generate responses offline, then train without keeping the teacher resident. Student hardware requirements still depend on its size, precision, optimizer, and sequence length. Response generation can also be refreshed during training.
3. Intermediate explanation supervision: Written rationales supply additional target tokens to learn. Logit distillation can also operate on rationale prefixes; neither channel establishes that a written rationale reveals the teacher's internal reasoning.

Ordinary response SFT presents selected tokens rather than the full teacher distribution. Multiple samples, weights, or preference filtering can supply other information, but teacher probabilities aren't automatically calibrated truth estimates. If a wrong policy exception becomes an unfiltered target, cross-entropy encourages the student to reproduce it. Use checks where available and reviewed quality criteria where exact verification isn't possible.

The original January 2025 R1 report gives DeepSeek-R1-Distill-Qwen-32B an AIME 2024 pass@1 of 72.6% and MATH-500 accuracy of 94.3%.^{[7]Reference 7DeepSeek-R1: Original January 2025 Report (v1)https://arxiv.org/html/2501.12948v1} Its response-SFT mixture includes about 600k reasoning samples and 200k general samples; some selection used model judgments rather than programmatic proof. Evaluation used temperature 0.6, top-p 0.95, and a 32,768-token generation limit. These are reported historical results at a stated budget, not present-day model rankings or proof that every explanation is valid. Sampled-correctness averages estimate pass@1; the separately reported majority-vote cons@64 aggregates answers.

Response distillation and synthetic-data generation overlap when a stronger model generates the training targets. Calling the recipe "distillation" doesn't remove the need for strict evaluation. Teacher generations can be factually wrong, stylistically repetitive, contaminated by benchmark answers, or out of scope for the student's parameter budget.

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.

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)

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_{m(l)}^{\text{teacher}} - g_l(f_l^{\text{student}})\|^2$

$m(l)$ selects the teacher layer paired with student layer $l$; their depths needn't match. The learned projection $g_l$ maps the student's hidden width to the teacher's. Compare aligned token positions, mask padding, and choose whether to average over tokens, features, and layers. That reduction sets the loss scale. Train the projection with the student while keeping teacher features detached.

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.

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 divergence for autoregressive language models, direction matters. The two directions create starkly different supervision incentives:

Both directions have zero divergence when the student matches the teacher. A constrained model family may be unable to represent that distribution, so the objectives can choose different compromises. Parameter count alone doesn't prove a particular teacher distribution is unrepresentable. The following tendencies depend on the model family and optimization:

1. Forward KL emphasizes teacher-supported outcomes: A fixed positive teacher mass incurs an unbounded penalty as its student mass approaches zero. Under a restricted family, fitting several separated teacher modes can put student mass between them. This motivates the name mode-covering. It doesn't establish that a language model must hallucinate, or that every teacher-supported answer is valid.
2. Reverse KL emphasizes student-generated outcomes: A fixed positive student mass incurs an unbounded penalty as the corresponding teacher mass approaches zero. A zero student mass contributes zero at that outcome, but reallocating the mass changes other terms. If the teacher gives $A$ and $B$ each 0.5 and the student selects only $A$, total reverse KL is $\log 2$, not zero. A constrained optimum may concentrate on fewer teacher modes, motivating mode-seeking. High teacher probability doesn't certify truth, calibration, or safety.

MiniLLM studies sequence-level reverse KL using policy gradients.^{[8]Reference 8MiniLLM: Knowledge Distillation of Large Language Modelshttps://arxiv.org/abs/2306.08543} Its token reward is a teacher/student log-probability ratio, with future-token credit, not teacher likelihood alone. It also uses single-step decomposition, teacher-mixed sampling, and length normalization. The paper's ablations observe repeated or meaningless outputs while reverse KL decreases when stabilizing components are removed. Lower divergence isn't a guarantee of faithful language.

Generalized Knowledge Distillation (GKD) separates the data-sampling rule from the divergence.^{[9]Reference 9On-policy Distillation of Language Models: Learning from Self-Generated Mistakeshttps://arxiv.org/abs/2306.13649} Its results are task-dependent: generalized Jensen-Shannon choices outperform both KL directions in its smaller-student translation experiments, while its headline on-policy GSM8K comparison uses forward KL. This directly contradicts a universal reverse-KL ranking. Measure the loss, sampling rule, model family, and decoding setup together.

The equations above can describe a distribution over complete sequences or a next-token distribution at one fixed prefix. A full autoregressive sequence KL decomposes into token KLs averaged over prefixes sampled from the distribution on the left. Teacher-prefix forward KL and student-prefix reverse KL therefore visit different states. Token KL on a fixed corpus isn't automatically the exact complete-sequence divergence.

Make that weighting visible with two-token sequences over A and B. The teacher starts with probabilities [0.8, 0.2]; the student starts with [0.5, 0.5]. Their next-token distributions also depend on that first token. Predict which prefixes get more weight in forward KL: A gets 80%, not 50%. Enumerating all four sequences must give the same result as this teacher-weighted decomposition.

import math

teacher_first, student_first = [0.8, 0.2], [0.5, 0.5]
teacher_next = [[0.9, 0.1], [0.2, 0.8]]
student_next = [[0.5, 0.5], [0.4, 0.6]]

def kl(left, right):
    return sum(a * math.log(a / b) for a, b in zip(left, right))

teacher_joint = [teacher_first[i] * teacher_next[i][j] for i in range(2) for j in range(2)]
student_joint = [student_first[i] * student_next[i][j] for i in range(2) for j in range(2)]
forward = kl(teacher_joint, student_joint)
forward_parts = kl(teacher_first, student_first) + sum(
    teacher_first[i] * kl(teacher_next[i], student_next[i]) for i in range(2)
)
reverse = kl(student_joint, teacher_joint)
reverse_parts = kl(student_first, teacher_first) + sum(
    student_first[i] * kl(student_next[i], teacher_next[i]) for i in range(2)
)
uniform_prefix_forward = kl(teacher_first, student_first) + sum(
    kl(teacher_next[i], student_next[i]) for i in range(2)
) / 2
assert math.isclose(forward, forward_parts) and math.isclose(reverse, reverse_parts)
assert not math.isclose(forward, uniform_prefix_forward)
print(f"forward joint={forward:.6f}; teacher-weighted tokens={forward_parts:.6f}")
print(f"reverse joint={reverse:.6f}; student-weighted tokens={reverse_parts:.6f}")
print(f"uniform-prefix forward={uniform_prefix_forward:.6f}; differs=True")

forward joint=0.505499; teacher-weighted tokens=0.505499
reverse joint=0.530881; student-weighted tokens=0.530881
uniform-prefix forward=0.422535; differs=True

These are sums over exactly two generated tokens, with no variable-length or stop-token issue. Averaging over positions would divide them by two. A corpus that weights A and B equally optimizes another useful possible objective, but its token loss isn't this sequence forward KL. Changing the collection distribution changes which conditional mismatches matter, even with the same teacher, student, and token-level divergence.

For another illustrative distribution, the teacher prefers 7d (0.58) over 3d (0.40) and gives never 0.02. Compare students assigning 0.10 and 0.25 to never. The latter puts more mass in a teacher-low-probability region, and reverse KL rises more in this example. This isn't an out-of-support event: all three teacher probabilities are positive.

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

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:
    if left.keys() != right.keys():
        raise ValueError("Both distributions must use the same outcomes.")
    for distribution in (left, right):
        if any(not math.isfinite(p) or p < 0 for p in distribution.values()) or not math.isclose(sum(distribution.values()), 1.0):
            raise ValueError("Probabilities must be finite, nonnegative, and sum to one.")
    terms = []
    for token, probability in left.items():
        if probability == 0:
            continue
        if right[token] == 0:
            return math.inf
        terms.append(probability * (math.log(probability) - math.log(right[token])))
    return sum(terms)

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

two_modes = {"A": 0.5, "B": 0.5}
one_mode = {"A": 1.0, "B": 0.0}
forward = kl(two_modes, one_mode)
reverse = kl(one_mode, two_modes)
assert math.isinf(forward) and math.isclose(reverse, math.log(2))
print(f"drop one mode: forward={forward}; reverse={reverse:.3f}")
# Matching a confidently wrong teacher still gives zero KL.
wrong_teacher = {"correct": 0.01, "wrong": 0.99}
assert kl(wrong_teacher, wrong_teacher) == 0.0
print("matched wrong teacher: KL=0.000; correctness not established")

covers_teacher forward: 0.051 reverse: 0.084
dumps_on_never forward: 0.218 reverse: 0.436
choose objective from evaluation, not slogan
drop one mode: forward=inf; reverse=0.693
matched wrong teacher: KL=0.000; correctness not established

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 distillation (OPD) 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 of Language Models: Learning from Self-Generated Mistakeshttps://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 includes predominantly synthetic teacher responses and distillation on the student's distribution.^{[2]Reference 2Gemma 2: Improving Open Language Models at a Practical Sizehttps://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf} Keep the sampling rule separate from the divergence: reversing KL on a fixed dataset doesn't make training on-policy. GKD treats sampled sequences as fixed data for its update; it doesn't backpropagate through discrete sampling. MiniLLM uses a different sequence-level optimization procedure, so swapping the arguments to a token-level KL function isn't a complete MiniLLM implementation.^{[9]Reference 9On-policy Distillation of Language Models: Learning from Self-Generated Mistakeshttps://arxiv.org/abs/2306.13649[8]Reference 8MiniLLM: Knowledge Distillation of Large Language Modelshttps://arxiv.org/abs/2306.08543}

This set-based fixture isolates prefix coverage. Both policies answer the same prompt; their partial answers differ. The strings are hand-written, not sampled from a model, and the union represents prefixes that would need teacher scoring rather than calls executed here.

prompt = "How long is this disabled key restorable? Answer: "
fixed_teacher_prefixes = {prompt + "7", prompt + "7 days"}
student_generated_prefixes = {
    prompt + "7",
    prompt + "30",
    prompt + "30 days",
}

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

fixed prefixes: 2
student prefixes needing new teacher scores: ['How long is this disabled key restorable? Answer: 30', 'How long is this disabled key restorable? Answer: 30 days']
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.

Current trainer choices

Checked October 10, 2026: TRL documents DistillationTrainer for student-generated completions and a chunked full-vocabulary generalized Jensen-Shannon loss.^{[10]Reference 10TRL Distillation Trainerhttps://huggingface.co/docs/trl/distillation_trainer} Chunking reduces the peak logit tensor; it doesn't remove teacher scoring or all memory costs. Its agent recipe masks tool-result tokens. The older GKD guide now uses trl.experimental.gkd, with lmbda controlling student-data collection and beta controlling the divergence.^{[11]Reference 11TRL Generalized Knowledge Distillation Trainerhttps://huggingface.co/docs/trl/gkd_trainer} Those are separate choices, and that beta isn't the earlier teacher/task mixing weight $\alpha$. Pin the library version and consult its matching guide before adopting either API. These guides were read; neither trainer was run here.

The current DistillationTrainer also constrains adapter placement. Its loss reads lm_head.weight directly and invokes the student backbone, so it rejects LoRA adapters on lm_head and prompt-learning methods that this path would bypass. To train the output head, its guide specifies modules_to_save=["lm_head"] instead.^{[10]Reference 10TRL Distillation Trainerhttps://huggingface.co/docs/trl/distillation_trainer} A configuration accepted by ordinary SFT isn't automatically a valid configuration for this distillation implementation.

A practical distillation training loop

A typical loop freezes the teacher, sends the same input through both models, and updates only the student. This CPU example uses small positionwise PyTorch networks, not full Transformers. Both start randomly initialized, including the teacher: lowering this loss verifies optimization mechanics, not transferred language ability. The networks see only the current token, not the preceding context.

For this response-only recipe, prompt tokens provide context and aren't supervised targets. A pretraining or full-sequence distillation recipe may intentionally supervise them. Padding isn't a target. PyTorch cross-entropy accepts class-ID ignore_index=-100; F.kl_div has no corresponding mask argument. Select valid next-token positions explicitly for both branches using valid = (shifted_labels != -100).^{[12]Reference 12CrossEntropyLosshttps://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html[13]Reference 13torch.nn.functional.kl_divhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.kl_div.html} For agent trajectories, tool-result text can be context while only student-generated tokens are supervised.

Flattening valid tensors to [valid_tokens, vocabulary] makes batchmean a mean per target token. It divides by the leading dimension, not every non-vocabulary dimension.^{[13]Reference 13torch.nn.functional.kl_divhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.kl_div.html} On [batch, sequence, vocabulary], it instead returns the summed token KL per sequence, even if ignored terms are zeroed. Relative to a token mean, that multiplies the loss by valid_tokens / batch_size; padding and response lengths can change the factor. Choose the intended reduction explicitly. Real causal models also need attention masks for padding; a loss mask doesn't change attention context.

import math
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,
            temperature: float = 3.0, alpha: float = 0.7) -> torch.Tensor:
    if not math.isfinite(temperature) or temperature <= 0 or not math.isfinite(alpha) or not 0 <= alpha <= 1:
        raise ValueError("Need positive finite temperature and alpha in [0, 1].")
    if student_logits.ndim != 3 or student_logits.shape != teacher_logits.shape:
        raise ValueError("Student and teacher need matching [batch, sequence, vocabulary] shapes.")
    if labels.shape != student_logits.shape[:2] or labels.dtype != torch.long:
        raise ValueError("Labels must be int64 [batch, sequence].")
    if student_logits.size(-1) == 0:
        raise ValueError("Vocabulary cannot be empty.")
    shifted_labels = labels[:, 1:]
    valid = shifted_labels != -100
    if not bool(valid.any()):
        raise ValueError("No supervised next-token positions remain; skip this batch.")
    labels_flat = shifted_labels[valid]
    if bool(((labels_flat < 0) | (labels_flat >= student_logits.size(-1))).any()):
        raise ValueError("Supervised labels must be valid token IDs.")
    student_flat = student_logits[:, :-1, :][valid]
    teacher_flat = teacher_logits[:, :-1, :].detach()[valid]
    if student_flat.dtype in (torch.float16, torch.bfloat16):
        student_flat = student_flat.float()
    teacher_flat = teacher_flat.to(student_flat.dtype)
    if not bool(torch.isfinite(student_flat).all() and torch.isfinite(teacher_flat).all()):
        raise ValueError("Supervised logits must be finite.")
    teacher_log_probs = F.log_softmax(teacher_flat / temperature, dim=-1)
    student_log_probs = F.log_softmax(student_flat / temperature, dim=-1)
    soft_loss = F.kl_div(student_log_probs, teacher_log_probs,
                         log_target=True, 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)
torch.set_num_threads(1)
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()
labels[:, :2] = -100  # Prompt tokens aren't supervised targets.
labels[1, -1] = -100  # One illustrative padding target.
optimizer = torch.optim.AdamW(student.parameters(), lr=0.05)

with torch.no_grad():
    teacher_logits = teacher(input_ids)
teacher_before = {name: value.clone() for name, value in teacher.state_dict().items()}
student_before = {name: value.clone() for name, value in student.state_dict().items()}
student.train()

student_before_logits = student(input_ids)
before = kd_loss(student_before_logits, 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)

assert all(torch.equal(value, teacher_before[name]) for name, value in teacher.state_dict().items())
assert all(parameter.grad is None for parameter in teacher.parameters())
assert any(not torch.equal(value, student_before[name]) for name, value in student.state_dict().items())
print("before:", round(float(before.detach()), 4))
print("after:", round(float(after.detach()), 4))
print("has_grad:", has_grad)
print("after_is_finite:", bool(torch.isfinite(after)))
print("improved:", bool(after < before))
print("teacher unchanged:", True)

# Ignored target positions cannot change this response-only loss.
valid = labels[:, 1:] != -100
ignored = ~valid
changed_teacher = teacher_logits.clone()
changed_teacher[:, :-1, :][ignored] = 1000.0
assert torch.allclose(before.detach(), kd_loss(student_before_logits, changed_teacher, labels))
soft_mean = kd_loss(student_before_logits, teacher_logits, labels, alpha=1.0)
terms = F.kl_div(
    F.log_softmax(student_before_logits[:, :-1, :] / 3, dim=-1),
    F.log_softmax(teacher_logits[:, :-1, :] / 3, dim=-1),
    log_target=True, reduction="none",
).sum(-1) * 9
sequence_mean = terms.masked_fill(ignored, 0).sum() / labels.size(0)
factor = valid.sum().item() / labels.size(0)
assert torch.allclose(sequence_mean, soft_mean * factor)
try:
    kd_loss(student_before_logits, teacher_logits, torch.full_like(labels, -100))
except ValueError as error:
    assert "No supervised" in str(error)
else:
    raise AssertionError("An all-ignored batch must not become a silent optimizer step")
print("ignored-target perturbation leaves loss unchanged=True")
print(f"sequence batchmean / valid-token mean={factor:.1f}")
print("all-ignored batch rejected=True")

before: 1.1329
after: 0.952
has_grad: True
after_is_finite: True
improved: True
teacher unchanged: True
ignored-target perturbation leaves loss unchanged=True
sequence batchmean / valid-token mean=2.5
all-ignored batch rejected=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.

The masked example weights each valid token equally, so longer responses contribute more terms. Per-sequence averaging is a different objective. For distributed training or gradient accumulation, a global token-mean objective needs the global valid-token count, accounting for any gradient averaging. An unweighted mean of unequal microbatch means changes the objective. If a batch has no targets, skip the optimizer step too: AdamW weight decay can change weights with a zero-valued gradient tensor. Parameters with grad=None are skipped by PyTorch AdamW.^{[14]Reference 14PyTorch 2.14: AdamWhttps://docs.pytorch.org/docs/2.14/generated/torch.optim.AdamW.html}

import math

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)

# Renormalizing a stored subset hides its total probability mass.
teacher, student = [0.60, 0.30, 0.10], [0.30, 0.15, 0.55]
def forward_kl(q, p):
    return sum(qi * math.log(qi / pi) for qi, pi in zip(q, p))
q_top = [value / sum(teacher[:2]) for value in teacher[:2]]
p_top = [value / sum(student[:2]) for value in student[:2]]
full, truncated = forward_kl(teacher, student), forward_kl(q_top, p_top)
assert full > 0.45 and math.isclose(truncated, 0.0, abs_tol=1e-12)
print(f"full forward KL={full:.3f}; renormalized top-2 KL={truncated:.3f}")

dense cache GiB: 2980.2
top-k cache GiB: 17.9
storage reduction: 166.7 x
quality must still be evaluated: True
full forward KL=0.453; renormalized top-2 KL=0.000

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.

Top-k storage changes more than the file size. Suppose the teacher is [0.60, 0.30, 0.10] and the student is [0.30, 0.15, 0.55]. Renormalize just the first two entries and both become [2/3, 1/3], giving zero truncated KL. Full forward KL is about 0.453 nats because the student put 55% into the omitted tail. Preserve normalization or tail-mass metadata if the chosen approximation requires it, and document the loss actually optimized. Temperature changes the tail mass too.

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

For response distillation, data coverage, student capacity, and training budget constrain results. None supplies a universal accuracy 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.

1. Seed: Start with a small set of high-quality, human-written prompts (for example, 100 restore-window and key-rotation questions).
2. Expand: Ask the teacher to generate new, diverse variations of those prompts.
3. Generate: Have the teacher answer the 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.

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

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. This wrapper uses a constant fake teacher and a lookup verifier, not an API or a model judge. Lowercasing is appropriate only for this fixture; it can change code or case-sensitive identifiers in real prompts.

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 not normalized or 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.get(prompt) == response  # Unknown prompts aren't approved.

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

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

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.

Split source documents, tasks, or seed families before generating variants, and keep all descendants in the same split. Otherwise a rewritten sibling can enter evaluation even if exact strings don't overlap. Keep final test labels out of generation, filtering, and hyperparameter selection. The next check catches only exact normalized overlap; passing it doesn't certify independence.

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

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.

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.

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.

Compare an unmodified student, a trusted-label SFT baseline, the distilled student, and the teacher on the same held-out examples. Match decoding budgets and report output length, sample counts, and uncertainty. For latency and cost, specify hardware, precision, concurrency, input/output lengths, and percentile; parameter count alone doesn't establish a speedup. The next gate uses assigned numbers only. It neither measures latency nor estimates a real invoice.

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

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