Model Distillation for LLMs

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

Why distill a large language model

A large teacher LLM (say, a frontier-scale model) is expensive to serve at high volume — both in latency and compute cost. Knowledge distillation trains a smaller student model to reproduce the teacher's behavior, aiming to capture as much of its capability as possible in a fraction of the parameters. For LLMs specifically, distillation has become central to the economics of deployment: many widely-used 'small' models in production are distilled, in part, from larger siblings rather than trained from scratch on raw data alone.

Soft-label distillation: matching distributions, not just answers

The classical distillation objective (Hinton et al., 2015) doesn't just train the student to predict the correct next token — it trains the student's full output probability distribution to match the teacher's full output distribution, using a temperature-scaled softmax. The loss is typically KL divergence between the teacher's and student's token distributions:

$$\mathcal{L} = \text{KL}(p_{\text{teacher}}^{(T)} \,\|\, p_{\text{student}}^{(T)})$$

with temperature $T > 1$ softening both distributions. Why does this help over just training on hard labels (the single correct token)? Because the teacher's probabilities over *incorrect* tokens still carry information — if the teacher assigns much higher probability to a near-synonym than to a completely unrelated word, that relative ranking is a useful training signal the hard label alone throws away.

distillation_loss.py

python

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, true_labels, temperature=2.0, alpha=0.5):
    # Soft loss: match teacher's softened output distribution
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)

    # Hard loss: standard supervised cross-entropy on the true next token
    hard_loss = F.cross_entropy(student_logits, true_labels)

    return alpha * soft_loss + (1 - alpha) * hard_loss

For LLMs, data distillation often matters more than loss-function distillation

In modern LLM practice, a huge amount of what's loosely called 'distillation' is actually simpler than the KL-divergence formulation above: generate a large synthetic dataset of (prompt, response) pairs using the teacher model, then fine-tune the student on those teacher-generated responses with ordinary supervised cross-entropy. This sidesteps needing access to the teacher's internal logits at all (useful when the teacher is only available through an API) and has proven extremely effective in practice — much of the capability gap closing seen in smaller open models has come from training on large volumes of frontier-model-generated data rather than from sophisticated logit-matching losses.

Logit/distribution distillation vs response (data) distillation

Property	Logit/soft-label distillation	Response (synthetic data) distillation
Requires teacher internals	Yes — needs output logits/probabilities	No — just needs teacher's generated text
Works with API-only teacher	No	Yes
Signal richness per example	Rich — full distribution over vocabulary	Sparse — a single generated response
Common use today	Research settings, open-weight teacher/student pairs	Most practical commercial small-model training

What gets lost in distillation

Distilled students reliably approach the teacher on tasks well-represented in the distillation data, but typically show a larger capability gap on tasks requiring deep multi-step reasoning, rare/long-tail knowledge, or robustness to adversarial/out-of-distribution prompts — precisely the capabilities that are hardest to transfer through imitation because they depend on capacity (parameter count, depth) the smaller student simply doesn't have, rather than on information that imitation alone can recover.

Distillation can silently inherit and amplify teacher biases and errors

Because the student is trained to imitate the teacher rather than to independently verify truth, any systematic bias, factual error, or harmful pattern present in the teacher's outputs gets directly transferred into the training data of the student — and training on a narrower, teacher-generated distribution can sometimes *sharpen* these patterns rather than merely preserving them, since the student never sees the broader, more varied distribution the teacher itself was originally trained on.

What's Next?

Distillation is one lever among several for getting a smaller, cheaper model to behave well — it's frequently combined with quantization and parameter-efficient fine-tuning techniques, each attacking a different axis (capability transfer, numerical precision, trainable parameter count) of the same underlying cost-versus-capability tradeoff.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →