Knowledge Distillation

August 2026 – Vladislav Kruglikov

Knowledge distillation is a teacher to student paradigm where a smaller student model learns not only from ground truth labels but from the full output distribution of a larger teacher effectively performing a lossy compression of the teacher’s behavior. This allows model compression and faster inference, often improves the quality of small models beyond standard training, and serves as a form of transfer learning by transferring uncertainty and structure learned by the teacher

Loss formulation

Distillation uses:

Ldistill=α⋅CE+(1−α)⋅KL(Pteacher ∥ Qstudent)\mathcal{L}_{\text{distill}} = \alpha \cdot \text{CE} + (1 - \alpha) \cdot \mathrm{KL}(P_{\text{teacher}} \,\|\, Q_{\text{student}})

The weighting coefficient α\alpha controls the trade off between factual correctness matching ground truth and imitation of the behavior of the teacher.

Hard labels

Hard labels are the standard ground truth answers the teacher model was originally trained on. They provide only the single correct target and do not encode uncertainty or relationships between alternative outputs.

Soft labels

Soft labels are the full probability distributions produced by the teacher model. They are typically obtained using temperature scaling which smooths the distribution and makes it less peaky exposing additional information. soft labels are more informative. Unlike hard labels soft labels convey:

For example if teacher outputs 70% for cat and 25% for dog and 5% for bird this soft distribution represents distilled knowledge. It tells the student that the input is somewhat similar to both a cat and a dog but clearly not a bird. Information that is lost when training only on the hard label cat.

Temperature scaling

Teacher distribution is too peaky. Collapses to hard labels because low probability classes have almost no weights into gradient. Use temperature greater then 1 usually around 2 to 4 to smooth distribution. If model was overfitted to something we can use this to reveal other hidden knowledge model has to learn from.

Implementations

There are multiple implementations:

Granulariry

Vocabulary compatibility

Token-level logit matching is straightforward only when the teacher and student output spaces align. If they use different tokenizers or vocabularies, their logits refer to different tokens and cannot be compared directly. Distillation then requires token mapping, sequence-level supervision, or another method for aligning the two output spaces.

Logit truncation

Teacher outputs are usually stored as logits and not probabilities because logits more numeraiclly stable and temperature scaling is applied to logits thus softmax later can be recomputed with different temperatures.

Storing full logits over the entire vocabulary is prohibitively expensive. Most probability mass is covered by top probable tokens (that is also called popularity bias). Tail tokens contribute little to gradients. KL over top tokens is a good approximation of full KL.

Quality based selection

Filtering step applied to teacher generated data. It decides which teacher outputs are kept for training and which are discarded. Distillation blindly imitates the teacher. That means hallucinations, incorrect reasoning, low quality or inconsistent outputs will be copied by the student unless filtered out. Quality based selection exists to remove junk before training.

  1. Teacher generates one or more candidate completions for a prompt.
  2. Each candidate is scored or evaluated.
  3. Low quality candidates are rejected.
  4. Only accepted samples are used for distillation.

The student never sees rejected outputs.

Common signals:

On policy distillation of language models

Suppose the prompt contains tokens x0:tx_{0:t} and the completion contains tokens from positions t+1t+1 through t+mt+m. If the teacher generates the completion

yt+1T,…,yt+mTy^T_{t+1}, \ldots, y^T_{t+m}

then, when the student computes its distribution at position t+it+i, it is conditioned on the prompt and the teacher's preceding completion tokens:

x0:t,  yt+1:t+i−1Tx_{0:t},\; y^T_{t+1:t+i-1}

The student is therefore trained to answer the question «What would you predict after the teacher's previous tokens?» At inference, however, the student is conditioned on its own preceding tokens:

x0:t,  yt+1:t+i−1Sx_{0:t},\; y^S_{t+1:t+i-1}

There is no mismatch at the first completion position because both models are conditioned only on the prompt. The mismatch begins after the student chooses a different token from the teacher. From that point onward, ordinary distillation trains the student on prefixes it may not encounter when generating by itself.

On-policy distillation instead generates a student completion ySy^S first and compares the teacher and student distributions under the same student prefix:

pT(⋅∣x,y<iS)andpS(⋅∣x,y<iS)p_T(\cdot \mid x, y^S_{<i}) \quad\text{and}\quad p_S(\cdot \mid x, y^S_{<i})

The student is thus trained under its own preceding tokens, matching the situation it faces during inference.

In paper authors suggest:

Industry

Gemma 3 Technical Report:

References