Knowledge Distillation in LLMs: transferring intelligence, not just weights

How to train a small model to reason like a large one — theory, mathematics, architectures, and real-world examples.

The problem

Large language models like GPT-4 or LLaMA 70B are extraordinarily capable. But they are also extraordinarily expensive: they demand hundreds of A100 GPUs for inference, introduce significant latency, and carry price tags that are simply out of reach for most organizations.

The question researchers keep asking is: can we capture 80% of the capability with 10% of the resources? Quite often, the answer is yes — thanks to knowledge distillation.

The heart of the method: soft labels

The foundational concept is the distinction between hard labels and soft labels. In standard training, labels are binary: the correct answer is 1, everything else is 0. The teacher, by contrast, produces full probability distributions over the entire vocabulary.

Given the sentence “the cat sat on the ___”, the teacher doesn’t simply say “mat.” It says something far richer:

The key parameter: temperature

To further amplify the informational signal in the distributions, a temperature parameter T is introduced into the softmax function:

Practical note: During training, a high T (typically 2–10) is applied to both teacher and student. At inference time, you revert to T = 1. The distillation loss must be scaled by a factor of T² to compensate for the reduction in gradient magnitude introduced by the temperature.

The loss function

The student’s total loss combines two signals with tunable weights:

Distillation variants

There is no single approach. The main variants differ in what is transferred and when the transfer occurs:

The distillation pipeline in practice

• Dataset preparation

Collect a corpus of prompts representative of the target task. The more diverse the dataset, the better the student generalises. For Sequence-KD, the teacher directly generates the completions used for training.

• Logit extraction from the teacher

For each token position, the full logit vector is saved — or the top-k logits for memory efficiency. This is the main computational bottleneck: a 70B-parameter teacher demands substantial time and GPU memory. Logits are then normalised with temperature T.

• Student training with combined loss

The student performs a forward pass, computes the combined loss (KL + CE), and updates its weights via standard backpropagation. The T² factor rescales the gradients to compensate for the temperature.

• Hyperparameter tuning

T, α, β, and the learning rate are the most critical knobs. Typical starting points: T ∈ [2, 10], α ∈ [0.7, 0.95]. A large capacity gap between teacher and student generally calls for higher T to soften the target distributions.

• Benchmark evaluation

Measure performance on MMLU (reasoning), HumanEval (coding), HellaSwag (commonsense), and domain-specific tasks. Training loss is not a reliable proxy for real-world performance — always evaluate end-to-end.

The technical challenges nobody talks about

1. Vocabulary mismatch

When teacher and student use different tokenisers, logits are not directly comparable position by position. Solutions include projecting across logit spaces (as in MiniLLM) or restricting supervision to the top-k most probable tokens shared by both vocabularies.

2. Capacity gap

If the teacher is significantly more capable than the student, the student cannot faithfully mimic the teacher’s complex distributions — and vanilla distillation can actually underperform plain fine-tuning. The fix is to use intermediate teacher assistants, or to raise T to simplify the target signal enough for the student to follow.

3. Logit access

For closed-source models like GPT-4, the API exposes only text — not raw logits. This forces a fallback to Sequence-KD, which sacrifices the rich soft-label signal. Some providers offer partial log-probabilities (top-5 or top-20 tokens) as a workable middle ground.

4. Distributional shift

Teacher and student are both evaluated on the same training distribution, but the student generates different text at inference time. MiniLLM addresses this with on-policy KD: the student generates text, the teacher scores it, creating a feedback loop that is more stable and better aligned with actual deployment conditions.

Notable real-world examples

Distillation vs. other compression techniques

Distillation is the only technique that transfers semantic knowledge from one model to another. Every other approach reduces the computational cost of an existing model without extracting and moving what it knows. In production, the most powerful combinations are KD + quantisation and KD + LoRA.

Have you ever worked with distillation or model compression techniques?

I’d love to hear where you’re starting from — and what challenges you’ve run into along the way. Drop a comment: talking through this with people who actually apply it is always more useful than the literature alone.

If this post was useful, consider sharing it with anyone getting into the LLM space!