Knowledge distillation: method selection¶
Scope: how to pick a distillation method, and the white-box mechanics the choice rests on. Organises the field on two axes (what signal the teacher emits and who generated the data), covers logit KD and the distillation infrastructure (teacher servers, generation buffers, top-k logit truncation) in depth, and routes to the two existing method pages for the leaves they already own: on-policy distillation and reasoning distillation via SFT. A stage inside fine-tuning and post-training; the trainers live in TRL.
The TRL snippets below are unexecuted reference templates pinned to TRL v1.8.0 (released 2026-07-09). Every trl.experimental.* namespace is explicitly unstable ("APIs may change without notice") and GKD has already left the top-level namespace, so verify import paths against your installed version. The numpy blocks labelled "core math, runnable" were executed with system python3 (numpy 2.4.6) and their asserts pass; the pasted output is real.
What it is¶
Knowledge distillation trains a student model to reproduce a teacher's behaviour. Every method in the field is a choice on two independent axes, and almost all confusion in the literature comes from conflating them.
- Signal: what the teacher emits. Hard (the teacher's sampled text, one token per position, a black-box signal available from any API) or soft (the teacher's full next-token distribution, requiring white-box logit access).
- Data distribution: which states are scored. Offline/off-policy uses a fixed corpus or teacher-generated outputs; on-policy uses student rollouts, so supervision lands on states the student visits.
The 2x2 that results has four leaves, and each is a distinct engineering job:
| Offline / off-policy states | On-policy student states | |
|---|---|---|
| Hard signal (text or accept/reject) | Sequence-level KD (Kim & Rush 2016; DeepSeek-R1-Distill). Generate traces, filter, then SFT. No teacher is needed during student training. | Adjacent method: rejection sampling / expert iteration. The student samples; a teacher or verifier filters; the survivors become SFT data. This is not necessarily KD. |
| Soft signal (logits) | Logit KD (Hinton 2015). The teacher scores a fixed corpus with temperature-softened targets. | GKD / on-policy distillation. The student samples and the teacher scores each token. |
Self-distillation is orthogonal to both axes: it is a statement about who the teacher is, not about signal or data source. The teacher can be a privileged copy of the student (same weights, extra context), an earlier checkpoint of the same lineage, or absent entirely.
A useful correction to the common model: the teacher is often not a bigger model. Across several 2026 training reports, teachers are frequently checkpoints of the same base, the same size as the student, each pushed further in one domain with RL. "What makes them good teachers is specialization rather than scale."4 That reframes distillation from compression into capability integration.
Why use it¶
- Dense signal. Full-logit KD supplies
O(NV)probability values forNtokens and vocabulary sizeV; a top-k approximation suppliesO(Nk), versus an episodic RL reward withO(1)scalar values. Qwen3 puts distillation at roughly 1/10 the GPU hours of RL, with better results.5 - It is how capability gets merged. DeepSeek-V4 trains a per-domain expert (SFT, then GRPO) for each domain, then unifies them into one model by on-policy distillation with a reverse KL loss against the specialists.46 Nemotron 3 Ultra uses "more than ten" specialised teachers the same way.47 This is one reported use alongside model compression.
- It repairs what sequential training breaks. GLM-5 applies a final distillation pass across training stages, with the teacher being an earlier checkpoint of the same lineage, to recover capability that degraded through its sequential RL phases.48
- It buys behaviour without paying for it at inference. Cursor's Composer 2.5 injects a hint describing the desired behaviour into the context, makes the hinted model the teacher for the unhinted one, and pulls the unhinted policy toward its hint-conditioned self with a per-token KL. The behaviour survives; the hint does not need to be in the prompt at inference.410
When to use it (and when not)¶
Pick the leaf, not the library. The decision collapses to two questions: do you have logit access, and can you afford teacher inference inside the training loop.
- No logit access (the teacher is an API): you are restricted to the hard-signal column. Use sequence-level KD: generate traces, filter, SFT. See reasoning distillation via SFT.
- Logit access, but teacher inference in the loop is too expensive: use logit KD. Score a fixed corpus once, cache the (top-k) teacher logits, train offline against them. The teacher never runs during training. This is the cheapest soft-signal option and the right default when the student is close to the teacher's distribution already.
- Logit access and you can run the teacher in the loop: consider on-policy distillation (GKD) when exposure bias on a fixed corpus is limiting quality. Its cost and benefit must be measured against offline KD and RL on the target workload. Depth in on-policy distillation.
- No suitable teacher exists: use RLVR when a ground-truth reward is available. Teacher supervision constrains the training signal, but it is not a strict evaluation ceiling: students can outperform teachers on a downstream metric through regularization, data selection, or architectural differences. Self-distillation can improve generalization or repair forgetting without introducing an external teacher.
- Different tokenizers: this used to be a hard stop. It is not any more. Use GOLD (below), which distills across incompatible vocabularies.
- Do not conflate compression mechanisms. Quantization reduces precision and storage without teacher training. Model merging combines checkpoints but does not inherently compress them. Distillation is appropriate when behavior transfer or a smaller student architecture is the objective.
Architecture¶
The axes route the method choice; the resulting training systems are not otherwise identical. Hard-signal methods use a token loss on selected text, while soft-signal methods add teacher inference or cached logits and a divergence objective.
flowchart TB
P["Prompt corpus"] --> SRC{"Who generates the tokens?"}
SRC -->|"teacher (off-policy)"| TG["Teacher rollouts"]
SRC -->|"student (on-policy)"| SG["Student rollouts"]
TG --> SIG{"What does the teacher emit?"}
SG --> SIG
SIG -->|"hard: text only"| HARD["Sequence-level KD<br/>(SFT on traces)"]
SIG -->|"soft: full logits"| SOFT["Divergence loss<br/>(temperature, beta, top-k)"]
HARD --> STU["Student update"]
SOFT --> STU
subgraph INFRA["Distillation infrastructure (soft signal only)"]
TSRV["Teacher: vLLM server or co-resident"]
BUF["Generation buffer<br/>(batch prompts across accum steps)"]
TOPK["Top-k logprobs + tail bucket"]
TSRV --> BUF --> TOPK
end
TOPK -.->|"teacher logprobs"| SOFTLogit KD: the white-box mechanism¶
Hinton's original formulation trains the student against the teacher's temperature-softened distribution.1 Raising the softmax temperature T flattens both distributions, exposing the "dark knowledge" in the teacher's relative rankings of wrong answers, which is where most of the information lives. The subtlety that trips people: softening also shrinks the gradient, so the standard loss multiplies the soft term by T^2 to keep its magnitude comparable to the hard-label term.
This block asserts the three properties the loss depends on: T=1 recovers the plain softmax, raising T strictly increases entropy without leaving the simplex, and the gradient norm decays as 1/T^2 so that T^2 * grad stays roughly constant. The adversarial case is the one people actually hit: omit the T^2 factor and a T=4 run silently trains with a gradient about 16x weaker, which is an effective learning-rate change, not a loss-shape change. Run: python3 hinton_temperature.py.
# hinton_temperature.py - core math, runnable (numpy only).
# Why the classic KD soft loss carries a T^2 factor (Hinton et al. 2015, sec 2.1).
import numpy as np
def softmax(z, T=1.0):
z = z / T
z = z - z.max()
e = np.exp(z)
return e / e.sum()
def entropy(p):
return float(-np.sum(p * np.log(p)))
def soft_grad(zs, zt, T):
"""d/d zs of KL(teacher_T || student_T). Analytic: (student_T - teacher_T)/T."""
return (softmax(zs, T) - softmax(zt, T)) / T
teacher_logits = np.array([4.0, 2.0, 1.0, 0.5, -1.0])
student_logits = np.array([3.0, 2.5, 0.0, 1.0, -0.5])
# 1) T=1 recovers the plain softmax exactly (the limit the loss must reduce to).
assert np.allclose(softmax(teacher_logits, 1.0), softmax(teacher_logits)), "T=1 must be plain softmax"
# 2) Softening: raising T strictly increases entropy (flattens toward uniform).
ents = [entropy(softmax(teacher_logits, T)) for T in (1.0, 2.0, 4.0, 8.0, 16.0)]
assert all(b > a for a, b in zip(ents, ents[1:])), ents
assert ents[-1] < np.log(len(teacher_logits)), "must stay below the uniform bound"
# 3) The gradient scale falls off as 1/T^2, so T^2 * grad stays O(1). This is the
# whole reason the standard KD loss multiplies the soft term by T^2.
rows = []
for T in (1.0, 2.0, 4.0, 8.0, 16.0, 32.0):
g = float(np.linalg.norm(soft_grad(student_logits, teacher_logits, T)))
rows.append((T, g, g * T**2))
for a, b in zip(rows, rows[1:]): # each doubling of T cuts |grad| ~4x
ratio = a[1] / b[1]
assert 3.0 < ratio < 5.0, (a, b, ratio)
compensated = [r[2] for r in rows[2:]] # T^2-compensated norms converge
assert max(compensated) / min(compensated) < 1.25, compensated
# 4) Adversarial: WITHOUT the T^2 factor, a T=4 soft loss delivers a gradient ~16x smaller
# than T=1, silently shrinking the effective learning rate on the distillation term.
assert rows[0][1] / rows[2][1] > 10.0, "un-compensated T=4 must be >10x weaker than T=1"
# 5) Boundary: T -> large flattens but never leaves the simplex.
p_hot = softmax(teacher_logits, 1e6)
assert np.isclose(p_hot.sum(), 1.0) and np.allclose(p_hot, 1.0 / len(teacher_logits), atol=1e-4)
print("T entropy |grad| T^2*|grad|")
for (T, g, gc), e in zip(rows, [entropy(softmax(teacher_logits, T)) for T, _, _ in rows]):
print("%-5s %.4f %.6f %.6f" % (T, e, g, gc))
print("un-compensated T=4 gradient is %.1fx weaker than T=1" % (rows[0][1] / rows[2][1]))
print("temperature scaling: PASS")
T entropy |grad| T^2*|grad|
1.0 0.6583 0.353126 0.353126
2.0 1.2697 0.097034 0.388137
4.0 1.5194 0.022481 0.359697
8.0 1.5871 0.005345 0.342072
16.0 1.6039 0.001307 0.334601
32.0 1.6081 0.000324 0.331412
un-compensated T=4 gradient is 15.7x weaker than T=1
temperature scaling: PASS
Note the default temperatures differ between TRL's two soft-signal trainers: GKDConfig.temperature=0.9, DistillationConfig.temperature=1.0. Both apply T to the sampling and the loss, so changing it moves two things at once.
The beta contract, and a discontinuity in it¶
beta is frequently inverted in distillation code and documentation. In TRL the contract is: In TRL:
beta = 0.0is forward KL,KL(teacher || student): mass-covering. The student spreads probability to cover everything the teacher does.beta = 1.0is reverse KL,KL(student || teacher): mode-seeking. The student concentrates on the teacher's dominant modes. This is MiniLLM's choice, Thinking Machines' choice, and DeepSeek-V4's choice.beta = 0.5is the symmetric JSD, and isGKDConfig's default.
The reason this gets inverted so often is a genuine API trap: gkd_trainer.py implements the endpoints with F.kl_div, and PyTorch's F.kl_div(input, target) computes KL(target || input), so the arguments read backwards relative to the mathematical definition. TRL's own source carries a comment saying exactly that. GKDConfig's docstring only says the loss is "the KL divergence" at beta=0.0 and "the Inverse KL Divergence" at beta=1.0, without ever saying which of those is forward and which is reverse; DistillationConfig states it unambiguously in prose as forward KL / reverse KL, and is the authority to quote.3
A finding, and its provenance. TRL hard-codes the two endpoints but computes the interior as a mixture:
# trl/experimental/gkd/gkd_trainer.py (v1.8.0), generalized_jsd_loss.
# Pseudocode, not runnable: the real branch calls F.kl_div, whose arguments invert.
if beta == 0: jsd = KL(teacher || student) # full magnitude
elif beta == 1: jsd = KL(student || teacher) # full magnitude
else: m = (1 - beta) * student + beta * teacher
jsd = beta * KL(teacher || m) + (1 - beta) * KL(student || m) # NOT rescaled
The interior expression is not divided by beta * (1 - beta), so as beta approaches either endpoint the mixture m collapses onto one of the two distributions and the whole loss decays toward zero, while the endpoints themselves return the full un-scaled KL. The result is a jump discontinuity in loss magnitude at both ends of the range. Sweeping beta therefore rescales the gradient by roughly beta * (1 - beta), which changes the effective learning rate on the distillation term by orders of magnitude.
Get that constant right, because it is symmetric, and the folk version of it ("roughly beta") inverts the intuition on exactly the side practitioners use. beta * (1 - beta) peaks at beta=0.5 (a quarter of the KL) and collapses toward zero at both ends. A beta=0.001 run is not "almost forward KL"; it is forward KL at about 1/1000th the learning rate. And a beta=0.95 run is not "almost reverse KL" at 95% of full magnitude; it carries about 5% of it, the same near-zero magnitude as beta=0.05. The executed sweep below prints exactly that.
Honesty about this claim: it is a KB finding, not a cited one. It is derived from reading TRL v1.8.0's source together with PyTorch's documented kl_div semantics, and reproduced independently in numpy below. It was not executed against torch (torch is not installed in this environment) and it appears in no paper or vendor doc. Treat it as a strong prediction to verify on your own install before you act on it, not as a measured result.
# gkd_beta_discontinuity.py - core math, runnable (numpy only).
# Reimplements TRL generalized_jsd_loss (trl/experimental/gkd/gkd_trainer.py, v1.8.0) in numpy.
# F.kl_div(input, target) computes KL(target || input): the arguments invert.
import numpy as np
def kl(p, q): # KL(p || q), the standard definition
return float(np.sum(p * (np.log(p) - np.log(q))))
def trl_jsd(student, teacher, beta):
"""TRL's branch structure, verbatim. Endpoints are hard-coded; the interior is a mixture."""
if beta == 0: # F.kl_div(student_lp, teacher_lp) -> KL(teacher || student)
return kl(teacher, student) # forward KL, mass-covering
if beta == 1: # F.kl_div(teacher_lp, student_lp) -> KL(student || teacher)
return kl(student, teacher) # reverse KL, mode-seeking
m = (1.0 - beta) * student + beta * teacher # mixture
return beta * kl(teacher, m) + (1.0 - beta) * kl(student, m) # NOT rescaled by 1/(beta(1-beta))
student = np.array([0.60, 0.25, 0.15])
teacher = np.array([0.30, 0.45, 0.25])
fwd, rev = kl(teacher, student), kl(student, teacher)
# 1) Endpoint identity: beta=0 is EXACTLY forward KL, beta=1 is EXACTLY reverse KL.
assert np.isclose(trl_jsd(student, teacher, 0.0), fwd), "beta=0 must be KL(teacher||student)"
assert np.isclose(trl_jsd(student, teacher, 1.0), rev), "beta=1 must be KL(student||teacher)"
# 2) Adversarial: the two endpoints are NOT the same number (the forward/reverse mix-up).
assert not np.isclose(fwd, rev), "forward and reverse KL must differ"
# 3) The interior does NOT approach the endpoints: it decays to zero on both sides.
assert trl_jsd(student, teacher, 1e-3) < 0.01 * fwd, "interior must collapse as beta->0"
assert trl_jsd(student, teacher, 1 - 1e-3) < 0.02 * rev, "interior must collapse as beta->1"
# 4) Quantify the jump: the one-sided limits are ~0 while the endpoints are the full KL.
jump_lo = fwd - trl_jsd(student, teacher, 1e-6)
jump_hi = rev - trl_jsd(student, teacher, 1 - 1e-6)
assert jump_lo > 0.99 * fwd and jump_hi > 0.99 * rev, "jump must be nearly the whole KL"
# 5) Boundary: every value is finite and non-negative across the sweep.
assert all(np.isfinite(trl_jsd(student, teacher, b)) and trl_jsd(student, teacher, b) >= 0
for b in np.linspace(0.0, 1.0, 21))
# 6) THE SCALING CONSTANT IS beta*(1-beta), NOT beta. Dividing the interior loss by
# beta*(1-beta) restores an O(1) quantity that interpolates fwd -> rev; dividing by
# beta does not. The factor is SYMMETRIC, so beta=0.95 carries roughly the same tiny
# magnitude as beta=0.05, not 95% of the full KL.
renorm = [trl_jsd(student, teacher, b) / (b * (1 - b)) for b in (0.05, 0.2, 0.5, 0.8, 0.95)]
assert min(renorm) > 0.99 * min(fwd, rev) and max(renorm) < 1.01 * max(fwd, rev), renorm
assert abs(trl_jsd(student, teacher, 0.95) / trl_jsd(student, teacher, 0.05) - 1.0) < 0.10, \
"beta=0.95 and beta=0.05 must have near-EQUAL magnitude (the symmetric factor)"
# 7) Adversarial: the folk model "roughly beta" fails on the reverse-KL side by ~20x.
assert trl_jsd(student, teacher, 0.95) < 0.10 * (0.95 * rev), "'roughly beta' overstates beta=0.95"
print("forward KL(teacher||student) = %.6f reverse KL(student||teacher) = %.6f" % (fwd, rev))
print("beta TRL loss ratio to beta=0 beta*(1-beta) loss / (beta*(1-beta))")
for b in (0.0, 0.001, 0.05, 0.5, 0.95, 0.99, 1.0):
v = trl_jsd(student, teacher, b)
if 0.0 < b < 1.0:
print("%-8s %.6f %.4f %.4f %.6f" % (b, v, v / fwd, b * (1 - b), v / (b * (1 - b))))
else:
print("%-8s %.6f %.4f %-14s %s" % (b, v, v / fwd, "n/a (endpoint)", "n/a (endpoint)"))
print("jump at beta=0: %.6f (%.1f%% of the loss vanishes)" % (jump_lo, 100 * jump_lo / fwd))
print("jump at beta=1: %.6f (%.1f%% of the loss vanishes)" % (jump_hi, 100 * jump_hi / rev))
print("the last column is O(1) across the sweep: the factor is beta*(1-beta), not beta")
print("beta discontinuity: PASS")
forward KL(teacher||student) = 0.184266 reverse KL(student||teacher) = 0.192318
beta TRL loss ratio to beta=0 beta*(1-beta) loss / (beta*(1-beta))
0.0 0.184266 1.0000 n/a (endpoint) n/a (endpoint)
0.001 0.000184 0.0010 0.0010 0.184262
0.05 0.008744 0.0475 0.0475 0.184093
0.5 0.046288 0.2512 0.2500 0.185153
0.95 0.009085 0.0493 0.0475 0.191257
0.99 0.001902 0.0103 0.0099 0.192098
1.0 0.192318 1.0437 n/a (endpoint) n/a (endpoint)
jump at beta=0: 0.184266 (100.0% of the loss vanishes)
jump at beta=1: 0.192318 (100.0% of the loss vanishes)
the last column is O(1) across the sweep: the factor is beta*(1-beta), not beta
beta discontinuity: PASS
The operational consequence: do not sweep beta on a fixed learning rate and read the eval curve as a divergence-shape effect. You are confounding the divergence with an LR sweep. Compare only the endpoints (0.0 and 1.0) against each other, or renormalise the interior loss by beta * (1 - beta) yourself before comparing.
How to use it¶
TRL v1.8.0 ships six distillation trainers, all under trl.experimental. The KB's older advice named only two.
| Trainer | Namespace | What it is |
|---|---|---|
GKDTrainer |
trl.experimental.gkd |
The original on-policy method. lmbda (on-policy fraction), beta (divergence), temperature. Defaults 0.5 / 0.5 / 0.9. |
DistillationTrainer |
trl.experimental.distillation |
The scalable GKD implementation: teacher server, generation buffer, top-k logits. Defaults lmbda=1.0, beta=1.0, temperature=1.0 (fully on-policy reverse KL out of the box). |
GOLDTrainer |
trl.experimental.gold |
Cross-tokenizer distillation, extending ULD. use_uld_loss, teacher_tokenizer_name_or_path, uld_use_hybrid_loss. |
MiniLLMTrainer |
trl.experimental.minillm |
Reverse-KLD with long-horizon credit assignment: rkl_advantage, single_step_decomposition, gamma. |
SDFTTrainer |
trl.experimental.sdft |
Privileged-context self-distillation: the teacher is the same network shown a privileged_context the student never sees. |
SSDTrainer |
trl.experimental.ssd |
Simple Self-Distillation: self-distillation with no teacher model at all. |
The lmbda semantics are worth stating precisely, because the natural reading is wrong. The trainer executes if random.random() <= self.lmbda: once per training step, and on success swaps the whole batch to student-generated tokens. It is a per-step Bernoulli coin-flip over batches, not a per-example fraction within a batch. DistillationConfig words it correctly: each gradient-accumulation slice is randomly assigned as on-policy or off-policy.
# distillation_trainer.py - REFERENCE TEMPLATE, pinned to TRL v1.8.0. Not executed here.
# trl.experimental APIs change without notice: verify the import path on your install.
from trl.experimental.distillation import DistillationConfig, DistillationTrainer
trainer = DistillationTrainer(
model="Qwen/Qwen3-1.7B", # student
teacher_model="Qwen/Qwen3-8B", # frozen; scores the student's tokens
args=DistillationConfig(
lmbda=1.0, # per-step Bernoulli: 1.0 = every step is student-generated
beta=1.0, # 1.0 = reverse KL (mode-seeking). 0.0 = forward KL. 0.5 = JSD.
temperature=1.0,
loss_top_k=1, # top-k logit truncation; 0 = exact full vocabulary (local teacher only)
loss_add_tail=True, # keep the residual mass as one bucket, so the rows stay normalised
bf16=True,
),
processing_class=tokenizer,
train_dataset=prompts, # prompts only: the student generates, the teacher grades
)
trainer.train()
For incompatible tokenizers, GKDTrainer now hard-errors on a vocab_size mismatch, and the error message names the fix. Use GOLDTrainer: GOLD (General Online Logit Distillation) extends ULD (Universal Logit Distillation)2 by incrementally decoding both token streams, grouping spans with identical visible text, and merging the teacher's multi-token probabilities through the chain rule, so no completion tokens are dropped. TRL's working example distills a Llama-3.2-1B-Instruct student from a Qwen2.5-0.5B-Instruct teacher, which also demonstrates that the teacher need not be larger.
# gold_trainer.py - REFERENCE TEMPLATE, pinned to TRL v1.8.0. Not executed here.
from trl.experimental.gold import GOLDConfig, GOLDTrainer
trainer = GOLDTrainer(
model="meta-llama/Llama-3.2-1B-Instruct", # student
teacher_model="Qwen/Qwen2.5-0.5B-Instruct", # DIFFERENT tokenizer, and smaller
args=GOLDConfig(
output_dir="gold-model",
use_uld_loss=True, # the cross-tokenizer path
teacher_tokenizer_name_or_path="Qwen/Qwen2.5-0.5B-Instruct",
uld_use_hybrid_loss=True, # exact-match vocab compared directly, ULD for the rest
),
)
trainer.train()
How to develop with it¶
- Start at an endpoint, not in the interior. Because of the discontinuity above,
beta=1.0(reverse KL) andbeta=0.0(forward KL) are the only two settings whose loss magnitudes are directly comparable. Reverse KL is what DeepSeek-V4, MiniLLM, and Thinking Machines all use; start there. - Choose the teacher before tuning the loss. Teacher calibration, domain coverage, tokenizer compatibility, and scoring throughput can dominate the loss hyperparameters. Compare specialized and general teachers empirically; on-policy versus offline data is a separate variable.
- Keep an off-policy control.
lmbda=0.0is supervised distillation on teacher tokens. Retain it as a regression anchor so you can always prove the on-policy contribution is still positive. - Self-distillation needs no second model.
SDFTTrainertakes aprivileged_contextfield per example (a hint, a demonstration, or the answer) and makes the context-conditioned model its own teacher: this is exactly the Cursor Composer 2.5 recipe.1110SSDTrainerdrops the teacher entirely.12 Both can improve calibration, generalization, or retention without a separate teacher checkpoint. - Use multiple teachers when specialization is the objective. MiMo-V2-Flash names it MOPD (Multi-Teacher On-Policy Distillation), studied in its own paper.9 One domain expert per domain, all grading the same student rollouts.
How to maintain it¶
- Pin TRL and re-verify import paths in CI.
trl.experimentalis documented as unstable, and GKD has already moved out of the top-level namespace once. A minor bump can relocate every symbol in the table above. - Re-run the numpy blocks on upgrade. They are version-independent and pin down the divergence identities; if TRL's
generalized_jsd_lossbranch structure changes, the discontinuity finding above needs re-deriving against the new source. - Track the student-minus-teacher eval gap. Once the student approaches the teacher, distillation has given what it can. The next gain must come from a stronger teacher, more teachers, or RLVR.
- Watch output diversity, not just loss. A falling reverse-KL loss with collapsing generation diversity is mode collapse, not convergence.
How to run it in production¶
DistillationTrainer exists because the naive shape (teacher weights co-resident with the student, full-vocabulary logits, one generation call per microbatch) does not scale. Three mechanisms address it.
- Generation buffer. Decouples the training microbatch size from the generation batch size, so vLLM can batch many prompts into a single call across gradient-accumulation steps. TRL's doc claims this "alone can speed up training by up to 40x".
- Teacher server.
use_teacher_server=Trueplusteacher_model_server_url="http://teacher-host:8000"moves the teacher onto an external vLLM server, so it never has to fit on the training GPUs. TRL cites 100B+ teachers as the motivating case. This is the mechanism that lets a small student distill from a very large teacher without co-locating both sets of weights. - Binary-encoded logprob payloads. Log-probabilities are packed into base64-encoded NumPy arrays instead of nested JSON lists, which TRL says shrinks transfer payloads by "~5x".
Honesty about those two numbers: the "up to 40x" and "~5x" are vendor documentation claims with no published benchmark, methodology, or baseline. Cite them as claimed, never as measured, and benchmark on your own workload before you size a cluster around them.
The teacher-server path imposes hard constraints that will reject your config at construction time:
| Constraint | Rule |
|---|---|
beta == 0 (forward KL) |
requires loss_top_k > 0. The server only exposes top-k logprobs, so exact full-vocabulary loss is impossible. |
beta > 0 (any reverse/JSD) |
requires loss_top_k == 1 exactly. |
reverse_kl_top_1_mode="argmax" |
unsupported with a teacher server. "sampled" only. |
| Liger kernels | unsupported with a teacher server. |
loss_top_k = 0 |
exact, full-vocabulary, local teacher only. |
Top-k truncation and the tail bucket¶
loss_top_k (default 1) restricts the loss to the top-k tokens. The support rule follows the divergence: teacher top-k for forward KL, student top-k for reverse KL, the union for mixed JSD. loss_add_tail (default True) appends one extra bucket holding all the residual probability mass outside the top-k, which is what keeps both rows on the simplex.
This block quantifies both branches implemented by TRL v1.8.0. With loss_add_tail=True, a residual bucket preserves probability mass outside the selected support and converges toward the full-vocabulary loss as k grows. With loss_add_tail=False, TRL renormalizes each retained row. That remains a valid divergence, but at k=1 both rows become [1] and the loss is identically zero, discarding all information about how much probability either model placed on that token. The final assert covers a separate numerical hazard: an exact zero in a teacher distribution makes an unguarded reverse KL infinite. Run: python3 topk_tail_bucket.py.
# topk_tail_bucket.py - core math, runnable (numpy only).
# DistillationConfig.loss_top_k restricts the loss support; loss_add_tail=True (the default)
# appends one bucket holding the residual mass. What each choice costs, plus the zero-mass guard.
import numpy as np
def kl(p, q):
return float(np.sum(p * (np.log(p) - np.log(q))))
def topk_rows(teacher, student, k, add_tail):
"""Restrict both rows to the teacher's top-k: the forward-KL support rule."""
idx = np.argsort(teacher)[::-1][:k]
t, s = teacher[idx], student[idx]
if add_tail: # one bucket for everything outside top-k
return np.append(t, 1.0 - t.sum()), np.append(s, 1.0 - s.sum())
return t / t.sum(), s / s.sum() # TRL renormalizes retained support
rng = np.random.default_rng(7)
V = 128
teacher = np.exp(-0.35 * np.arange(V)); teacher /= teacher.sum() # peaked, like a real next-token row
student = teacher * rng.uniform(0.05, 0.30, V) # a student that has over-concentrated
student[0] = 0.0; student[0] = 1.0 - student.sum() # ... onto the teacher's argmax
exact = kl(teacher, student) # loss_top_k=0: full vocab, exact
print("teacher top-1 mass %.4f | student top-1 mass %.4f" % (teacher[0], student[0]))
print("full-vocab forward KL (loss_top_k=0, exact) = %.6f" % exact)
print("k add_tail sum(teacher row) loss abs error")
for k in (1, 8, 32):
for add_tail in (True, False):
t, s = topk_rows(teacher, student, k, add_tail)
loss = kl(t, s)
print("%-4d %-8s %.6f %+.6f %.6f"
% (k, str(add_tail), t.sum(), loss, abs(loss - exact)))
# 1) The tail bucket keeps BOTH rows on the simplex. That is what makes the loss a divergence.
for k in (1, 8, 32):
t, s = topk_rows(teacher, student, k, True)
assert np.isclose(t.sum(), 1.0) and np.isclose(s.sum(), 1.0), k
assert kl(t, s) >= -1e-12, "a normalised divergence can never be negative"
# 2) With the tail bucket, widening the support converges to the exact full-vocab loss.
errs = [abs(kl(*topk_rows(teacher, student, k, True)) - exact) for k in (1, 8, 32)]
assert errs[0] > errs[1] > errs[2], errs
assert errs[2] < 0.01 * exact, "top-32 must be within 1% of exact"
# 3) ADVERSARIAL: without a tail, TRL renormalizes the retained support. At k=1 both
# distributions become [1], so the divergence is exactly zero even though the original
# top-1 masses differ by 0.56. The branch is numerically valid but information-free.
t1, s1 = topk_rows(teacher, student, 1, False)
assert np.array_equal(t1, np.array([1.0])) and np.array_equal(s1, np.array([1.0]))
assert kl(t1, s1) == 0.0
assert abs(teacher[0] - student[0]) > 0.5
# 4) ADVERSARIAL: a zero-mass teacher token under reverse KL must not silently emit inf/NaN.
# Truncation and quantisation both zero out tokens the student may still sample.
teacher_gap = teacher.copy(); teacher_gap[3] = 0.0
with np.errstate(divide="ignore", invalid="ignore"):
naive = kl(student, teacher_gap) # KL(student || teacher): log(0) -> -inf
assert not np.isfinite(naive), "the unguarded path really does blow up"
floored = np.clip(teacher_gap, 1e-8, None); floored /= floored.sum()
guarded = kl(student, floored)
assert np.isfinite(guarded) and guarded > kl(student, teacher), "guarded: finite, and costlier"
print("k=1, add_tail=False: %+.6f (renormalized singleton: no signal)" % kl(t1, s1))
print("zero-mass teacher token -> unguarded %s | floored %.6f" % (naive, guarded))
print("top-k + tail bucket: PASS")
teacher top-1 mass 0.2953 | student top-1 mass 0.8539
full-vocab forward KL (loss_top_k=0, exact) = 0.862134
k add_tail sum(teacher row) loss abs error
1 True 1.000000 +0.795324 0.066810
1 False 1.000000 +0.000000 0.862134
8 True 1.000000 +0.859221 0.002914
8 False 1.000000 +0.855832 0.006302
32 True 1.000000 +0.862132 0.000002
32 False 1.000000 +0.862127 0.000008
k=1, add_tail=False: +0.000000 (renormalized singleton: no signal)
zero-mass teacher token -> unguarded inf | floored 0.756371
top-k + tail bucket: PASS
The default loss_top_k=1 is aggressive: with the default tail bucket, this example differs from the exact forward KL by 0.067 on a loss of 0.862, roughly 8%. On the external teacher-server path, beta > 0 requires loss_top_k == 1. If the loss error matters more than the memory, keep the teacher local and set loss_top_k=0.
Cluster shape¶
- Two models, or one plus a serving pool. Budget memory for the student's training state (optimizer plus activations under FSDP) and a forward-only teacher. If they do not fit together, the teacher server is the answer, not a smaller teacher.
- Teacher throughput can be the bottleneck. The teacher scores every student token in the loop, so an under-provisioned teacher caps trainer step time. Size the teacher pool against the student's rollout rate; an under-provisioned teacher starves the trainer (performance tuning).
- Multi-teacher capacity depends on routing. Domain teachers may use separate pools, shared batching infrastructure, or selective routing rather than every teacher scoring every rollout. Size each teacher path from the tokens it actually scores.
Failure modes¶
- Forward/reverse KL inverted. The
F.kl_div(input, target)argument order computesKL(target || input). Getting this backwards trains a mass-covering student when you wanted a mode-seeking one, and the loss curve looks fine either way.beta=0is forward,beta=1is reverse. - Sweeping
betathrough the interior. The loss magnitude is discontinuous at both endpoints (above), so an interior sweep silently rescales the gradient by roughlybeta * (1 - beta). That factor is symmetric and collapses at both ends, sobeta=0.95carries about 5% of the full magnitude, not 95%. You will read an LR effect as a divergence effect. - Top-k truncation without a tail bucket. TRL renormalizes the retained rows. At
k=1, both become[1]and the loss collapses to zero regardless of the original top-token masses. - Zero-mass teacher tokens. Truncation or quantisation zeroes a token the student still samples; reverse KL returns
inf. Floor the teacher distribution. lmbdamisread as a within-batch fraction. It is a per-step Bernoulli over whole batches. Atlmbda=0.5half your steps are pure off-policy SFT, not half of each batch.- Tokenizer mismatch.
GKDTrainerhard-errors on avocab_sizemismatch. That is the good case (loud, at startup). UseGOLDTrainerrather than forcing vocabularies to align. - Treating the teacher as a strict evaluation ceiling. The teacher limits the direct supervision available, but student metrics can exceed teacher metrics through regularization, data selection, or architecture. Use RLVR when a verifiable reward supplies information the teacher does not.
- Missing
T^2on the soft term. A temperature-scaled KD loss without theT^2factor trains at a fraction of the intended learning rate (about 1/16 atT=4). trl.experimentaldrift. Every trainer here lives in an explicitly unstable namespace. Unpinned, an upgrade breaks the import, not the math.
References¶
- Distilling the Knowledge in a Neural Network (Hinton, Vinyals, Dean 2015; temperature scaling, the
T^2factor): https://arxiv.org/abs/1503.02531 - Sequence-Level Knowledge Distillation (Kim & Rush 2016): https://arxiv.org/abs/1606.07947
- On-Policy Distillation of Language Models (GKD, ICLR 2024): https://arxiv.org/abs/2306.13649
- MiniLLM: On-Policy Distillation of Large Language Models (reverse KLD): https://arxiv.org/abs/2306.08543
- Towards Cross-Tokenizer Distillation: the Universal Logit Distillation Loss (ULD): https://arxiv.org/abs/2402.12030
- Self-Distillation Enables Continual Learning (SDFT; Shenfeld, Damani, Hübotter, Agrawal): https://arxiv.org/abs/2601.19897 (code: https://github.com/idanshen/Self-Distillation)
- Embarrassingly Simple Self-Distillation Improves Code Generation (SSD): https://arxiv.org/abs/2604.01193
- MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training: https://arxiv.org/abs/2606.30406
- Distillation in 2026 (Sergio Paniego, Hugging Face; the frontier-lab survey this page's evidence section follows): https://huggingface.co/blog/sergiopaniego/distillation-2026
- On-Policy Distillation (Thinking Machines Lab): https://thinkingmachines.ai/blog/on-policy-distillation/
- TRL GKD Trainer: https://huggingface.co/docs/trl/gkd_trainer
- TRL Distillation Trainer (teacher server, generation buffer,
loss_top_k): https://huggingface.co/docs/trl/distillation_trainer - TRL GOLD Trainer (cross-tokenizer): https://huggingface.co/docs/trl/gold_trainer
- DeepSeek-R1 (sequence-level distillation into Qwen/Llama students): https://arxiv.org/abs/2501.12948
- DeepSeek-V4 (per-domain experts unified by on-policy distillation, reverse KL): https://arxiv.org/abs/2606.19348
- Gemma 3 Technical Report (distillation from a large instruction-tuned teacher): https://arxiv.org/abs/2503.19786
- Qwen3 Technical Report (strong-to-weak distillation, ~1/10 the GPU hours of RL): https://arxiv.org/abs/2505.09388
- MiMo-V2-Flash Technical Report (coins MOPD): https://arxiv.org/abs/2601.02780
- GLM-5 (distillation across training stages; teacher is an earlier checkpoint): https://arxiv.org/abs/2602.15763
- Nemotron 3 Ultra (more than ten specialised teachers): https://arxiv.org/abs/2606.15007
- Cursor Composer 2.5 (privileged-context self-distillation): https://cursor.com/blog/composer-2-5
Related: On-policy distillation · Reasoning distillation via SFT · Fine-tuning and post-training · TRL · RLVR · GRPO · SFT/LoRA · RLSD · Synthetic data · Training-data curation · Model merging · Quantization for inference · Speculative decoding · RL libraries · Glossary
-
Hinton, Vinyals, Dean: soften teacher and student with a temperature
T, train on the soft targets, and multiply the soft-target gradient byT^2because it scales as1/T^2. https://arxiv.org/abs/1503.02531 ↩ -
Boizard et al., Universal Logit Distillation: distillation across models with different tokenizers by comparing sorted probability distributions rather than aligned vocabulary indices. TRL's GOLD extends it with incremental decoding and chain-rule merging of multi-token spans. https://arxiv.org/abs/2402.12030 ↩
-
TRL
DistillationConfig(v1.8.0) states the contract unambiguously in prose:beta0.0is "the forward KL divergence",1.0is "the reverse KL divergence",0.5is the standard JSD; andlmbdais the "probability of using on-policy (student-generated) data for each gradient accumulation slice".GKDConfigsays only "the loss is the KL divergence" (beta=0.0) and "the loss is the Inverse KL Divergence" (beta=1.0), naming neither as forward or reverse, which is the source of the widespread inversion. https://huggingface.co/docs/trl/distillation_trainer ↩ -
Paniego (Hugging Face), "Distillation in 2026": the survey of frontier-lab reports this page's evidence follows, including the observation that teachers "are checkpoints of the same base, the same size as the student, each pushed further in a single domain with RL. What makes them good teachers is specialization rather than scale." The post contains no results table, no benchmark numbers, no code, and no TRL parameter names. Its claim that Gemma 4 also uses distillation is explicitly an inference ("we can infer"), not a reported fact. https://huggingface.co/blog/sergiopaniego/distillation-2026 ↩↩↩↩↩
-
Qwen3 Technical Report: strong-to-weak distillation for the smaller models "significantly outperforms reinforcement learning in performance and training efficiency", at roughly 1/10 of the GPU hours of the multi-stage RL pipeline. https://arxiv.org/abs/2505.09388 ↩
-
DeepSeek-V4: each domain gets its own expert (SFT then GRPO), after which "a single unified model is trained through on-policy distillation", the student optimising a reverse KL loss against the specialist teachers. https://arxiv.org/abs/2606.19348 ↩
-
Nemotron 3 Ultra: adopts the multi-teacher form at scale, with more than ten specialised teachers each carrying its own domain pipeline, giving the student dense token-level guidance on its own rollouts. https://arxiv.org/abs/2606.15007 ↩
-
GLM-5: applies distillation across training stages rather than across domains. After its sequential RL phases, a final distillation pass recovers capability degraded along the way, with the teacher being an earlier checkpoint of the same lineage. https://arxiv.org/abs/2602.15763 ↩
-
MiMo-V2-Flash names the multi-teacher form MOPD (Multi-Teacher On-Policy Distillation); it is studied in its own paper. Domain teachers provide a dense token-level signal on whatever the student generates. https://arxiv.org/abs/2606.30406 ↩
-
Cursor Composer 2.5: a hint describing the desired behaviour is injected into the context, and the hinted model becomes the teacher for the same model without the hint. A per-token KL pulls the unhinted policy toward its hint-conditioned self, so the behaviour persists without the hint at inference. https://cursor.com/blog/composer-2-5 ↩↩
-
Shenfeld et al., Self-Distillation Fine-Tuning: the student generates from the plain prompt, and the same network conditioned on a
privileged_context(a demonstration, hint, or feedback) re-scores those tokens. Teacher and student are one network differing only in what they see. Reported to reduce catastrophic forgetting versus SFT. https://arxiv.org/abs/2601.19897 ↩ -
Simple Self-Distillation (SSD): self-distillation for code generation with no separate teacher model. https://arxiv.org/abs/2604.01193 ↩