What is knowledge distillation in machine learning?

Knowledge distillation is a model compression technique where a smaller student model learns to mimic the behavior of a larger teacher model by matching its output distributions, not just its hard predictions. The core insight is that the teacher’s probability distribution over classes — especially the low-probability assignments — carries “dark knowledge” about inter-class relationships that transfers more signal per training example than one-hot labels alone. This lets you deploy models that are 2-10x smaller with minimal accuracy loss.

How it works

The standard setup trains a student network to minimize a weighted combination of two losses: the standard cross-entropy against ground-truth labels, and a distillation loss that matches the teacher’s softened probabilities. The temperature parameter T controls how soft the distribution becomes — higher T reveals more structure in the tail probabilities.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    """
    student_logits: [batch, num_classes] raw outputs from student
    teacher_logits: [batch, num_classes] raw outputs from teacher (no grad)
    labels: [batch] integer class indices
    temperature: softens the probability distribution
    alpha: weight on distillation loss vs hard label loss
    """
    # Soft targets from teacher
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    
    # KL divergence scaled by T^2 (per Hinton et al. 2015)
    distill_loss = F.kl_div(student_log_probs, teacher_probs, reduction='batchmean') * (temperature ** 2)
    
    # Hard label loss
    hard_loss = F.cross_entropy(student_logits, labels)
    
    return alpha * distill_loss + (1 - alpha) * hard_loss

The temperature scaling is critical. At T=1, the teacher’s distribution is peaky — the top class gets ~0.99 probability and everything else near zero. At T=4 or T=8, the relative probabilities between incorrect classes become visible: “this image is 90% dog, 5% wolf, 3% fox, 2% cat” teaches the student more than “this image is dog.” The T² scaling on the KL term keeps gradient magnitudes stable as you change temperature.

In practice, you pre-train the teacher (or grab a checkpoint), freeze it, then train the student from scratch or from a smaller pretrained checkpoint. The teacher never updates during distillation.

Why it matters for LLM deployment

Large language models are the canonical use case. A 70B parameter model serves great quality but costs ~140 GB in FP16 — too big for most single-GPU inference, too slow for latency-sensitive paths. Distilling to 7B or 3B puts the model on a single A10G or even an Apple Silicon MacBook with 4-8 bit quantization.

The quality gap has narrowed dramatically. In 2023, distilled 7B models lagged their 70B teachers by 10-15 points on MMLU. By late 2024, careful distillation recipes (better data, longer training, reverse KL, logit matching at multiple layers) closed that to 2-4 points on many benchmarks. The tradeoff is now favorable for most product use cases.

But distillation isn’t free. You need:

• Teacher inference compute to generate soft labels (or logits) for your training corpus
• Student training compute — often 10-30% of pretraining from scratch
• A curriculum: random initialization works but starting from a smaller pretrained checkpoint converges faster and reaches higher quality

Concrete example: distilling a classifier

Suppose you have a BERT-large teacher (340M params) fine-tuned on a legal document classification task with 50 categories. You want a DistilBERT student (66M params) for edge deployment.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
from torch.utils.data import DataLoader

# Load teacher and student
teacher = AutoModelForSequenceClassification.from_pretrained("my-org/legal-bert-large")
student = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=50)

teacher.eval()
for param in teacher.parameters():
    param.requires_grad = False

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def collate(batch):
    texts = [b["text"] for b in batch]
    labels = torch.tensor([b["label"] for b in batch])
    enc = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
    return {**enc, "labels": labels}

# Training loop sketch
optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5, weight_decay=0.01)
temperature = 5.0
alpha = 0.5

for epoch in range(3):
    for batch in DataLoader(train_dataset, batch_size=16, collate_fn=collate, shuffle=True):
        labels = batch.pop("labels")
        
        with torch.no_grad():
            teacher_logits = teacher(**batch).logits
        
        student_logits = student(**batch).logits
        
        loss = distillation_loss(student_logits, teacher_logits, labels, temperature, alpha)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

A few practical notes from running this pattern in production:

Logit matching vs probability matching. Matching logits directly (MSE on logits) sometimes works better than KL on probabilities, especially when the teacher is overconfident. Try both.

Intermediate layer matching. Adding MSE loss between teacher and student hidden states at matched layers (e.g., every 3rd layer of teacher to every layer of student) helps when the capacity gap is large. This is “feature-based distillation.”

Data augmentation matters. The student sees the same data as the teacher. If you augment (back-translation, synonym replacement, dropout noise), the teacher’s soft labels on augmented examples provide a consistency signal that regularizes the student.

Common misconceptions

“Distillation is just label smoothing”

Label smoothing replaces one-hot targets with a uniform distribution over incorrect classes: (1-ε) on the true class, ε/(K-1) elsewhere. Distillation uses the teacher’s actual predictions, which encode semantic similarity — “wolf” gets higher probability than “toaster” for a dog image. That structure is the whole point.

“You need the teacher’s training data”

You need some data to generate teacher predictions, but it doesn’t have to be the teacher’s original training set. In-domain unlabeled data works well. For LLMs, people often distill on a mix of the original pretraining corpus (sampled) plus task-specific data. The teacher’s logits on out-of-domain data still transfer useful priors.

“Distillation always wins over training from scratch”

If you have compute to train a 7B model from scratch on 2T tokens, that 7B model will beat a 7B model distilled from 70B on the same 2T tokens — if the 70B teacher was trained on the same data distribution. Distillation shines when:

• You lack compute for full pretraining
• The teacher has seen data you can’t access (proprietary, licensed, or just massive)
• You need a model now and can’t wait for pretraining

“Higher temperature is always better”

Temperature is a hyperparameter. Too high (T>10) flattens the distribution so much that the signal-to-noise ratio drops. Too low (T<2) and you lose the dark knowledge. Sweep T in {2, 3, 4, 5, 6, 8} on a validation set.

“The student architecture must mirror the teacher”

Architecture mismatch is fine — BERT teacher to DistilBERT student, LLaMA teacher to Mamba student, transformer teacher to RNN student. The logits are architecture-agnostic. What matters is capacity: the student must have enough parameters to absorb the teacher’s function. A 100M parameter student distilling a 70B teacher will hit a ceiling no matter what.

Advanced variants worth knowing

Reverse KL distillation. Standard KL minimizes KL(teacher || student), which forces the student to cover all modes the teacher assigns mass to — mode-covering behavior. Reverse KL minimizes KL(student || teacher), which is mode-seeking: the student picks the highest-probability modes and ignores the rest. For generation tasks, reverse KL (or a symmetric JS divergence) often produces sharper, less hallucinated outputs.

Online distillation. Teacher and student train simultaneously. The teacher isn’t frozen; it’s an ensemble or a moving average of the student (EMA). This avoids the two-stage pipeline and lets the teacher adapt to the student’s learning dynamics. Used in self-distillation (e.g., BYOL, DINO for vision; SEED for LLMs).

Data-free distillation. Generate synthetic inputs that maximize teacher activation (via gradient ascent on input space), then distill on those. Useful when you can’t access the training data due to privacy or licensing. Quality lags data-based distillation but works surprisingly well for classification.

Quantization-aware distillation. Combine distillation with quantization: the student learns to match the teacher while simulating low-precision arithmetic (fake quantization during training). The student adapts its weights to be quantization-robust. This is how you get a 4-bit 7B model that doesn’t degrade.

When to skip distillation

• You have unlimited training compute and clean data. Train the target size from scratch. It’s simpler, no teacher dependency, and often matches or beats distillation at equal compute.
• The teacher is worse than your target. Distilling a bad teacher bakes in its errors. Verify teacher quality on your eval set first.
• Extreme compression ratios. Going from 70B → 1B loses too much capacity. The student cannot represent the teacher’s function class. Consider pruning + distillation, or accept a larger student.
• Latency-critical path with no batch size. If you’re serving batch=1 on CPU, even a distilled 3B model may be too slow. Look at quantization, speculative decoding, or non-autoregressive architectures instead.

A note on routing and fallbacks

In production systems that serve multiple model sizes, you often route easy queries to a small distilled model and escalate to the teacher (or a larger model) only when confidence is low or the query is complex. This cascading approach captures most of the teacher’s quality at a fraction of the average cost. The routing logic can be as simple as “if student entropy > threshold, forward to teacher” or a learned classifier on the student’s hidden states. Some gateways handle this routing automatically — forwarding requests to the best available model based on client directives and provider health — but the distillation itself remains a model-building concern, not an infrastructure one.

Summary checklist

If you’re evaluating distillation for a project:

1. Define the student budget — parameters, latency, memory, quantization target
2. Pick a teacher — best available model for your task, ideally same architecture family
3. Prepare distillation data — in-domain unlabeled data + task labels if available
4. Choose loss formulation — KL + CE, reverse KL, logit MSE, intermediate feature matching
5. Sweep temperature and alpha — validate on held-out set, not training loss
6. Evaluate on your actual metrics — not just perplexity or accuracy; test calibration, latency, edge cases
7. Quantize the student — distillation and quantization compose well; do both

Distillation is the most reliable lever for shrinking models without retraining from scratch. It’s not magic — it’s a second training run with a better supervision signal. But that signal makes all the difference when you’re pushing the Pareto frontier of quality vs. cost.