42. Knowledge Distillation Loss With Temperature Scaling

Problem

Distillation trains a (usually smaller or architecturally different) student model to match a larger teacher model's output distribution, not just the ground-truth hard labels. The classic formulation blends two terms: an ordinary cross-entropy loss against the hard labels, and a KL-divergence loss between the student's and teacher's temperature-softened output distributions. Dividing both logits by a temperature T > 1 before the softmax flattens the distribution, exposing more of the teacher's relative confidence across wrong classes ("dark knowledge") instead of just its single top prediction. The soft term is rescaled by T^2 so its gradient magnitude stays comparable to the hard term's regardless of which temperature is chosen.

Implement the blended loss. This is directly the training technique behind Cartesia's own public Llamba work, distilling a Transformer teacher into a Mamba-architecture student.

Source: src/42_distillation_loss_temperature.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:
>>> torch.manual_seed(0)
>>> student = torch.randn(4, 5)
>>> 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. Compute the hard-label term first, since it's the simple part: F.cross_entropy(student_logits, labels), completely independent of the teacher.
2. Soften both distributions by dividing logits by the temperature before the softmax: teacher_logits / T and student_logits / T. Higher T flattens the distribution more.
3. Compute the teacher's softened distribution with a plain softmax (it's a fixed target, no gradient needs to flow through it in a real training loop).
4. Compute the student's softened distribution with log_softmax, not log(softmax(...)) — this is the numerical-stability trick: taking the log of an already-computed softmax can underflow to -inf for near-zero probabilities, while log_softmax computes log-probabilities directly and stays stable.
5. Combine them with F.kl_div(student_log_probs, teacher_probs, reduction="batchmean"), matching PyTorch's expected argument order: input is log-probabilities, target is probabilities. "batchmean" divides by batch size only, not batch times classes.
6. Rescale the soft loss by T^2 — softening flattens gradients too (they scale roughly as 1/T), so without this the soft term's influence on training would silently shrink as T grows, changing what alpha effectively means.
7. Blend the two terms: alpha * hard_loss + (1 - alpha) * soft_loss. Sanity check both extremes: alpha=1 should reduce to plain supervised CE (teacher ignored entirely), and alpha=0, T=1 should reduce to an untempered KL divergence between the two models' raw distributions.

The key insight is that two easy-to-miss details, log_softmax instead of log(softmax()) for stability, and the T^2 rescale for gradient-magnitude consistency, are exactly the parts that separate a "looks right" implementation from one that actually reproduces the textbook Hinton distillation loss.

Reference solution

def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Hard term: ordinary supervised cross-entropy against ground-truth labels.
    hard_loss = F.cross_entropy(student_logits, labels)

    # Soft term: KL-divergence between temperature-softened teacher and
    # student distributions. Dividing by T > 1 flattens both softmaxes,
    # exposing the teacher's relative confidence across *wrong* classes.
    #
    # NUMERICAL-STABILITY TRICK: use log_softmax for the student side (not
    # log(softmax(...))) -- computing softmax then taking its log can
    # underflow to log(0) = -inf; log_softmax stays numerically stable.
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student_log = F.log_softmax(student_logits / temperature, dim=-1)

    # F.kl_div expects (input=log-probs, target=probs); "batchmean" divides
    # the summed-over-classes KL by the batch size only.
    soft_loss = F.kl_div(soft_student_log, soft_teacher, reduction="batchmean")

    # T^2 RESCALE: keeps the soft term's gradient magnitude comparable to
    # the hard term's across different temperature choices.
    soft_loss = soft_loss * (temperature ** 2)

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

Key Functions & Tricks

• F.log_softmax vs. torch.log(F.softmax(...)) — the numerical-stability fix; the fused version avoids the intermediate underflow.
• F.kl_div(input, target, reduction="batchmean") — note the non-obvious argument order (log-probs first, probs second) and that "batchmean" differs from the (deprecated) default "mean" reduction.
• temperature ** 2 rescale — keeps alpha's meaning stable across different temperature choices.
• logits / temperature before softmax — the softening step; note it's applied to both teacher and student consistently.
• Independent test oracle via a manual elementwise KL sum ((p * (log(p) - log_q)).sum(-1).mean()) — validates the loss's math without depending on the same F.kl_div call path the solution uses.

How to Recognize This Pattern