Distillation: train on what another model knows

Knowledge distillation trains a student model using signals from a teacher. The student may be smaller, cheaper to run or better suited to a device. Its weights are learned through training; they are not simply a shortened copy of the teacher's weights.

A teacher can provide more information than one correct label. Its probability distribution shows which alternatives it considers plausible. For a language model, that can mean scores over the next-token vocabulary. Other methods use generated responses, intermediate representations or combinations of signals.

Learn from the whole distribution

This PyTorch loss accepts student and teacher logits of shape batch × classes, plus integer target labels. Both models must score the same classes in the same order. The teacher is detached, so this loss updates only the student side.

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels,
                      temperature=2.0, mix=0.5):
    teacher = F.softmax(teacher_logits.detach() / temperature, dim=-1)
    student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(student, teacher, reduction="batchmean")
    soft_loss *= temperature ** 2
    hard_loss = F.cross_entropy(student_logits, labels)
    return mix * soft_loss + (1 - mix) * hard_loss

The temperature softens both distributions. KL divergence compares their probabilities; the temperature ** 2 factor compensates for the gradient scaling introduced by softening. Cross-entropy supplies ordinary feedback from the target labels. mix chooses the balance between those objectives.

For token-level language modeling, valid token positions can be flattened into the batch dimension after excluding padding or other masked targets. Different vocabularies need an additional alignment strategy.

A smaller model still has to learn

Distillation transfers selected behaviour on the training inputs. It does not guarantee that every teacher capability fits into the student. Evaluate the student on the tasks and devices it will actually serve.

Quantization changes numerical precision. Distillation changes how a model is trained. They can be combined.