Training & Alignment

Knowledge Distillation

Knowledge distillation trains a small student model to mimic a larger teacher model, transferring the teacher's behavior through its full probability distribution so the student keeps much of the quality at a fraction of the cost.

What Knowledge Distillation Is

Knowledge distillation is a model compression technique in which a small, fast student model is trained to reproduce the behavior of a larger, more accurate teacher model. Rather than learning only from the dataset's ground-truth labels, the student also learns to match the teacher's output, which carries far more information than a single correct answer. The goal is a student that retains most of the teacher's quality while costing less to run.

The idea was formalized by Geoffrey Hinton, Oriol Vinyals, and Jeff Dean in the 2015 paper "Distilling the Knowledge in a Neural Network" (arXiv:1503.02531). Their motivation was deployment: a large ensemble of models can be accurate but is too cumbersome and computationally expensive to serve, so they compress its knowledge into one smaller model that approximates the ensemble's predictions.

Distillation differs from ordinary supervised learning because the supervision signal is the teacher's full prediction, not just a label. A teacher that classifies an image of a dog might assign a tiny probability to "wolf" and almost zero to "car"; those relative probabilities encode how the teacher relates the classes. Hinton and colleagues called this extra signal dark knowledge, and capturing it is what lets a student exceed what it could learn from labels alone.

Teacher: a large, accurate, expensive model whose behavior is the training target.
Student: a smaller, cheaper model trained to imitate the teacher's outputs.
Supervision signal is the teacher's prediction, which is richer than a one-hot label.
Primary motivation is cheaper, faster deployment at near-teacher quality.

Soft Targets, Temperature, and Dark Knowledge

A trained classifier outputs logits that a softmax converts into class probabilities. Normally the softmax produces a sharp distribution that puts nearly all mass on the predicted class, which hides the relationships between the other classes. Knowledge distillation softens this distribution by dividing the logits by a temperature T greater than 1 before the softmax, spreading probability across classes so the secondary, non-top predictions become visible.

These softened outputs are called soft targets, and the relative probabilities among the wrong classes are the dark knowledge the student learns from. With T equal to 1 the softmax is standard; as T grows the distribution becomes softer and reveals more about which classes the teacher found similar. Hinton and colleagues used temperatures ranging from 1 to 20 in their experiments, choosing T to expose useful structure without flattening the distribution into noise.

Soft targets transfer information that hard labels cannot. A model told only "this is a 7" learns nothing about how 7s resemble 1s; a model told the teacher assigned 7 a high probability and 1 a small but nonzero one learns the geometry of the classes. This is why a student trained on soft targets can generalize from far fewer examples than one trained on labels alone, and why distillation often outperforms training the small model directly.

p_i = exp(z_i / T) / Σ_j exp(z_j / T)

Softmax with temperature T. Logits z_i are divided by T before exponentiation; T = 1 is the standard softmax and larger T produces a softer distribution.

Dividing logits by temperature T before softmax softens the distribution.
Soft targets expose the teacher's relative probabilities across all classes.
Dark knowledge is the information carried in those non-top probabilities.
Hinton et al. used temperatures from 1 to 20 in their original experiments.

The Distillation Loss

The student is trained on a weighted combination of two objectives. The first is a soft-target term: the divergence between the student's softened probabilities and the teacher's softened probabilities, both computed at the same high temperature T. The second is a standard hard-label term: ordinary cross-entropy between the student's predictions at T equal to 1 and the true labels. A coefficient alpha balances the two.

A subtlety arises from the temperature. The gradients produced by the soft-target term scale with 1/T-squared, so when soft and hard targets are mixed the soft term is multiplied by T-squared to keep the two contributions roughly balanced as T changes. Without this correction, raising the temperature would silently shrink the influence of the teacher signal.

In practice the soft-target divergence is computed as a KL divergence (or equivalently a cross-entropy) between the temperature-softened student and teacher distributions. The hard-label term anchors the student to the ground truth, which matters when the teacher is wrong; the soft-target term supplies the dark knowledge. The relative weight alpha and the temperature T are the two main hyperparameters of basic distillation.

L = α · T² · KL(soft_teacher ‖ soft_student) + (1 − α) · CE(hard_labels, student)

The distillation loss combines a temperature-softened soft-target KL term, scaled by T² to correct the 1/T² gradient shrinkage, with a hard-label cross-entropy term.

python

import torch.nn.functional as F

def kd_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):
    # Soft-target term: KL between temperature-softened distributions.
    soft_student = F.log_softmax(student_logits / T, dim=1)
    soft_teacher = F.softmax(teacher_logits / T, dim=1)
    soft = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (T * T)

    # Hard-label term: ordinary cross-entropy at T = 1.
    hard = F.cross_entropy(student_logits, labels)

    return alpha * soft + (1.0 - alpha) * hard

Minimal PyTorch knowledge-distillation loss. F.kl_div expects log-probabilities for the input and probabilities for the target; the (T * T) factor restores the temperature-scaled gradient magnitude.

Soft-target term: KL divergence between softened student and teacher, scaled by T-squared.
Hard-label term: ordinary cross-entropy against the true labels at T = 1.
Alpha balances the two; the T-squared factor keeps gradients balanced across temperatures.
Hard labels keep the student grounded when the teacher is mistaken.

Variants and Where Knowledge Lives

Distillation methods differ in what part of the teacher the student imitates. The survey "Knowledge Distillation: A Survey" by Gou, Yu, Maybank, and Tao groups them into three kinds of knowledge. Response-based distillation, the original Hinton formulation, matches the teacher's final output probabilities. Feature-based distillation matches intermediate representations such as hidden activations or attention maps. Relation-based distillation matches the relationships between layers or between examples rather than any single output.

Methods also differ in how teacher and student are trained relative to each other. Offline distillation uses a fixed, pre-trained teacher, which is the most common setup. Online distillation trains a group of models together so they teach each other, useful when no strong pre-trained teacher exists. Self-distillation trains a model to transfer knowledge into itself, for example from deeper to shallower layers, with no separate teacher network.

DistilBERT, introduced by Victor Sanh and colleagues in 2019 (arXiv:1910.01108), is a widely cited application. It distills BERT using a triple loss that combines the soft-target distillation loss, the masked-language-modeling loss, and a cosine-embedding loss aligning the student and teacher hidden states. The result reduces the size of BERT by 40 percent and runs 60 percent faster while retaining 97 percent of its language-understanding capabilities.

Response-based: match final outputs; feature-based: match hidden activations; relation-based: match cross-layer or cross-sample relations.
Offline uses a fixed teacher; online trains peers together; self-distillation teaches a model from itself.
DistilBERT (Sanh et al., 2019) uses a distillation, masked-LM, and cosine-embedding triple loss.
DistilBERT cuts BERT's size by 40 percent and is 60 percent faster while keeping 97 percent of its language understanding.

Distillation for Large Language Models

For modern large language models the classic logit-matching recipe is often replaced by sequence-level distillation, sometimes called data distillation. A strong teacher model generates outputs, such as completions, reasoning traces, or instruction-following responses, and a smaller student is trained on those generated sequences as if they were ground-truth data. The student learns to imitate the teacher's text rather than its raw probability vectors, which sidesteps the need for matching logits across a huge vocabulary.

This approach has become a standard way to build smaller, cheaper models that approximate the behavior of a frontier model. Because the supervision is the teacher's generated text, the student inherits the teacher's style, formatting, and much of its task competence, though it can also inherit the teacher's mistakes and biases. The quality of the generated training data is the dominant factor in the result.

Sequence-level and response-based distillation sit on a spectrum: response-based distillation matches a probability distribution over next tokens, while data distillation trains on sampled sequences. Both aim at the same outcome, a compact model that mimics a larger one, and both are frequently combined with parameter-efficient fine-tuning so the student can be adapted at low cost.

Sequence-level (data) distillation trains a student on text generated by a teacher.
It avoids vocabulary-wide logit matching and transfers style, format, and task skill.
The student inherits the teacher's strengths and its errors and biases.
Often paired with parameter-efficient fine-tuning to keep adaptation cheap.

Distillation Among Compression Techniques

Knowledge distillation is one of three main families of model compression, alongside quantization and pruning, and the three are complementary rather than mutually exclusive. Distillation produces a genuinely smaller architecture by training a new, compact model. Quantization keeps the architecture but stores and computes weights at lower numerical precision, such as 8-bit or 4-bit integers instead of 16- or 32-bit floats. Pruning removes weights or whole structures that contribute little to the output.

Because they act on different aspects of a model, they stack. A common pipeline distills a large model into a smaller student, then quantizes the student to shrink its memory footprint and speed up inference further, and may prune redundant parameters as well. Each step trades a little accuracy for efficiency, and the combination can compress a model far more than any single technique.

Choosing among them depends on the constraint. Distillation is attractive when a faster architecture is needed and a strong teacher is available to supervise it. Quantization is attractive when the architecture must stay fixed but memory and latency must drop. Pruning targets redundancy within an existing network. Distillation is unique in that it can also improve a small model's accuracy beyond what training it directly would achieve, by injecting the teacher's dark knowledge.

Distillation trains a new smaller model; quantization lowers numeric precision; pruning removes weights.
The three techniques are complementary and commonly stacked in a pipeline.
Pick distillation for a faster architecture, quantization for fixed-architecture memory and latency.
Distillation can raise a small model's accuracy, not only shrink it.

Key takeaways

Knowledge distillation trains a small student to mimic a large teacher, keeping much of the quality at lower cost.
The teacher's soft targets, its full softened probability distribution, carry dark knowledge that hard labels cannot.
Temperature T softens the softmax to expose secondary class probabilities; T = 1 is the standard softmax.
The KD loss mixes a T-squared-scaled soft-target KL term with a hard-label cross-entropy term, weighted by alpha.
DistilBERT retains 97 percent of BERT's language understanding with 40 percent fewer parameters and 60 percent faster inference.
Distillation, quantization, and pruning are complementary compression techniques that can be combined.

Frequently asked questions

The teacher is a large, accurate, expensive model whose behavior serves as the training target. The student is a smaller, cheaper model trained to reproduce the teacher's outputs. The aim is for the student to retain most of the teacher's quality while being faster and less costly to deploy.

A standard softmax produces a sharp distribution that hides how the teacher relates the non-top classes. Dividing the logits by a temperature T greater than 1 softens the distribution so those secondary probabilities, the dark knowledge, become visible. Hinton, Vinyals, and Dean used temperatures from 1 to 20 in their 2015 paper.

The gradients from the soft-target term scale with 1/T-squared, so raising the temperature would otherwise shrink the teacher signal's influence. Multiplying the soft-target loss by T-squared cancels that shrinkage, keeping the soft and hard target contributions roughly balanced as the temperature changes.

DistilBERT, from Sanh and colleagues in 2019, reduces BERT's size by 40 percent and runs 60 percent faster while retaining 97 percent of its language-understanding capabilities. It is trained with a triple loss combining the distillation loss, the masked-language-modeling loss, and a cosine-embedding loss that aligns student and teacher hidden states.

Modern LLM distillation often uses sequence-level or data distillation: a strong teacher generates outputs such as completions or reasoning traces, and a smaller student is trained on that generated text as if it were ground truth. This avoids matching logits across a large vocabulary and transfers the teacher's style and task competence, along with its errors.

All three compress models but act differently and complement each other. Distillation trains a new smaller architecture, quantization lowers numeric precision while keeping the architecture, and pruning removes low-value weights. They are commonly stacked, for example distilling a model and then quantizing the student to shrink it further.

Put the idea into practice

MemX is an AI memory agent built on these ideas: store anything, skip the folders, and find it again by asking in plain English.

Try MemX Free