Distillation

Distillation trains student model using a frozen teacher’s output distributions. For an input , the teacher produces a probability distribution over classes, . The student is trained to match this distribution, often using a loss function like KL divergence.

Teacher outputs is called soft labels. Labels other than the correct class are called dark knowledge. Why do we not teach the student to only match the correct class, but rather the entire distribution? Because dark knowledge contains more information than just the correct class e.g., what the input is similar to, which classes are more likely than others, etc.

Examples:

• Google Gemini distills into Gemma series
• Llama Behemoth distills into Scout and Maverick

Pros: The teacher needs to be big during training because it needs a big parameter space to wiggle around. At inference time, this costs latency so we distill into a smaller student model so inference is fast and cheap.

Cons: Soft labels take up a lot space if vocabulary is large so either need to compute on the fly or just store the top-k classes. Also, distillation is a lossy compression so student will never be as good as teacher.

Reference: Distilling the Knowledge in a Neural Network