KD takes advantage of the fact that artificial neural networks are “universal approximators”: given enough training data, and a large enough hidden layer, a neural network can approximate any function to arbitrary precision.⁵

In conventional machine learning, the “knowledge” of a trained model is identified with its learned parameters: the variable weights (and biases), applied to the different mathematical operations occurring across the neural network, that amplify or diminish the influence a certain part of the network’s output has on another part. This view of knowledge makes it hard to see how one model can absorb the knowledge of another model of a different size and structure.

Instead, Hinton et al applied a more abstract, flexible view of knowledge as simply “a learned mapping from input vectors to output vectors.” In other words, KD interprets a model’s knowledge not as the strictly mathematical parameters it learns in training, but as how it generalizes to new data after that training.

Through this alternate understanding of knowledge, knowledge distillation methods aim to train student models to mimic not just the teacher model’s final output for a given input, but also the reasoning steps the teacher model takes to arrive at that final output. Conceptually, this works similarly to instruction tuning through chain-of-thought (CoT) prompts, which improves the quality of LLM responses by teaching them to articulate their “step by step” rationale.

In conventional supervised or self-supervised learning, a loss function produces a vector representing the divergence (or loss) between the model’s outputs and the “correct” outputs (or ground truth) across different inputs. By adjusting model parameters to minimize the slope (or gradient) of this vector through an optimization algorithm like gradient descent, the model’s outputs come closer to those correct outputs. While the model’s reasoning steps are “important” in that they influence its final output, they are not typically measured by a conventional loss function.

Knowledge distillation, conversely, also trains the student model to mimic the teacher model’s reasoning process through the addition of a specialized type of loss function, distillation loss, that uses discrete reasoning steps as soft targets for optimization.

The output of any AI model can be understood as predictions: an autoregressive LLM predicts the next word(s) in a specified sequence; a computer vision model used for image classification predicts the category of a certain image. To arrive at these final predictions, called “hard targets” in this context, deep learning models typically make multiple preliminary predictions and use a softmax function to output the prediction with the highest probability. During training, a cross-entropy loss function is used to maximize the probability assigned to the correct output and minimize the probability assigned to incorrect outputs.

For example, an image classification model predicts the probability of an input image belonging to each known class the model is trained to recognize, then outputs the class with the highest probability value. In the mathematical parlance of machine learning, these individual classwise predictions are called logits. Similarly, an autoregressive LLM predicts multiple possibilities for each next word and (depending on its temperature setting) samples one of those possibilities for its output.

In knowledge distillation, these intermediate predictions—the “soft targets”—generated by the teacher model often provide the principal training data for the student model. The relative probabilities assigned to these preliminary predictions provide valuable insight into how the teacher model tends to generalize. For example, an image classification model is many times more likely to misclassify an image of a fox as “dog” than as “sandwich.” Soft targets thus provide far more information per training case than hard targets alone.

Soft targets also provide more consistency than hard targets: a model’s final prediction might ultimately hinge on a minuscule difference between two logit values, but the logit values themselves have much less variance in the gradient between each training example.

Because of the richness and stability of the information provided by soft targets, the student model can be trained on fewer training examples, using a higher learning rate, than were used to train the original teacher model.

To bring the student network’s generalization tendencies closer to those of the teacher network, knowledge distillation typically uses two loss functions. The first is a standard loss function that operates on “hard loss,” measuring the student model’s final outputs against the ground truth labels (in supervised learning) or against the original data sample (in self-supervised learning). The second is distillation loss, a “soft loss” measuring the student model’s soft targets against those of the teacher.

Because there can be multiple soft targets for each training example, distillation loss measures the difference between the probability distribution of the teacher network’s soft targets and the probability distribution of the student’s. Kullback-Leibler divergence (or “KL divergence”) is commonly used for this purpose.

These different forms of knowledge generally fall into one of three categories: response-based knowledge, feature-based knowledge or relation-based knowledge.

Response-based knowledge, the most common genre of knowledge distillation, focuses on transferring information from the final output layer of the teacher model. In a typical response-based KD method, the student model is trained to output logits that match the teacher model’s predictions.

When the teacher model’s soft targets have low entropy—in other words, when the predictions are extremely “confident,” like if a classification model outputs a logit very close to 1 (representing certainty) for one class and logits approaching 0 for all others—they do not provide as much information. Response-based methods thus often use a high temperature setting for model outputs, which increases the entropy of model predictions. This ensures a more variable probability distribution and thus a greater quantity of information from each training example.

Feature-based knowledge focuses on information that is conveyed in the intermediate layers, or “hidden layers,” of a neural network. This is where neural networks tend to perform feature extraction, the identification of distinct characteristics and patterns of the input data that are relevant to the task at hand.

For example, in the convolutional neural networks used predominantly for computer vision tasks like image segmentation, each successive hidden layer captures progressively richer detail as data is transmitted across the network. In a model used to classify images of animals by species, the earliest hidden layers might simply discern the presence of an animal shape in one part of the photo; the middle hidden layers might discern that the animal is a bird; the final hidden layers, just before the output layer, would discern the nuanced details differentiating one species of bird from another closely related species.

The goal of feature-based knowledge distillation methods is thus to train the student model to learn the same features as the teacher network. Feature-based distillation loss functions are used to measure and then minimize the difference between the two networks’ feature activations.

Whereas both response-based and feature-based knowledge focus on the outputs of specific model layers, relation-based knowledge distillation focuses on the relationships between different layers or between feature maps representing the activations at different layers or locations.

In essence, relation-based knowledge represents perhaps the comprehensive approach to training the student network to emulate the teacher model’s “thought process.” These relationships and correlations can be modeled in various ways, including correlations between feature maps, matrices representing the similarity between different layers, feature embeddings or probabilistic distributions of feature representations.

In offline distillation, the teacher network is already pre-trained and its model weights are frozen to prevent further changes. Offline distillation is typical of many KD approaches for LLMs, in which the teacher is often a larger proprietary model for which model weights cannot be changed.

In some circumstances, a suitably pre-trained and adequately performing teacher model might not be available, or a data scientist might want to tailor the teacher network to their specific use case. Online distillation schemes aim to simultaneously train both the teacher and student networks.

For example, Cioppa et al proposed an online distillation scheme for semantic segmentation models used in live sporting events, where visual circumstances might change throughout a match. It aimed to circumvent the tradeoff between a smaller network’s speed and a larger network’s accuracy by continuously training a slow, well-performing model on live match data while simultaneously distilling that larger model’s knowledge into a smaller, faster model deployed to generate outputs in real time.⁶

In self-distillation, one network acts as both teacher and student. Whereas conventional knowledge distillation entails the transfer of knowledge from one model to another, self-distillation can be understood as the transfer of knowledge from a network’s deeper layers to the same network’s shallow layers.⁷

In self-distillation, multiple attention-based “shallow classifiers” are added to the model’s intermediate layers at varying depths. During training, the deeper-lying classifiers act as the teacher models and guide the training of the other attention-based modules through two kinds of distillation losses: a KL divergence metric loss on the outputs and an L2 regularization loss on the feature maps.

After the model is trained and ready for inference, all these shallow classifiers are dropped from the model. Essentially, this allows for the model to be larger and have greater capacity for pattern recognition during training, but then be smaller and consequently faster and more efficient when deployed.

The LLMs with the highest capabilities are, in most cases, too costly and computationally demanding to be accessible to many would-be users like hobbyists, startups or research institutions. Furthermore, despite their advanced performance and unique abilities, proprietary LLMs by their nature cannot be tailored to niche applications and specific use cases.

Furthermore, most commercially viable LLMs are too large and computationally demanding to be used locally on mobile phones or other edge devices. This presents various logistical, computational and privacy complications that would otherwise be circumvented with a smaller model that could be run directly on mobile devices. KD’s model compression thus presents a promising means to transfer the emergent qualities of large models to models small enough to be run on-device.