Knowledge Distillation

Teach a small model what a big one knows

Knowledge distillation trains a small student model to imitate a large teacher model — not just its final answers, but its full probability distribution over classes. A hard label says "this is a 7." The teacher's softmax says "95% a 7, 4% a 1, 1% everything else," and that 4%-toward-1 is information: it tells the student which classes look alike. That extra signal is what Hinton called dark knowledge, and it's why a distilled student often beats an identical-size model trained from scratch on hard labels alone.

How the loss is built

You subclass keras.Model and override train_step. Each step you run the (frozen) teacher and the student on the same input, then combine two losses: a hard loss against the true labels, and a soft loss (KL divergence) between the two probability distributions. The alpha weight balances them.

The temperature knob

temperature controls how soft those distributions get. Dividing the logits by a temperature > 1 before softmax flattens the peak, which surfaces the small inter-class probabilities the student is supposed to learn from — at temperature=1 the teacher's near-one-hot output would tell the student almost nothing extra. It's a hyperparameter; 3-5 is a common starting range.

Code

Distiller: a keras.Model with a custom train_step·python

class Distiller(keras.Model):
    def __init__(self, student, teacher, temperature=3.0, alpha=0.1):
        super().__init__()
        self.student = student
        self.teacher = teacher
        self.temperature = temperature
        self.alpha = alpha

    def train_step(self, data):
        x, y = data
        # Teacher predictions (soft labels)
        teacher_pred = self.teacher(x, training=False)

        # Student predictions
        student_pred = self.student(x, training=True)

        # Distillation loss (soft) + student loss (hard)
        loss = (
            self.alpha * hard_loss(y, student_pred) +
            (1 - self.alpha) * soft_loss(teacher_pred, student_pred)
        )
        ...

Teach a small model what a big one knows

How the loss is built

The temperature knob

Code

External links

Exercise

Progress

Comments 0