Knowledge distillation for LLMs, explained

Knowledge distillation trains a smaller student model to imitate a larger teacher. For LLMs it comes in two main forms: sequence-level distillation, where the teacher writes responses and the student is fine-tuned on them with ordinary SFT; and token-level (logit) distillation, where at every position of a training sequence the student is trained to match the teacher's full next-token distribution rather than a single correct token. It is how most small open models with strong reasoning were produced, how several model families train their smaller sizes, and the most compute-efficient way to move a capability from a model you cannot afford to serve into one you can. In the post-training stack it usually replaces or augments SFT, and it is followed by preference tuning, RL, or its on-policy refinement, on-policy distillation.

Where it came from#

The idea of compressing a large model or ensemble into a small one by training on its outputs goes back to Caruana and collaborators' work on model compression (Bucilă et al., 2006). Hinton, Vinyals and Dean (2015) gave it the name and the form most people use: train the student on the teacher's softened class probabilities, produced by dividing logits by a temperature $T>1$, because those probabilities carry "dark knowledge" about which wrong answers are nearly right. A teacher that says an image of a 2 is 0.9 "2", 0.09 "3" and 0.01 "7" tells the student much more than the label "2" does.

Kim and Rush (2016) adapted the idea to autoregressive sequence models for machine translation. They showed that word-level distillation works, and introduced sequence-level distillation: train the student on the teacher's beam-search outputs. Their best student ran 10 times faster than its teacher with little loss in BLEU, and sequence-level training seemed to remove the student's need for beam search. DistilBERT (Sanh et al., 2019) applied distillation during pretraining to make BERT 40% smaller while retaining 97% of its language-understanding performance. In the LLM era, distillation became central at both ends of training. Gemma 2 (Gemma Team, 2024) trained its 2B and 9B models with token-level distillation from a larger model instead of plain next-token prediction. DeepSeek-R1 (DeepSeek-AI, 2025) fine-tuned Qwen and Llama models on about 800k samples curated with R1, a purely sequence-level recipe: DeepSeek-R1-Distill-Qwen-7B reached 55.5% on AIME 2024, and the distilled 32B model scored 72.6% against 47.0% for the same base trained with large-scale RL directly. MiniLLM (Gu et al., 2023) and GKD (Agarwal et al., 2023) moved the field toward reverse-KL and on-policy objectives, which have their own article.

The objectives#

Let $p_T(\cdot \mid c)$ and $q_{\theta }(\cdot \mid c)$ be the teacher's and student's next-token distributions over a vocabulary $\mathcal{V}$ given a context $c=(x,y_{<t})$. Let $z^T,z^S\in {\mathbb{R}}^{|\mathcal{V}|}$ be their logits and $\tau$ a temperature, so $p_T^{\tau }=\mathrm{softmax}(z^T/\tau )$ and similarly for the student.

Token-level (logit) distillation

The sequences $y$ can come from a fixed dataset, from the teacher, or (in the on-policy variant) from the student. In practice it is often mixed with the ordinary cross-entropy on the ground-truth token:

The gradient explains both the temperature and the ${\tau }^2$. With respect to the student's logits at one position,

Soft targets replace the one-hot $e_{y_t}$ with the teacher's full distribution, so every logit gets a meaningful target, not just the observed token. Raising $\tau$ flattens both distributions and shrinks the differences between them, which shrinks gradients roughly as $1/{\tau }^2$; multiplying by ${\tau }^2$, as Hinton et al. recommend, keeps the soft term's scale comparable to the hard term's as $\tau$ changes. In the high-temperature limit the loss becomes matching logits (up to their mean) in squared error.

Sequence-level distillation

The ideal sequence-level objective is to match the teacher's distribution over whole responses, ${\min }_{\theta }\mathrm{KL}(p_T(\cdot \mid x)\Vert q_{\theta }(\cdot \mid x))$, which equals $-{\mathbb{E}}_{y\sim p_T}[\log q_{\theta }(y\mid x)]$ plus a constant. Approximating the expectation with teacher samples $\hat{y}\sim p_T(\cdot \mid x)$ (or its beam-search mode, as Kim and Rush did) gives

which is just SFT on teacher-written data. It needs only the teacher's text, not its logits, so it works across different tokenizers and with teachers available only as a text API. Its weakness is information per token: each position supervises one sampled token where token-level KD supervises a distribution.

Which divergence: forward, reverse, or in between

The KL above is the forward KL, $\mathrm{KL}(p_T\Vert q_{\theta })$, which is mass-covering: the student is penalised wherever the teacher has mass that it lacks, so a student too small to represent every teacher mode spreads its probability across all of them, including the gaps between. The reverse KL, $\mathrm{KL}(q_{\theta }\Vert p_T)$, is mode-seeking: the student is penalised for putting mass where the teacher has little, so it concentrates on the modes it can fit. Gu et al. (2023) argue that reverse KL suits generative distillation because it stops the student from overestimating the teacher's low-probability regions, which is where hallucinated or degenerate text lives. Generalised Jensen-Shannon and skew-KL divergences (used in GKD and DistiLLM respectively) interpolate between the two.

We fitted a single discretised Gaussian to a two-mode teacher on 100 outcomes by grid search. The forward-KL fit centred at 45.5 with standard deviation 22, putting 36.9% of its mass in the valley between the modes where the teacher has 4.0%. The reverse-KL fit locked onto the larger mode (mean 30, standard deviation 5) and put 2.9% there. For a language model, the valley is text the teacher would never write.

Worked example: one position, three tokens#

Let a position have three candidate tokens, with teacher logits $z^T=(2,1,0)$ and student logits $z^S=(1,1,0.5)$, and let the ground-truth token be the first. At $\tau =1$ the teacher's distribution is $(0.665,0.245,0.090)$ and the student's is $(0.384,0.384,0.233)$. The forward KL is

and the hard-label cross-entropy is $-\ln 0.384=0.958$. The logit gradients differ in a telling way. The hard loss gives $q-e_1=(-0.616,0.384,0.233)$: push token 1 up, push tokens 2 and 3 down in proportion to their current probability. The soft loss gives $q-p_T=(-0.282,0.139,0.143)$: push token 1 up less, and push token 3 down slightly more than token 2, because the teacher considers token 2 a plausible alternative and token 3 not. That ranking information is absent from the hard label. At $\tau =2$ the distributions become $(0.506,0.307,0.186)$ and $(0.360,0.360,0.280)$, the KL drops to 0.0484, and the gradient $(q-p)/\tau$ shrinks accordingly, which the ${\tau }^2$ factor compensates. All numbers computed in numpy.

A toy run: how much is a soft label worth?#

We simulated distillation with a log-linear teacher over a 64-token vocabulary, $p_T(\cdot \mid c)=\mathrm{softmax}(W_T\varphi (c))$ with 48-dimensional context features, and a smaller student that sees only the first 32 feature dimensions. The student cannot match the teacher exactly: trained on 20,000 contexts with full soft targets, its held-out KL to the teacher bottomed out at 1.022 nats (the teacher's own entropy is 1.524 nats). We then trained students from zero on $N$ contexts with three kinds of target: one hard token sampled from the teacher per context (sequence-level distillation in miniature), the teacher's top-5 probabilities renormalised, and the full teacher distribution. Each run used SGD for 3,000 steps with batch 64, and we report held-out KL at the best checkpoint, which flatters the hard-label runs since it amounts to oracle early stopping.

Full soft targets on 800 contexts (1.322) matched hard labels on 3,200 (1.321): a 4× gain in sample efficiency in this setting. Measured as excess over the capacity floor, at 3,200 contexts hard labels were 0.30 nats above it and soft labels 0.07. Keeping only the top 5 of 64 probabilities retained most of the benefit, which is why storing top-$k$ logits is the standard way to make offline token-level distillation affordable. A second check: with 200 contexts, sampling 64 hard tokens per context and training on their empirical distribution gave 1.872, still worse than one soft label per context (1.829). A soft label is worth many samples.

The algorithm, step by step#

# 1. Data
prompts = collect(task_prompts + general_prompts)                # 50k to 1M+
if sequence_level:
    for x in prompts:
        ys = teacher.sample(x, n=1..8, temperature=0.6..1.0)
        keep = [y for y in ys if passes_filters(x, y)]            # format, verifier, length, dedup
        dataset += [(x, y) for y in keep]
if token_level_offline:
    for (x, y) in dataset:
        logits = teacher.forward(x + y)
        store(topk(logits / tau, k=20..100))                     # ids + log-probs per position

# 2. Train the student
for batch in loader:
    s_logits = student(batch.ids)
    if sequence_level:   loss = masked_cross_entropy(s_logits, batch.labels)
    else:
        t_logp = batch.teacher_topk_logprobs  or  teacher(batch.ids) / tau   # offline or online
        loss = alpha * tau**2 * kl(t_logp, log_softmax(s_logits / tau)) + (1 - alpha) * ce
    loss.backward(); opt.step(); opt.zero_grad()

Data sizes depend on what is being transferred. Narrow task distillation (one extraction or classification task, one domain) works with 10k to 100k prompts. Broad reasoning distillation is larger: DeepSeek-R1's distilled models used about 800k samples. Pretraining-scale distillation, as in Gemma 2, uses the full pretraining token budget. Prompt coverage matters more than the number of responses per prompt: the student only learns the teacher's behaviour on the kinds of inputs it sees. For sequence-level distillation with a verifier (maths answers, unit tests), filtering to correct teacher outputs is cheap and helps; this is rejection-sampling fine-tuning with the teacher as sampler.

Hyperparameters that matter#

The last row reflects Mirzadeh et al. (2019), who showed on image classification that student performance degrades when the gap between student and teacher is large, and that distilling through an intermediate-size "teacher assistant" recovers some of the loss. For LLMs the effect shows up as students that imitate the surface style of a far larger teacher (long, confident reasoning traces) without the competence underneath.

Compute and memory cost#

Sequence-level distillation costs one round of teacher generation plus an SFT run. Generation dominates when responses are long: a 70B teacher writing 800k responses of 4,000 tokens is 3.2 billion generated tokens. The student run is ordinary SFT at roughly $6P_{S}D$ FLOPs for $D$ training tokens.

Token-level distillation adds a teacher forward pass on every training token, about $2P_{T}D$ FLOPs. With a 70B teacher and an 8B student, $2\times 70=140$ against $6\times 8=48$: the teacher's forward pass costs about three times the student's training. Online, the teacher must also be resident in memory alongside the student's training state. Offline, the cost moves to storage: top-$k$ ids and log-probs at roughly 6 bytes per entry (4-byte id, 2-byte value) take

Storing full bf16 distributions over a 128k vocabulary (256 KB per token) would be about 850 times larger, which is why nobody does it. Compared with other post-training stages, sequence-level distillation is SFT-priced plus a one-off generation cost; token-level online distillation is two to four times SFT per token; both are far cheaper than RL, which is why DeepSeek-R1's authors concluded that for small models distillation is the more economical route to reasoning ability.

Failure modes and how to detect them#

Exposure bias

Both forms train on teacher-generated or dataset prefixes, never on the student's own. At inference the student makes mistakes the teacher would not, lands in unfamiliar states, and compounds errors over long generations. Symptom: per-token metrics look good, but quality degrades with response length, and long reasoning traces loop or derail. The remedy is on-policy distillation: sample from the student and let the teacher grade each token.

Tokenizer mismatch

Token-level KD assumes teacher and student share a vocabulary. If they do not, logits cannot be compared position by position. Use the same family for both, use sequence-level distillation, or use one of the cross-tokenizer alignment methods, and validate with a sanity check that the KL on teacher text is near zero when the student is the teacher.

Copying style without substance

The student adopts the teacher's verbosity, hedging and formatting, and its benchmark scores rise less than its output length. Track correctness and length separately. For reasoning, measure accuracy at matched token budgets.

Inheriting teacher errors

Distillation transfers the teacher's mistakes, biases and hallucinations as faithfully as its knowledge. Filter teacher samples with verifiers where possible, and evaluate the teacher on your target distribution before distilling from it.

Tail noise at high temperature

With large vocabularies, high $\tau$ gives weight to thousands of near-zero teacher probabilities whose relative values are noise. Symptom: the student's outputs become slightly more random (higher entropy, occasional odd tokens). Lower $\tau$ or truncate to top-$k$.

When to use it, and what to use instead#

Distil when a stronger model exists for your task and you need a smaller or cheaper one: to serve at lower cost or latency, to deploy on device, or to bootstrap reasoning in a small model that RL alone would struggle to teach. It is also the right first stage before RL for small models. Do not distil when no teacher is meaningfully better than your student on the target distribution, when the teacher's terms of use forbid training on its outputs, or when the capability you need is checkable and the teacher is not much better than the student (RL or RFT against the verifier will do more).

Minimal implementation in PyTorch#

The function below computes the token-level loss for a batch, given student logits and either full teacher logits (online) or stored top-$k$ teacher log-probs (offline). In POFE, distillation is a recipe that takes a teacher checkpoint and switches between these modes; the loss is the part worth understanding.

import torch, torch.nn.functional as F

def kd_loss_full(s_logits, t_logits, labels, tau=1.0, alpha=1.0):
    """s_logits, t_logits: [B, L, V]; labels: [B, L] with -100 on unsupervised positions."""
    mask = (labels != -100).float()
    s_logp = F.log_softmax(s_logits / tau, dim=-1)
    t_logp = F.log_softmax(t_logits / tau, dim=-1)
    kl = (t_logp.exp() * (t_logp - s_logp)).sum(-1)               # forward KL per position
    kd = tau**2 * (kl * mask).sum() / mask.sum()
    ce = F.cross_entropy(s_logits.flatten(0, 1), labels.flatten(), ignore_index=-100)
    return alpha * kd + (1 - alpha) * ce

def kd_loss_topk(s_logits, t_ids, t_logp, labels, tau=1.0):
    """Offline: t_ids, t_logp [B, L, k] are the teacher's top-k token ids and log-probs (at tau)."""
    mask = (labels != -100).float()
    t_p = F.softmax(t_logp, dim=-1)                                # renormalise over the k kept tokens
    s_logp = F.log_softmax(s_logits / tau, dim=-1).gather(-1, t_ids)
    kl = (t_p * (t_p.clamp_min(1e-12).log() - s_logp)).sum(-1)    # KL on the support of the top-k
    return tau**2 * (kl * mask).sum() / mask.sum()

# training step (online teacher)
with torch.no_grad():
    t_logits = teacher(input_ids=batch["ids"]).logits[:, :-1]
s_logits = student(input_ids=batch["ids"]).logits[:, :-1]
loss = kd_loss_full(s_logits.float(), t_logits.float(), batch["labels"][:, 1:], tau=1.0, alpha=0.9)
loss.backward(); opt.step(); opt.zero_grad()

The top-$k$ version computes the student's log-probabilities over the full vocabulary and then reads off the $k$ teacher tokens, so the student is still pushed to put its mass on those tokens rather than renormalising its own distribution. Materialising $[B,L,V]$ logits in fp32 is the main memory cost; chunk over the sequence dimension for long contexts.

Further reading#

Hinton, Vinyals and Dean (2015) for the original formulation and the temperature argument. Kim and Rush (2016) for word-level versus sequence-level distillation in autoregressive models. Sanh et al. (2019) for distillation during pretraining, and the Gemma 2 report (2024) for its use at scale in a modern model family. DeepSeek-AI (2025) for sequence-level reasoning distillation and the comparison against RL on the same base. Gu et al. (2023) for the reverse-KL argument, Ko et al. (2024) for skew-KL, and Mirzadeh et al. (2019) for the capacity-gap problem. The on-policy distillation article covers GKD.

Frequently asked questions#

References#

1. Geoffrey Hinton, Oriol Vinyals, Jeff Dean (2015). Distilling the Knowledge in a Neural Network. arxiv.org/abs/1503.02531
2. Yoon Kim, Alexander M. Rush (2016). Sequence-Level Knowledge Distillation. arxiv.org/abs/1606.07947
3. Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arxiv.org/abs/1910.01108
4. Seyed-Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, Hassan Ghasemzadeh (2019). Improved Knowledge Distillation via Teacher Assistant. arxiv.org/abs/1902.03393
5. Yuxian Gu, Li Dong, Furu Wei, Minlie Huang (2023). MiniLLM: On-Policy Distillation of Large Language Models. arxiv.org/abs/2306.08543
6. Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, et al. (2023). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. arxiv.org/abs/2306.13649
7. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, Se-Young Yun (2024). DistiLLM: Towards Streamlined Distillation for Large Language Models. arxiv.org/abs/2402.03898
8. Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, et al. (2024). Gemma 2: Improving Open Language Models at a Practical Size. arxiv.org/abs/2408.00118
9. DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, et al. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arxiv.org/abs/2501.12948

Related reading