How LLMs are distilled, step by step
A visual guide to understanding how LLMs are trained using model distillation.
Model distillation, also called Knowledge distillation (KD) or simply Distillation, is a popular technique that helps build smaller, faster, and more efficient models from larger and more capable models that require massive compute and memory resources for deployment.
The process of Knowledge distillation (KD) was described in a 2015 paper by Hinton et al. It involves training a weaker/smaller model called the “student model” to imitate the outputs of a stronger/ larger model called the “teacher model”.
The student learns from teacher-generated output trajectories, rather than from trajectories generated by itself. This is why conventional KD is an Off-policy post-training method (as opposed to On-policy distillation).
The process of KD goes like this:
Take a larger/stronger teacher model and a smaller/weaker student.
Create a dataset of prompts relevant to the training task.
Pass a prompt to the teacher model and collect the teacher's next-token probability distribution over the vocabulary. These probabilities are called “soft targets” (as opposed to “hard targets” or training labels).
Pass the prompt and the teacher’s response to the student model and collect the student’s next-token probability distribution conditioned on the same prefix as the teacher.
Use a loss function to measure how much the teacher’s and student’s probability distributions differ. A commonly used loss is the KL Divergence.
Optionally, a hyperparameter called Temperature is applied to the logits of both the teacher and the student models when calculating the distillation loss. This helps flatten the teacher’s probability distribution of predictions, which better reveals the relationships between different tokens.In many cases, alongside KL Divergence, a Cross-entropy loss between the student’s final prediction and the ground-truth label, or “hard target”, is also calculated and used to formulate the combined distillation loss. This combined loss helps the student model to learn from both the teacher model and from ground-truth data.
The teacher model is kept frozen, but the student model’s parameters are updated using backpropagation.
This process is repeated over all prompts in the dataset as the loss decreases and the student model learns to approximate the teacher model's outputs.
Distillation with real numbers
The process of knowledge distillation is shown in detail in the illustration below.
