44. Knowledge Distillation Loss (Soft + Hard)

Problem

A common way to compress a large trained model into a smaller, cheaper one is knowledge distillation: keep the large model frozen as a "teacher" and train the small "student" to match both the ground-truth labels (the "hard" loss) and the teacher's full output distribution (the "soft" loss). The soft loss is what actually transfers extra signal beyond the labels -- the teacher's relative confidence across wrong classes encodes information a one-hot label alone throws away.

Implement the classic (Hinton et al., 2015) combined objective, given student and teacher logits over a fixed set of classes and the integer target labels.

Source: src/44_knowledge_distillation_loss.py

def distillation_loss(
    student_logits: torch.Tensor, teacher_logits: torch.Tensor,
    targets: torch.Tensor, temperature: float = 4.0, alpha: float = 0.5,
) -> torch.Tensor: ...

>>> student = torch.randn(4, 10)
>>> teacher = torch.randn(4, 10)
>>> targets = torch.randint(0, 10, (4,))
>>> distillation_loss(student, teacher, targets).shape
torch.Size([])

Step-by-Step Approach

1. Soften both logit sets by dividing by temperature before any softmax -- a higher temperature spreads probability mass more evenly across wrong classes, which is exactly the extra signal distillation is trying to transfer.
2. Compute student_log_probs = log_softmax(student_logits / T) -- F.kl_div requires log-probabilities as its first argument, not raw probabilities.
3. Compute teacher_probs = softmax(teacher_logits / T) as the fixed target distribution for the KL term.
4. Compute the soft loss with F.kl_div(..., reduction="batchmean") -- this reduction sums over classes and averages only over the batch, not over the class dimension too.
5. Multiply the soft loss by T^2 to rescale its gradient back to the magnitude it would have at T=1 (dividing logits by T shrinks the gradient by roughly 1/T^2).
6. Compute the hard loss as ordinary F.cross_entropy(student_logits, targets) against the ground-truth labels, using the un-softened student logits.
7. Combine as alpha * soft_loss + (1 - alpha) * hard_loss and return the scalar.

The key insight is that the T^2 rescaling is not optional bookkeeping -- without it, raising the temperature (which is needed to soften the teacher's distribution) silently shrinks the soft loss's gradient contribution, making alpha not actually mean what it looks like it means.

Reference solution

def distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    targets: torch.Tensor,
    temperature: float = 4.0,
    alpha: float = 0.5,
) -> torch.Tensor:
    # F.kl_div wants log-probs as its first arg and plain probs as its
    # second; teacher stays a fixed target distribution.
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)

    # reduction="batchmean" sums KL over classes then averages over the
    # batch -- plain "mean" would ALSO divide by num_classes, silently
    # shrinking the loss as the class count grows.
    soft_loss = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")
    # dividing logits by T shrinks gradients by ~1/T^2, so T^2 restores
    # the soft-loss gradient to the same scale it would have at T=1.
    soft_loss = soft_loss * (temperature ** 2)

    hard_loss = F.cross_entropy(student_logits, targets)

    return alpha * soft_loss + (1.0 - alpha) * hard_loss

Key Functions & Tricks

• F.log_softmax / F.softmax — kl_div's two arguments are asymmetric: log-probs for the "input" side, plain probs for the "target" side.
• F.kl_div(..., reduction="batchmean") — the one reduction mode that matches the mathematical definition of a per-example KL divergence averaged over a batch; "mean" is a common, silent trap here.
• F.cross_entropy — combines log_softmax and nll_loss internally, so it takes raw logits directly, not probabilities.
• temperature ** 2 — the gradient-rescaling factor tied to dividing logits by T before the softmax.

How to Recognize This Pattern