9. Teacher-Student Distillation Loss for Sequence Models

Problem

Distilling a large Transformer teacher into a smaller, faster recurrent (SSM/Mamba-style) student is a real technique for cutting inference cost while keeping most of the teacher's quality — the student replaces attention with a linear recurrence, trading some capacity for the O(length) (vs. O(length²)) inference cost and constant per-token memory that a real-time system benefits from. Training such a student needs more signal than the final task loss alone: a standard recipe combines (a) a soft-label matching term — KL divergence between the student's and teacher's output distributions, temperature-scaled so neither distribution is too peaked to give useful gradient — with (b) an intermediate hidden-state matching term, so the student's internal representations track the teacher's, not just its final outputs.

Implement this two-term loss: KL(teacher || student) on temperature-scaled softmax output distributions (scaled by temperature², the standard correction so the gradient magnitude doesn't shrink as temperature grows), plus MSE between student and teacher intermediate hidden states, combined with a weight alpha.

Source: src/9_sequence_model_distillation_loss.py

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

Examples:
>>> logits = torch.randn(2, 3, 5)
>>> hidden = torch.randn(2, 3, 6)
>>> distillation_loss(logits, logits.clone(), hidden, hidden.clone())
tensor(-0.0000)  # ~0: identical student/teacher distributions and hidden states

Step-by-Step Approach

1. Temperature-scale both logit tensors by dividing by temperature before any softmax — a higher temperature softens the distribution, revealing more of the teacher's relative confidence across non-top classes ("dark knowledge") instead of just its argmax.
2. Compute student_log_probs = F.log_softmax(student_logits / T, dim=-1) and teacher_probs = F.softmax(teacher_logits / T, dim=-1) — F.kl_div expects log-probabilities as its first argument and probabilities as its target.
3. Call F.kl_div(student_log_probs, teacher_probs, reduction="batchmean"), then multiply by temperature**2 — the standard Hinton et al. correction: since the softmax gradient scales roughly like 1/T, without this factor a higher temperature would silently shrink the KL term's contribution to the total loss.
4. Compute the hidden-state matching term as plain F.mse_loss(student_hidden, teacher_hidden) — no temperature scaling here, since it's not operating on a softmax output.
5. Combine with the weighting: alpha * kl_term + (1 - alpha) * hidden_term, returning a single scalar.
6. Sanity-check with the degenerate case: if student_logits == teacher_logits and student_hidden == teacher_hidden exactly, both terms should be (numerically) zero, since KL divergence between identical distributions is zero and MSE between identical tensors is zero — a strong, cheap invariant to test before trusting the general case.

The key insight is that F.kl_div's argument convention is easy to get backwards (it wants log(Q) as input and P as target to compute KL(P || Q), not the reverse), and getting it backwards doesn't crash — it silently computes a different, still-plausible-looking loss, which is exactly the kind of subtle bug that's worth explicitly reasoning through rather than pattern-matching from memory.

Reference solution

def distillation_loss(
    student_logits, teacher_logits, student_hidden, teacher_hidden,
    temperature=2.0, alpha=0.5,
):
    T = temperature
    # temperature-scaled soft labels: dividing logits by T before softmax softens the
    # distribution, giving more gradient signal from the teacher's non-top predictions
    student_log_probs = F.log_softmax(student_logits / T, dim=-1)
    teacher_probs = F.softmax(teacher_logits / T, dim=-1)
    # F.kl_div expects log-probs as input, probs as target: computes KL(target || input)
    kl_term = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * (T ** 2)

    # feature/hidden-state matching: pulls the student's internals toward the teacher's,
    # not just its final output distribution
    hidden_term = F.mse_loss(student_hidden, teacher_hidden)

    return alpha * kl_term + (1 - alpha) * hidden_term

Key Functions & Tricks

• F.kl_div(input, target, reduction="batchmean") — input must be log-probabilities, target must be probabilities; batchmean normalizes by batch size (the mathematically correct KL divergence normalization)
• F.log_softmax / F.softmax — paired temperature-scaled distributions for student (log form) and teacher (prob form)
• F.mse_loss(a, b) — standard mean-squared-error, used here for hidden-state/feature matching
• Temperature² gradient correction — the Hinton et al. (2015) scaling factor that keeps KL gradient magnitude comparable across temperature settings
• Multi-term weighted loss combination — alpha trades off output-matching against feature-matching, a common pattern across distillation, multi-task, and auxiliary-loss training setups

How to Recognize This Pattern