Knowledge-Distillation Loss

Background

Knowledge Distillation (Hinton et al.) trains a small student to match a large teacher's softened logits, plus a small CE term on the hard ground-truth labels. Temperature $T$ softens both softmaxes so the student picks up dark knowledge about non-target classes.

Problem statement

Implement kd_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):

$$\mathcal L \;=\; \alpha\, T^2\, \mathrm{KL}\!\bigl(\sigma_T(\text{teacher}) \,\|\, \sigma_T(\text{student})\bigr) \;+\; (1-\alpha)\,\mathrm{CE}\!\bigl(\sigma_1(\text{student}),\ y\bigr)$$

where $\sigma_T(x) = \mathrm{softmax}(x/T)$ and the $T^2$ factor restores gradient scale.

Input

• student_logits — np.ndarray shape $(N, K)$.
• teacher_logits — np.ndarray shape $(N, K)$.
• labels — sequence of int, length $N$, in $[0, K)$.
• T — float > 0, temperature (default 4.0).
• alpha — float in $[0, 1]$, KD weight (default 0.9).

Output

Returns a Python float >= 0.

Examples

Example 1 — identical logits with alpha=1 → KL=0 → loss=0

Input:  student==teacher, alpha=1.0
Output: ≈ 0

Example 2 — alpha=0 reduces to plain cross-entropy on hard labels

Input:  perfect student logits (correct class very high), alpha=0
Output: near 0

Example 3 — returns plain Python float

Input:  any valid logits, labels
Output: isinstance(result, float)

Constraints

• student_logits and teacher_logits must agree in shape.
• labels length matches batch dimension.
• Use a numerically stable softmax (subtract max before exp).

Notes

• Why $T^2$. Soft targets shrink gradient magnitude by $1/T^2$; the $T^2$ multiplier brings it back so the KD term doesn't get drowned out by the CE term at high $T$.
• vs DistilBERT-style. Adds a hidden-state MSE term that pushes the student's intermediate representations to match the teacher — out of scope here.