Distillation & Knowledge Transfer

Overview

Knowledge distillation is a model compression technique where a smaller “student” model is trained to replicate the behavior of a larger, more capable “teacher” model. Rather than training the student from scratch on raw data, distillation transfers the teacher’s learned knowledge — including nuanced relationships between outputs that aren’t captured by simple labels — into a compact model suitable for production deployment.

This technique has become critical in the LLM era because frontier models are often too large and expensive to deploy directly. A 70B-parameter model might cost $1–2 per million tokens to serve, while a distilled 7B student could deliver 80–90% of the quality at 10% of the cost. Companies routinely distill capabilities from proprietary models like GPT-4 or Claude into smaller open-weight models optimized for specific tasks.

For AI engineers, distillation is a core tool for bridging the gap between what’s possible in research and what’s practical in production. It’s the primary technique behind most efficient open-source models and a key component of any cost-optimization strategy.

How It Works

The Teacher-Student Paradigm

The fundamental idea is straightforward: instead of training a model on hard labels (e.g., “this token is ‘cat’”), you train it on the teacher’s full probability distribution over all possible outputs. These “soft targets” contain far more information than hard labels because they encode the teacher’s uncertainty and the relationships between different outputs.

Temperature Scaling

A key insight in distillation is that raw softmax outputs are often too “peaked” — the top prediction dominates with very high probability, making the soft targets nearly indistinguishable from hard labels. Temperature scaling addresses this by dividing the logits by a temperature parameter T before applying softmax:

# Standard softmax (T=1): peaked distribution
softmax(logits / 1.0)  # → [0.95, 0.03, 0.01, 0.01, ...]

# High-temperature softmax (T=4): softer distribution
softmax(logits / 4.0)  # → [0.40, 0.20, 0.15, 0.12, ...]

# The softer distribution reveals more about the teacher's
# internal knowledge about relationships between outputs

Higher temperatures (typically T = 2–10) produce softer distributions that reveal more of the teacher’s knowledge. However, very high temperatures can over-flatten the distribution, making it harder for small students to learn. The optimal temperature depends on the capacity gap between teacher and student — smaller students benefit from lower temperatures.

The Distillation Loss Function

The student is trained with a combined loss that balances learning from the teacher (soft targets) with learning from the ground truth (hard labels):

# Distillation loss function
loss = alpha * KL_divergence(
    student_softmax(logits_s / T),
    teacher_softmax(logits_t / T)
) * T² + (1 - alpha) * cross_entropy(logits_s, true_labels)

# alpha: weight between soft and hard losses (typically 0.5–0.9)
# T: temperature (typically 2–10)
# T²: scaling factor to keep gradient magnitudes comparable

The T² factor is important: as temperature increases, the magnitudes of the softmax gradients decrease by 1/T², so multiplying by T² compensates and keeps the soft-target gradients on the same scale as the hard-label gradients.

Distillation Schemes

There are three primary approaches to when and how distillation happens:

• Offline distillation: The teacher is fully trained first, then its outputs are cached and used to train the student. This is the most common approach — simple, parallelizable, and works with black-box API teachers (you only need the output distributions).
• Online distillation: Teacher and student train simultaneously, with the teacher’s evolving knowledge continuously transferred. This can produce better students but requires more compute and careful coordination.
• Self-distillation: A model distills knowledge from itself — for example, from deeper layers to shallower layers, or from a larger version of the same architecture to a smaller one. This has shown surprising effectiveness, sometimes improving performance even without a separate teacher.

Beyond Logit Matching

Modern distillation goes beyond just matching output distributions:

• Intermediate layer distillation: The student also learns to match the teacher’s internal representations at specific hidden layers, capturing more structured knowledge about how the teacher processes information.
• Attention transfer: The student mimics the teacher’s attention patterns, learning which parts of the input the teacher focuses on for different predictions.
• Feature-based distillation: Alignment of intermediate feature maps between teacher and student using projection layers when their dimensions differ.

Practical Applications

LLM Distillation in Practice

Knowledge distillation is the primary technique behind many successful production models:

• DistilBERT: A 66M-parameter model distilled from BERT-base (110M), retaining 97% of BERT’s performance while being 60% faster. One of the earliest and most influential examples of LLM distillation.
• Phi models (Microsoft): The Phi series of small language models (1.3B–14B parameters) are trained using synthetic data generated by larger models — a form of distillation where the teacher’s knowledge is embedded in the training data rather than transferred via soft targets directly.
• OpenAI’s distillation API: OpenAI provides a distillation API that lets you fine-tune smaller models (like GPT-4o-mini) using outputs from GPT-4o, making it straightforward to create task-specific distilled models.
• DeepSeek-R1-Distill series: DeepSeek distilled their R1 reasoning model into Qwen and Llama-based models at 1.5B, 7B, 8B, 14B, 32B, and 70B sizes, demonstrating that complex chain-of-thought reasoning can be effectively transferred to much smaller models.

When to Use Distillation

• Cost reduction: When you have a task where a large model works well but is too expensive to serve at scale. Distill it into a smaller model that handles 80–90% of cases, and fall back to the large model for the rest.
• Latency requirements: When you need sub-100ms inference for real-time applications like autocomplete, classification, or routing decisions.
• Edge deployment: When models need to run on mobile devices, embedded systems, or in environments with limited compute resources.
• Domain specialization: When you want to create a small model that excels at a narrow domain by distilling a generalist model’s knowledge for that specific area.

Distillation Pipeline

A typical production distillation pipeline involves: (1) selecting a teacher model that performs well on your task, (2) generating a dataset of teacher outputs (either via API calls or by running inference on your data), (3) choosing a student architecture (usually the same family, just smaller), (4) training the student on the teacher’s outputs with temperature-scaled soft targets, and (5) evaluating the student against your quality benchmarks to ensure acceptable performance.

Key Takeaways

• Knowledge distillation trains a smaller student model to mimic a larger teacher model’s behavior, transferring “dark knowledge” encoded in soft probability distributions.
• Temperature scaling softens the teacher’s output distribution to reveal more relational information between classes, with typical values of T=2–10.
• The distillation loss combines KL divergence on soft targets (teacher knowledge) with cross-entropy on hard labels (ground truth), weighted by alpha.
• Three distillation schemes exist: offline (most common, teacher trained first), online (simultaneous training), and self-distillation (model learns from itself).
• Modern techniques go beyond output matching to include intermediate layer distillation, attention transfer, and feature alignment.
• Distillation is the primary technique behind cost-effective production deployment, enabling 80–90% of a large model’s quality at 10% of the compute cost.