What is distillation?

What is Distillation?

Distillation is a machine learning technique where a compact "student" neural network is trained to reproduce the behavior, representations, or reasoning steps of a larger, high-capacity "teacher" model. Rather than training the student strictly on raw labels or tokens from scratch, distillation leverages the rich probability distributions and intermediate artifacts generated by the teacher to transfer its learned knowledge. This approach yields smaller models that retain much of the larger model's accuracy while consuming substantially less memory and compute at inference time.

Think of distillation like learning from a master chef: instead of discovering culinary chemistry strictly by trial and error from raw recipes, an apprentice observes the master's subtle choices, scoring rubrics, and step-by-step techniques, mastering the craft far faster than they would alone.

How it works

Distillation operates through two primary paradigms: classical logit-based distillation and output-based (sequence-level) distillation.

1. Logit-Based Distillation (White-Box): In the framework established by Geoffrey Hinton, Oriol Vinyals, and Jeff Dean (2015), the student learns from the teacher's continuous output logits rather than hard binary targets.

   • Temperature Scaling: Standard softmax converts output logits into probabilities where the top class dominates near 1.0, masking relationships between other classes. A temperature parameter T > 1 softens these logits: p_i = exp(z_i / T) / sum(exp(z_j / T)). Higher temperatures expose the "dark knowledge"—such as whether a handwritten digit looks slightly like an 8 when classified as a 3.
   • Loss Computation: The training objective combines Kullback-Leibler (KL) divergence between softened student and teacher distributions with standard cross-entropy loss against hard ground-truth labels. The KL loss term is scaled by T^2 to keep gradient magnitudes balanced across temperature changes.
   • Feature Matching: Advanced setups can also match intermediate hidden states or attention matrices between corresponding layers using Mean Squared Error (MSE) or cosine similarity loss.

2. Sequence-Level / Response Distillation (Black-Box): Common in generative Large Language Models (LLMs), the teacher generates high-quality synthetic completions, detailed solutions, or chain-of-thought reasoning paths given prompts. The student is then fine-tuned on these generated text trajectories using standard autoregressive cross-entropy loss without requiring direct access to the teacher's internal weights or logits.

Distillation vs Quantization

Distillation and quantization both reduce model deployment costs, but they operate on fundamentally different principles. Quantization preserves the model's original architecture and parameter count while reducing numerical precision (for example, converting 16-bit floating-point weights to 8-bit or 4-bit integers). In contrast, distillation changes the model architecture entirely, training a smaller, fewer-parameter network to emulate the teacher.

Because their mechanics are distinct, engineering teams frequently combine them: a 70-billion parameter model can be distilled into an 8-billion parameter student, which is then quantized to 4-bit weights for edge deployment.

A practical example

The following PyTorch implementation demonstrates the standard logit distillation loss function combining softened KL divergence with ground-truth cross-entropy:

import torch
import torch.nn as nn
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Soft target loss: match student probabilities to teacher probabilities at temperature T
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    kl_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)

    # Hard target loss: standard cross-entropy on ground-truth labels
    ce_loss = F.cross_entropy(student_logits, labels)

    # Return weighted combination
    return alpha * kl_loss + (1.0 - alpha) * ce_loss

# Example forward pass with batch_size=4, num_classes=10
student_logits = torch.randn(4, 10, requires_grad=True)
teacher_logits = torch.randn(4, 10)
labels = torch.tensor([1, 4, 2, 7])

loss = distillation_loss(student_logits, teacher_logits, labels, temperature=3.0, alpha=0.7)
loss.backward()
print(f"Distillation Loss: {loss.item():.4f}")

Running this code outputs a scalar loss value that drives gradient updates through the student model while keeping teacher weights fixed, penalizing the student when its relative probability spread deviates from the teacher's distribution.

Tools and implementations

Tool	What it is for	Link
PyTorch torch.nn.KLDivLoss	Core deep learning primitive for computing Kullback-Leibler divergence between softened logits	https://pytorch.org
Hugging Face Transformers	Framework supporting task distillation and distilled model architectures like DistilBERT	https://github.com/huggingface/transformers
TextBrewer	PyTorch-based knowledge distillation toolkit designed specifically for NLP and Transformer models	https://github.com/airaria/TextBrewer
Microsoft DeepSpeed	Distributed training engine with integrated model compression and distillation pipelines	https://github.com/microsoft/DeepSpeed

Where it goes wrong

• Capacity Gap Collapse: If the student model is too small or architecturally constrained relative to the teacher, it cannot capture complex decision boundaries, leading to underfitting and degraded performance across subtle edge cases.
• Error and Hallucination Amplification: In sequence-level LLM distillation, students inherit the teacher's factual errors, hallucinations, and stylistic biases without learning the underlying verification mechanisms.
• Loss of Output Calibration: Training primarily on soft teacher distributions can alter prediction confidence, requiring temperature recalibration before deployment in safety-critical classification tasks.
• Distribution Shift Vulnerability: A distilled student often mimics the teacher accurately on data similar to the distillation set, but exhibits sharper performance degradation when faced with out-of-distribution prompts.

Why it matters now

Distillation bridges the gap between frontier model research and cost-effective production deployment, allowing organizations to run reasoning and language capabilities on consumer GPUs, mobile devices, and local servers without datacenter-scale latency or hosting costs.

In January 2025, DeepSeek demonstrated the power of this paradigm by releasing the DeepSeek-R1 series, showing that fine-tuning smaller base models on 800,000 reasoning trajectories curated from DeepSeek-R1 allowed DeepSeek-R1-Distill-Qwen-32B to score 72.6% on AIME 2024 and 94.3% on MATH-500, outperforming far larger non-distilled models.

FAQs

What temperature parameter should I use for logit distillation? Temperature values typically range between T = 2.0 and T = 5.0 for classification tasks. Higher temperatures produce smoother probability distributions over secondary classes, whereas temperatures approaching 1.0 collapse back to hard argmax distributions.

Can you distill an LLM if you do not have access to its logits? Yes, through sequence-level (black-box) distillation. You prompt the teacher model to generate step-by-step reasoning solutions, filter the outputs for quality, and fine-tune your student model directly on those input-output pairs using standard supervised fine-tuning (SFT).

How much smaller was DistilBERT compared to the original BERT model? In the 2019 NeurIPS workshop paper by Sanh et al., DistilBERT reduced BERT's parameter count by 40% (66M vs 110M parameters) while retaining approximately 97% of its language understanding performance on GLUE benchmarks and running 60% faster at inference.

Does distilling a model require retraining the teacher? No, the teacher model is kept completely frozen with gradient tracking disabled during the entire distillation process. Gradients are computed and applied exclusively to the student model's parameters.