Model distillation: building smaller, cheaper, specialized models

Model distillation trains a smaller “student” model to reproduce the behavior of a larger “teacher,” so you get most of the teacher’s quality on your task at a fraction of the inference cost. For LLMs this usually means fine-tuning the student on filtered teacher outputs, matching the teacher’s token probabilities, or having the teacher score the student’s own generations.

Distillation is not a separate training paradigm. The machinery is usually the same supervised fine-tuning you would run on human labels. The difference is where the targets come from, and that one change affects data quality, failure modes, cost and what you are legally allowed to do. This article covers the methods from 2015 to on-policy distillation, when each works, and a pipeline you can run.

Where did knowledge distillation come from?

Hinton, Vinyals and Dean (2015) framed the idea. A large model, or an ensemble, produces class probabilities, and a small model is trained to match them as soft targets. Soft targets carry more information than a one-hot label. In the paper’s MNIST example, one image of a 2 might get probability 10⁻⁶ of being a 3 and 10⁻⁹ of being a 7. Those ratios encode which inputs resemble which. Because soft targets have high entropy, the authors note, they give more information per example and less gradient variance than hard labels. The student can often learn from less data at a higher learning rate.

Those tiny probabilities barely affect a cross-entropy loss at normal temperature, so the paper raises the temperature T of the softmax, softmax(z/T), for both teacher and student during training, and uses T = 1 at inference. When true labels are available, the loss is a weighted average of the soft-target cross-entropy at temperature T and the ordinary hard-label cross-entropy. The soft term’s gradients scale as 1/T², so the authors multiply that term by T² to keep the balance stable as T changes.

In its next-token form for language models, the objective looks like this:

import torch
import torch.nn.functional as F


def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.9):
    """Hinton-style KD for next-token prediction.

    student_logits, teacher_logits: [batch, seq, vocab], already aligned with labels,
        from models that share one tokenizer and vocabulary
    labels: [batch, seq], with -100 on positions that should not be trained
    """
    mask = labels != -100
    t = temperature
    # Soft targets: KL(teacher || student) at temperature T, scaled by T^2 so its
    # gradient magnitude stays comparable to the hard-label term.
    soft = F.kl_div(
        F.log_softmax(student_logits[mask] / t, dim=-1),
        F.log_softmax(teacher_logits[mask] / t, dim=-1),
        log_target=True,
        reduction="batchmean",
    ) * (t * t)
    # Hard targets: ordinary cross-entropy on the reference tokens at T = 1.
    hard = F.cross_entropy(student_logits[mask], labels[mask])
    return alpha * soft + (1 - alpha) * hard

Minimizing the KL term is equivalent to the paper’s soft-target cross-entropy, because the two differ only by the teacher’s entropy, which is constant for the student. The default alpha=0.9 follows the paper’s finding that a much lower weight on the hard-label term worked best. Tune both alpha and temperature on your validation set.

EfficientML.ai Lecture 9 - Knowledge Distillation (MIT 6.5940, Fall 2023)Song Han, MIT HAN Lab · YouTube · 2023

What is sequence-level distillation?

Matching token distributions one position at a time ignores the fact that generation is a sequence of dependent choices. Kim and Rush (2016) proposed sequence-level knowledge distillation for machine translation. You run beam search with the teacher to generate a new dataset, then train the student on it with ordinary cross-entropy. Their best student ran 10x faster than its state-of-the-art teacher with little loss in performance. With weight pruning it had 13x fewer parameters at a cost of 0.4 BLEU.

This is the ancestor of most LLM distillation today: generate outputs with a strong model, then fine-tune a smaller one on them. It needs no access to the teacher’s logits, only its text.

How is LLM distillation done today?

Three patterns dominate.

1. Teacher-generated data (off-policy data distillation). Prompt the teacher on task inputs, filter its outputs, and run supervised fine-tuning on the student. DeepSeek’s R1 distilled models are the best-known recent example. DeepSeek fine-tuned open Qwen2.5 and Llama 3 checkpoints at 1.5B, 7B, 8B, 14B, 32B and 70B on reasoning samples generated by DeepSeek-R1. The model card says the Qwen-based models used 800K curated samples. DeepSeek reports that distilling a large model’s reasoning patterns into smaller models worked better than discovering those patterns with RL directly on the small models.

2. Distilling rationales, not just answers. Distilling Step-by-Step (Hsieh et al., 2023) extracts the teacher’s rationales and trains the student to produce both labels and rationales in a multi-task setup. On ANLI, one of their four benchmarks, a 770M T5 model trained this way outperformed few-shot prompted 540B PaLM using 80% of the available data, while standard fine-tuning of the same T5 model didn’t match PaLM even with all of it. How small a student could go varied by task: 220M sufficed on e-SNLI, while CQA needed 11B.

3. On-policy distillation. In off-policy distillation the student only ever learns from the teacher’s trajectories. At inference it has to continue from its own, which may drift into states the teacher never visited. Generalized Knowledge Distillation (Agarwal et al., ICLR 2024) addresses this mismatch: the student generates sequences and the teacher gives token-level feedback on them. GKD also lets you choose the divergence, which helps when the student lacks the capacity to mimic the teacher fully. Thinking Machines’ 2025 write-up of on-policy distillation uses a per-token reverse KL between student and teacher. They describe it as combining on-policy sampling, like RL, with dense per-token feedback, like SFT.

The reported efficiency is striking, but specific to its setup. Thinking Machines cite the Qwen3 technical report: on-policy distillation of a reasoning student reached 74.4% on AIME’24 in about 1,800 GPU-hours, against 67.6% for RL in about 17,920 GPU-hours. In their own experiment, on-policy distillation also recovered most of the instruction-following ability a model had lost during mid-training on internal documents, without erasing the new knowledge. Both results use Qwen-family models with a shared tokenizer, on math reasoning and an internal QA set. Treat them as evidence that the method is worth trying, not as a forecast for your task.

Logit distillation or data distillation?

	Logit (white-box) distillation	Data (black-box) distillation
What the student learns from	Teacher’s full or top-k token probabilities	Teacher’s generated text
Teacher access needed	Logits or log-probs on arbitrary tokens	Text outputs only; an API is enough
Tokenizer	Must match, or needs a cross-vocabulary mapping	Any
Signal per example	Dense: a full distribution per token	Sparse: one sampled sequence
Typical tools	Custom loss (above), TRL GKD trainer	Standard SFT on teacher outputs
Main risk	Compute and memory to run both models	Copying the teacher’s style and errors

On-policy distillation is a white-box method. It needs the teacher’s log-probabilities on tokens the student chose, so in practice it usually means an open-weight teacher that shares the student’s tokenizer. TRL ships this as an experimental trainer. lmbda controls the fraction of on-policy batches, and beta interpolates the generalized Jensen-Shannon divergence between forward KL (0.0) and reverse KL (1.0) (TRL GKD docs):

from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer
from trl.experimental.gkd import GKDConfig, GKDTrainer

STUDENT = "Qwen/Qwen3-4B-Instruct-2507"
TEACHER = "Qwen/Qwen3-30B-A3B-Instruct-2507"  # same tokenizer family as the student

# {"messages": [...]} rows: real task prompts plus a reference assistant reply
tokenizer = AutoTokenizer.from_pretrained(STUDENT)
train = load_dataset("json", data_files="data/train.jsonl", split="train")

args = GKDConfig(
    output_dir="out/student-gkd",
    lmbda=1.0,             # 1.0 = always train on the student's own samples (on-policy)
    beta=1.0,              # 1.0 approximates reverse KL, 0.0 forward KL
    temperature=1.0,       # sampling temperature for student generations
    max_new_tokens=512,
    teacher_model_init_kwargs={"dtype": "bfloat16"},
    model_init_kwargs={"dtype": "bfloat16"},
    learning_rate=1e-4,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    bf16=True,
    report_to="none",
)

trainer = GKDTrainer(
    model=STUDENT,
    teacher_model=TEACHER,
    args=args,
    processing_class=tokenizer,  # required: the trainer does not load it for you
    train_dataset=train,
    peft_config=LoraConfig(r=32, lora_alpha=32, target_modules="all-linear", task_type="CAUSAL_LM"),
)
trainer.train()
trainer.save_model(args.output_dir)

Both Qwen models above are Apache 2.0 and share an identical vocabulary. We ran this script end to end against TRL 1.17.1 with tiny stand-in models only, not the full-size pair. The teacher must fit in memory next to the student, so plan for its weights as well. The trainer lives under trl.experimental, so pin your TRL version.

When does distillation work, and when does it fail?

Distillation works best when:

• The task is narrow and its input distribution is stable. Examples include classification, extraction, routing, structured generation, domain-specific drafting, and a fixed style of reasoning.
• Teacher outputs can be checked. Exact answers, schemas, unit tests or validated rubrics let you keep only correct outputs. You then train on a filtered version of the teacher, not on its raw distribution.
• Volume justifies the effort. The saving is the inference price gap multiplied by traffic. Low-volume tasks rarely repay the data generation, training and evaluation work. How to reduce LLM costs covers the cost model.
• Prompts come from real traffic. The student learns the teacher’s behavior on the inputs you sample. If those differ from production, so will the student.

It fails, or quietly disappoints, when:

• You want general capability. Gudibande et al. (2023) fine-tuned 1.5B–13B models on up to 150M tokens of ChatGPT outputs. Crowd raters judged the results competitive, yet targeted evaluations showed the imitation models “are adept at mimicking ChatGPT’s style but not its factuality”. They closed little of the gap on tasks the imitation data didn’t heavily cover.
• The capacity gap is too large. A student too small for the task learns the format without the substance. Thinking Machines make the same point about off-policy distillation: the student can pick up the teacher’s style and confidence without its accuracy.
• Teacher errors go unfiltered. The student reproduces the teacher’s mistakes and may make them more confidently. Synthetic data has its own failure modes, covered in synthetic data for fine-tuning.
• The evaluation is weak. If you can’t tell a good answer from a fluent one automatically, you can’t filter the training data, and you can’t tell whether the student is good enough.

What do OpenAI, Anthropic and Google terms say about distillation?

This is a summary, not legal advice. Terms change, and enterprise agreements can differ from public terms, so have counsel read the version that governs your account.

• Anthropic. The Commercial Terms (the version effective June 17, 2025) forbid customers to “access the Services to build a competing product or service, including to train competing AI models”. They also forbid reverse engineering or duplicating the services, or supporting a third party’s attempt to do any of these. The same terms say the customer owns its outputs.
• Google. The Gemini API Additional Terms (effective March 23, 2026) state: “You may not use the Services to develop models that compete with the Services (e.g., Gemini API or Google AI Studio).” They also forbid attempts to “reverse engineer, extract or replicate any component of the Services,” including the underlying data or models such as parameter weights.
• OpenAI. Check the current terms that govern your account, including any business or API agreement, before training on outputs. We have not quoted them because we could not verify the current text on OpenAI’s site. Separately, OpenAI’s own fine-tuning guide recommends distilling a large OpenAI model’s outputs into a smaller OpenAI model, though that platform is winding down: it is closed to new users, and eligible existing customers can create jobs until January 6, 2027.

Neither the Anthropic nor the Google clause defines “competing” precisely, and neither is a blanket permission to train a general-purpose model on API outputs. The cleanest path is an open-weight teacher whose license explicitly allows it. DeepSeek’s model card, for example, says the DeepSeek-R1 series allows modifications and derivative works, including distillation for training other LLMs. Distilled models also inherit the licenses of their base checkpoints: the R1 distills built on Llama carry Llama license terms. The trade-offs are covered in open-weight vs proprietary LLMs.

How should you evaluate a distilled student?

Evaluate the student against the teacher and against your current production system, on the same frozen test set, with the same scorer.

1. Score against ground truth, not just agreement with the teacher. Teacher agreement shows how faithfully the student imitates. Ground truth shows whether it is any good. A student can agree with the teacher 95% of the time and still inherit all of its errors.
2. Use a test set the teacher never generated. Prompts and labels for testing should come from real traffic, never from the teacher’s synthetic pool, or you will measure memorization.
3. Slice the results. Report quality by input type, length and difficulty. Distilled students tend to match the teacher on common cases and fall off on rare ones.
4. Check what you didn’t train. Run a small regression set for adjacent behaviors: refusals, instruction following, formatting under unusual inputs.
5. Measure cost and latency under load, at your real request mix, so the saving is a measured number.
6. Use confidence intervals. On a few hundred examples, a two-point gap may be noise. How to measure whether fine-tuning actually improved your model covers sample sizes and paired comparisons.

Ship only if the student meets the quality bar the current system sets.

A practical distillation pipeline

1. Define the task and acceptance bar. Write the output spec, the scorer, and the minimum quality the student must reach on the test set.
2. Freeze evaluation data first. Draw a held-out test set from production, labeled by people or by verified checks.
3. Choose a teacher you’re allowed to use. Prefer an open-weight model with a permissive license and, if you want logit or on-policy methods, a tokenizer shared with the student.
4. Collect prompts from real traffic. Deduplicate them and remove anything that overlaps with the test set.
5. Generate, then filter hard. Sample several teacher outputs per prompt and keep only those that pass your checks. Include rationales if the task involves reasoning.
6. Run SFT on the student with LoRA or full fine-tuning (see the LoRA and QLoRA tutorial).
7. Optionally add on-policy distillation with the teacher as scorer, to fix the errors the student makes on its own samples.
8. Evaluate against the teacher and the production baseline, sliced and with intervals, then canary on live traffic before switching over.
9. Monitor and refresh. When inputs drift, re-collect prompts, regenerate and re-evaluate.

Distillation is one of the most reliable ways to make a narrow, high-volume task cheaper. It is also one of the easiest ways to ship a model that sounds right and is wrong. The difference between the two is the filter and the evaluation, not the training run.