Knowledge Distillation
Teacher-student training — black-box vs logit-matching, why on-policy distillation matters, and TRL's GKDTrainer/DistillationTrainer/MiniLLMTrainer
Knowledge Distillation
TL;DR
Instead of training a small model from scratch on raw text, teach it to imitate an already-capable teacher model. The simplest version trains on the teacher's outputs directly; the more effective modern versions have the student generate its own attempts, which the teacher then scores or corrects — fixing the mismatch between "text the teacher would write" and "the actual situations the student's own generations wander into." For a given compute budget, distillation from a decent teacher beats training from scratch on raw text almost every time.
| Property | Value |
|---|---|
| Level | Intermediate |
| Reading time | ~20 minutes |
| Prerequisites | Training Data & Curricula |
| You will understand | The distillation spectrum from black-box to on-policy, and when each is worth its extra complexity |
The Basic Idea
Take a large, already-capable teacher model and a smaller, untrained (or partially trained) student model. Instead of training the student purely on raw text and hoping it discovers good behavior on its own, train it to imitate the teacher.
This is the same core idea used in classic, non-LLM knowledge distillation (e.g. distilling a large vision model into a small one): the teacher produces outputs, the student is trained to match them, and the student ends up with more capability per parameter than training it on raw labels alone would produce.
For language models, "imitate the teacher" can mean several genuinely different things, and the differences matter a lot in practice.
Black-Box vs Logit-Matching Distillation
Two levels of access to the teacher
Black-box distillation
RecommendedGenerate text from the teacher — questions answered, tasks completed, reasoning written out — and train the student on that text as ordinary supervised data. Works with any teacher, including one only reachable through an API where you never see its internals. Simple to set up, and the default starting point.
Logit-matching (white-box) distillation
Train the student to match the teacher's full output distribution over the vocabulary at each step — not just the one word it happened to pick, but how confident it was in every alternative. Requires access to the teacher's internals (its logits), which rules out API-only teachers, but gives a much richer training signal than a single generated token ever could.
The intuition for why logit-matching helps: a teacher's chosen word is one sample from a whole probability distribution over what it considered saying. Training only on that one word throws away everything about how confident the teacher was, and what else it considered plausible. Matching the full distribution (via a KL-divergence-style loss) transfers much more of what the teacher actually "knows" at each step.
On-Policy vs Off-Policy: The Idea That Matters Most Here
This is the single most important distinction in modern LLM distillation, and it's worth sitting with.
Off-policy distillation trains the student on a fixed, pre-generated set of teacher outputs — text the teacher wrote, once, ahead of time. The student never influences what it's trained on.
On-policy distillation flips this: the student generates its own rollouts, and the teacher then scores or corrects those — the actual text the student produced, including its mistakes.
Why this distinction matters more than it sounds like it should
A student trained only on static teacher text learns to imitate the teacher's writing in situations the teacher chose to write about. But at inference time, the student will find itself in situations of its own making — partway through its own generated reasoning, having made its own small earlier mistakes. Off-policy training never showed it what to do from there, because that exact path never appears in the teacher's own (mistake-free) text. On-policy distillation trains the student to recover from exactly the states its own generation actually visits — which is where it needs the help most.
This is the core idea behind the current generation of TRL's distillation trainers.
TRL's Distillation Trainers
TRL (v1.13.0) ships three related trainers, each a different point on the on-policy/logit-matching spectrum:
| Trainer | Approach |
|---|---|
GKDTrainer | Generalized Knowledge Distillation — on-policy, from "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes." Mixes student-generated and teacher-generated data, and mixes forward/reverse KL objectives. |
DistillationTrainer | Also implements on-policy distillation in the same lineage — the student's own generations get corrected against the teacher. |
MiniLLMTrainer | Reverse-KL-divergence-based, on-policy, optimized via reinforcement learning — from "Knowledge Distillation of Large Language Models," aimed at better precision and quality than standard forward-KL approaches. |
A worked example with GKDTrainer
from trl import GKDConfig, GKDTrainer
trainer = GKDTrainer(
model=student_model,
teacher_model=teacher_model,
args=GKDConfig(
output_dir="out",
lmbda=0.5, # fraction of on-policy (student-generated) data used
beta=0.5, # balance between forward and reverse KL objectives
),
train_dataset=prompts_ds,
)
trainer.train()lmbda controls how much of training uses the student's own generations (on-policy) versus static teacher-generated data (off-policy) — a middle setting mixes both rather than committing fully to one. beta controls the KL objective's direction, which affects whether the student is pushed to cover the full breadth of the teacher's distribution or to concentrate on the teacher's most likely outputs.
You don't need to pick a single "correct" value for lmbda on the first attempt. Because on-policy distillation depends on what the student is currently generating, it's normal to start with a mostly off-policy mix early in training (when the student's generations are still low quality) and shift toward more on-policy data as the student improves.
When Distillation Beats Training From Scratch
Distillation vs. raw pretraining, for a given compute budget
A decent teacher already exists for your domain
RecommendedAlmost always the better use of compute. The teacher has already done the expensive work of learning good behavior; the student's job is the comparatively cheaper one of imitating it well.
No usable teacher exists (a genuinely novel domain or task)
Distillation has nothing to imitate. Training from a strong general pretrained base and fine-tuning on your data is the fallback — still not "training from scratch on raw text with no starting point," which is rarely the right move at small-model compute budgets either.
This is why almost none of the leading current small models were trained purely on undifferentiated raw text — they lean on distillation from a stronger sibling or a frontier model at some stage of their curriculum, exactly as described in Training Data & Curricula.
Concept Checks
Check yourself
Why can a student trained only on static, off-policy teacher text still perform poorly at inference time, even if the teacher's text was high quality?
Because the student encounters its own generation's mistakes during inference — states the teacher's clean, pre-written text never visited. The student was never shown what to do from a state reached by its own error, so it has no learned recovery behavior for exactly the situations where it needs one most. On-policy distillation trains directly on the student's own rollouts to close this gap.
When is black-box distillation the right choice over logit-matching, even though it transfers less information per example?
When the teacher is only reachable through an API and its internal logits are unavailable — logit-matching requires white-box access the black-box case simply doesn't have. It's also a reasonable simpler starting point even with a white-box teacher, since it's easier to set up and still captures much of the value of imitating a stronger model.
In GKDConfig, what does increasing `lmbda` toward 1.0 do, and why might you avoid that early in training?
It shifts training toward using more student-generated (on-policy) data and less static teacher-generated data. Early in training the student's generations are still low quality, so leaning heavily on-policy too soon means training mostly on the student's own poor attempts before it has learned much from the teacher at all — a mixed or lower lmbda early on, increasing as the student improves, is the more common approach.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| Teacher/student | A large, capable model whose behavior a smaller model is trained to imitate |
| Black-box distillation | Train on the teacher's generated text — works with any teacher, even API-only |
| Logit-matching | Train to match the teacher's full output distribution — richer, needs teacher internals |
| Off-policy distillation | Trains on static, pre-generated teacher text |
| On-policy distillation | The student's own generations are scored/corrected by the teacher |
GKDTrainer | TRL's Generalized Knowledge Distillation trainer, mixing on/off-policy and KL direction |
DistillationTrainer / MiniLLMTrainer | Sibling on-policy trainers, the latter using reverse-KL optimized via RL |
| The core reason on-policy wins | It trains recovery behavior for the states the student's own generation actually reaches |
Next
With a well-trained (or well-distilled) base model in hand, the next question is how to adapt it further to your specific task: Fine-Tuning Small Models.
Training Data & Curricula
Why data quality matters more as models shrink, the multi-stage curriculum pattern, and why small models are trained past the compute-optimal point on purpose
Fine-Tuning Small Models
Why full fine-tuning is back on the table at small scale, why catastrophic forgetting hits harder, and the extreme-efficiency PEFT methods worth knowing