Knowledge Distillation › The Distillation Objective

To understand why distillation works, we need to examine what information is encoded in a model's output probabilities and how the student is trained to capture it. The mechanism is elegant, but it takes some care to see why it provides more signal than ordinary supervised training.

A standard language model training objective uses cross-entropy loss against the true one-hot label: the correct next token gets probability 1.0, all others get 0.0. This is called training with hard labels. Hard labels are efficient and mathematically clean, but they discard an enormous amount of information.

Consider what happens when a language model processes the sentence "The cat sat on the ___." The model assigns high probability to "mat" but also substantial probability to "floor," "rug," "sofa," and "chair." These secondary probabilities are not noise. They reflect semantic knowledge: all of these completions are plausible, and the differences between their probabilities reflect subtle contextual cues about which is most natural in this exact phrasing. Hard-label training treats this entire distribution as equivalent to a one-hot vector pointing at "mat." All that relational information evaporates.

The same problem appears at finer scales. If the teacher assigns 60% probability to "cat" and 35% probability to "feline," that relationship between near-synonyms is discarded by hard-label training. The student only sees "cat: 1, everything else: 0." But the teacher's soft assignment reveals that "cat" and "feline" occupy the same semantic neighborhood in the model's representation space. A student that learns to match this distribution will develop a similar geometric structure in its own embeddings.

Distillation trains the student on soft labels: the teacher's full probability distribution over all possible outputs. Let $p^T_i$ be the teacher's probability for token $i$, and $p^S_i$ be the student's predicted probability for token $i$. The distillation loss is the cross-entropy between these two distributions:

$$\mathcal{L}_{\text{distill}} = -\sum_{i=1}^{V} p^T_i \log p^S_i$$

where:

• $V$: the vocabulary size
• $p^T_i$: the teacher's probability assigned to token $i$ (a soft label)
• $p^S_i$: the student's predicted probability for token $i$
• The sum runs over all $V$ tokens in the vocabulary

This is Kullback-Leibler divergence between the teacher and student distributions up to a constant, so minimizing the cross-entropy is equivalent to minimizing the KL divergence:

$$\mathcal{L}_{\text{distill}} = D_{\text{KL}}(p^T \| p^S) + H(p^T)$$

where $H(p^T)$ is the teacher's entropy, which is constant with respect to the student's parameters. Minimizing the distillation loss therefore pushes the student's distribution toward the teacher's, which is exactly what we want.

Why does this help? The teacher has been trained on vastly more computation than the student could afford. Its output distribution encodes the result of fitting billions of parameters to massive datasets. When the student is trained to match that distribution, it gets a rich supervisory signal that goes far beyond binary correct-or-wrong feedback. Every training example now communicates a full probability distribution rather than a single spike, and the student can learn the relative relationships between outputs from that distribution.

In practice, most distillation setups combine the distillation loss with the standard supervised loss against the ground-truth labels. This ensures the student learns from labeled data directly while also absorbing the teacher's knowledge:

$$\mathcal{L} = \alpha \cdot \mathcal{L}_{\text{distill}} + (1 - \alpha) \cdot \mathcal{L}_{\text{CE}}$$

where:

• $\mathcal{L}_{\text{distill}}$: the distillation loss (cross-entropy against teacher soft labels)
• $\mathcal{L}_{\text{CE}}$: the standard cross-entropy loss against ground-truth labels
• $\alpha \in [0, 1]$: a mixing coefficient that controls the balance between the two objectives

When $\alpha = 1$, the student learns purely from the teacher's soft labels. When $\alpha = 0$, it learns only from ground truth. Typical values in practice lie in the range $[0.5, 0.9]$, favoring the distillation signal while retaining some grounding in labeled data. The appropriate value depends on how well the teacher's distributions match the true underlying distribution and how much labeled data is available.

This combination is useful for a subtle reason. The teacher's soft labels encode structural knowledge about relationships between classes, but they are still the teacher's beliefs, not ground truth. The teacher makes errors. If the teacher systematically assigns too much probability to certain tokens due to biases in its training data, the student will inherit those biases. Mixing in the ground-truth loss acts as a corrective force, anchoring the student to the actual labels even while it absorbs the teacher's relational knowledge.