Knowledge Distillation: Learning from Teachers

Let the teacher and student produce logits $z_T(x)$ and $z_S(x)$ over a discrete output space.

With temperature $\tau>0$, define softened probabilities:

The soft-target term is usually written as the KL divergence between the teacher's and the student's softened distributions. With the teacher fixed, it differs from the soft-target cross-entropy only by the teacher's entropy, a constant, so both give the same gradients:

where:

Here $\mathcal L_{\text{hard}}$ is the ordinary cross-entropy with the true label, computed at $\tau=1$; for a language model it is the next-token loss on the training text. The soft-target gradients shrink as $1/\tau^2$, so the factor $\tau^2$ keeps the balance between the two terms roughly unchanged when $\tau$ changes. Hinton et al. found that a considerably smaller weight on the label term, which here means $\alpha$ close to 1, generally worked best.

Raising $\tau$ softens the teacher distribution and gives more weight to the classes it did not pick. The signal does not vanish at large $\tau$: both distributions approach uniform, but with the $\tau^2$ factor and logits centred to mean zero, the gradient becomes proportional to that of matching the logits directly, $\tfrac12\sum_i (z_{S,i}-z_{T,i})^2$ (Hinton et al., 2015, §2.1). Lower temperatures pay less attention to very negative logits, which may be noise; Hinton et al. found intermediate temperatures worked best when the student was much smaller than the teacher. In practice $\tau$ is tuned.