18. Distillation Loss for a Recurrent Student Model

Problem

Training an SSM/recurrent model from scratch to match a large Transformer's quality is expensive. A cheaper path — the one Cartesia's own Llamba work takes — is distillation: keep a pretrained Transformer as a frozen teacher, and train a much smaller recurrent student to match it, combining a "soft" loss against the teacher's output distribution with the usual "hard" loss against the ground-truth labels. Getting the two loss terms and the temperature scaling right is what makes distillation training actually transfer the teacher's knowledge rather than just retraining on labels with extra steps in between.

The classic (Hinton et al., 2015) combined objective: soft_loss = T^2 * KL(softmax(teacher/T) || softmax(student/T)), hard_loss = CrossEntropy(student, targets), loss = alpha * soft_loss + (1 - alpha) * hard_loss. The T^2 factor rescales the soft-target gradient back to its T=1 magnitude; the temperature softens the teacher's distribution so the student also learns from its relative confidence across incorrect tokens, not just its argmax prediction.

Source: src/18_recurrent_distillation_loss.py

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

>>> student = torch.randn(2, 5, 100)
>>> teacher = torch.randn(2, 5, 100)
>>> targets = torch.randint(0, 100, (2, 5))
>>> recurrent_distillation_loss(student, teacher, targets).shape
torch.Size([])

Step-by-Step Approach

1. Flatten the batch and sequence dimensions together: (batch, seq_len, vocab) → (batch*seq_len, vocab), since both the KL and cross-entropy losses are defined per-token, independent of sequence position.
2. Compute student_log_probs = F.log_softmax(student / temperature, dim=-1) and teacher_probs = F.softmax(teacher / temperature, dim=-1) — F.kl_div expects log-probabilities for its first argument and plain probabilities for its second.
3. Call F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") — "batchmean" sums over the vocabulary and averages only over the batch*seq_len rows, unlike plain "mean" which would also divide by vocab_size and silently shrink the loss as vocabulary grows.
4. Multiply the KL term by temperature ** 2 to restore the gradient scale that dividing both logits by T otherwise shrinks.
5. Compute the hard loss with F.cross_entropy(student_flat, targets_flat) directly on the un-tempered student logits (temperature only applies to the distillation term).
6. Combine: alpha * soft_loss + (1 - alpha) * hard_loss, and sanity-check the two extremes — alpha=0 should reduce to plain cross-entropy, alpha=1 to pure distillation.

The key insight is that the teacher's role is purely to define a fixed target distribution — only the student's logits carry gradients in a real training loop, and the temperature is applied symmetrically to both sides of the KL term but never to the hard cross-entropy term.

Reference solution

import torch
import torch.nn.functional as F


def recurrent_distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    target_ids: torch.Tensor,
    temperature: float = 2.0,
    alpha: float = 0.5,
) -> torch.Tensor:
    batch, seq_len, vocab_size = student_logits.shape
    student_flat = student_logits.reshape(-1, vocab_size)
    teacher_flat = teacher_logits.reshape(-1, vocab_size)
    targets_flat = target_ids.reshape(-1)

    student_log_probs = F.log_softmax(student_flat / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_flat / temperature, dim=-1)
    # batchmean: sum over vocab, then average over the batch*seq_len rows --
    # NOT averaged over vocab_size too, which is what plain "mean" would do
    soft_loss = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean")
    soft_loss = soft_loss * (temperature ** 2)

    hard_loss = F.cross_entropy(student_flat, targets_flat)

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

Key Functions & Tricks

• F.log_softmax / F.softmax — F.kl_div's asymmetric API requires log-probabilities on the input side and plain probabilities on the target side; mixing these up silently produces a wrong-but-finite loss.
• F.kl_div(..., reduction="batchmean") — the reduction mode that matches the mathematical definition of mean KL divergence per example; plain "mean" additionally divides by vocab_size.
• temperature ** 2 rescaling — compensates for the gradient-shrinking effect of dividing logits by T before the softmax, standard in the original Hinton et al. distillation formulation.
• F.cross_entropy — combines log-softmax and negative-log-likelihood internally, applied at T=1 regardless of the distillation temperature.
• tensor.reshape(-1, vocab_size) — collapses batch and sequence axes together so per-token losses are computed independent of position.

How to Recognize This Pattern