44. Knowledge Distillation Loss (Soft + Hard Target)

Problem

Distillation trains a smaller/cheaper "student" model to match a larger "teacher" model's behavior, so the student can serve at a fraction of the teacher's cost/latency. The classic recipe blends two signals: a "soft" loss that pulls the student's output distribution toward the teacher's full softened probability distribution (encoding which wrong classes the teacher considers "less wrong"), and a "hard" loss against the ground-truth labels the model must ultimately get right.

Implement the combined loss: a temperature-softened KL divergence term plus a standard cross-entropy term, weighted by alpha.

Source: src/44_knowledge_distillation_loss.py

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

Examples:
>>> student = torch.randn(4, 5, requires_grad=True)
>>> teacher = torch.randn(4, 5)
>>> labels = torch.randint(0, 5, (4,))
>>> loss = distillation_loss(student, teacher, labels)
>>> loss.shape
torch.Size([])

Step-by-Step Approach

1. Soften both distributions by dividing logits by the temperature T before the softmax — this widens the probability mass on non-argmax classes, exposing the teacher's relative confidence across wrong answers ("dark knowledge").
2. Compute the teacher's softened probabilities with F.softmax(teacher_logits / T, dim=-1) and the student's softened log-probabilities with F.log_softmax(student_logits / T, dim=-1).
3. Feed both into F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") — note the argument order: PyTorch's kl_div takes (input=log-probs, target=probs).
4. Rescale the KL term by T²: softened gradients shrink by roughly 1/T², so this keeps the soft loss's gradient magnitude comparable to the hard loss's regardless of what temperature was chosen.
5. Compute the hard loss as plain F.cross_entropy(student_logits, labels) at the original (unsoftened) logits.
6. Combine as alpha * soft_loss + (1 - alpha) * hard_loss and return the scalar.

The key insight is that temperature scaling and the T² gradient correction aren't optional details — without them the soft-target signal either vanishes (too little useful gradient at low T) or is scaled inconsistently against the hard-label term as you tune T.

Reference solution

def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # soft loss: KL(teacher || student) at temperature T
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")
    # T^2 rescale: softened gradients shrink by ~1/T^2, this corrects for it
    soft_loss = soft_loss * (temperature ** 2)

    # hard loss: standard supervised cross-entropy at original logits
    hard_loss = F.cross_entropy(student_logits, labels)

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

Key Functions & Tricks

• F.log_softmax / F.softmax — computed at temperature-scaled logits, not the raw logits, to soften the distributions.
• F.kl_div(..., reduction="batchmean") — PyTorch's KL divergence expects log-probabilities as input and probabilities as target; batchmean divides by batch size only, matching the textbook KD formula (not mean, which would also divide by the class count).
• F.cross_entropy — combines log-softmax and negative log-likelihood internally; takes raw logits and integer class labels directly.
• temperature ** 2 gradient correction — the detail most likely to be missed if implementing this from memory.

How to Recognize This Pattern