Knowledge distillation explained: teacher and student

Knowledge distillation is a model compression technique where a smaller student network learns to approximate the output probability distribution of a larger teacher network, typically by minimizing a loss function that combines hard labels with the teacher’s soft targets. The teacher’s logits contain “dark knowledge” — the relative probabilities across incorrect classes — that encodes structural relationships the student would struggle to learn from ground-truth labels alone. This transfers the teacher’s generalization ability into a model small enough for latency- or cost-constrained deployment.

How knowledge distillation works

The standard formulation comes from Hinton, Vinyals, and Dean (2015). Given a teacher model $T$ and student model $S$, both producing logits $z_T$ and $z_S$ for input $x$, the distillation loss combines two terms:

$$\mathcal{L} = \alpha \cdot \mathcal{L}{CE}(y, \sigma(z_S)) + (1-\alpha) \cdot \tau^2 \cdot \mathcal{L}{KL}(\sigma(z_T/\tau), \sigma(z_S/\tau))$$

Where $\mathcal{L}{CE}$ is cross-entropy with hard labels $y$, $\mathcal{L}{KL}$ is KL divergence between softened teacher and student distributions, $\tau > 1$ is a temperature parameter that sharpens or flattens the softmax, and $\alpha$ balances the two objectives.

The temperature $\tau$ is critical. At $\tau=1$, the softmax concentrates probability mass on the top class, discarding the inter-class relationships. At higher $\tau$ (typically 2–10), the distribution spreads, revealing that “this image is 90% dog, 8% wolf, 2% cat” — information the student exploits to learn decision boundaries that mirror the teacher’s.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    """
    Standard knowledge distillation loss.
    
    Args:
        student_logits: [batch, num_classes] raw outputs from student
        teacher_logits: [batch, num_classes] raw outputs from teacher (no grad)
        labels: [batch] ground truth class indices
        temperature: softmax temperature for soft targets
        alpha: weight for hard label loss (1-alpha for soft target loss)
    """
    # Hard label loss
    hard_loss = F.cross_entropy(student_logits, labels)
    
    # Soft target loss: KL(student || teacher) with temperature scaling
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
    
    # Scale by T^2 per Hinton et al. to maintain gradient magnitude
    return alpha * hard_loss + (1 - alpha) * soft_loss * (temperature ** 2)

Training proceeds in two stages typically: first train the teacher to convergence (or use a pre-trained checkpoint), freeze its weights, then train the student on the same data using the combined loss. The teacher never sees gradients from the student — it’s a one-way transfer.

Why it matters for LLM inference

For large language models, distillation addresses three production constraints simultaneously: latency, memory footprint, and per-token cost. A 7B parameter model distilled from a 70B teacher can retain 90–95% of the teacher’s benchmark performance while running 5–10× faster on the same hardware. The student fits on a single GPU where the teacher requires tensor parallelism across multiple devices.

The economics are straightforward. If your workload serves 1M tokens/day at $0.002/1K tokens for a 70B model versus $0.0003/1K for a distilled 7B equivalent, that’s roughly $60/day vs $9/day — a difference that compounds to ~$18K/year per workload. At scale, distillation is often the highest-ROI optimization available.

But the gains aren’t free. Distillation quality depends on:

• Data coverage: The student only learns what the teacher demonstrates on your training distribution. Out-of-distribution behavior degrades faster than the teacher’s.
• Capacity gap: A 1B student cannot fully absorb a 70B teacher. The “capacity gap” literature suggests staying within 10–20× parameter ratio for best results.
• Training compute: You still need to run forward passes through the teacher on your entire training set. For 1T tokens, that’s non-trivial GPU-hours.

Concrete example: distilling a code generation model

Suppose you have a 34B parameter code model (the teacher) that’s too slow for your IDE autocomplete feature. You want a 1.3B student that runs locally on a developer’s machine. Here’s a minimal training loop using Hugging Face Transformers:

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from torch.utils.data import Dataset
import torch

class DistillationDataset(Dataset):
    def __init__(self, tokenizer, texts, max_length=2048):
        self.tokenizer = tokenizer
        self.texts = texts
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            truncation=True,
            max_length=self.max_length,
            padding='max_length',
            return_tensors='pt'
        )
        return {k: v.squeeze(0) for k, v in encoding.items()}

# Load models
teacher = AutoModelForCausalLM.from_pretrained("bigcode/starcoder2-34b", device_map="auto", torch_dtype=torch.bfloat16)
student = AutoModelForCausalLM.from_pretrained("bigcode/starcoder2-1b", device_map="auto", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoder2-1b")
tokenizer.pad_token = tokenizer.eos_token

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

class DistillationTrainer(Trainer):
    def __init__(self, teacher_model, temperature=2.0, alpha=0.5, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher = teacher_model
        self.temperature = temperature
        self.alpha = alpha
    
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        input_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask")
        
        # Student forward
        student_outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        student_logits = student_outputs.logits
        
        # Teacher forward (no grad)
        with torch.no_grad():
            teacher_outputs = self.teacher(input_ids=input_ids, attention_mask=attention_mask)
            teacher_logits = teacher_outputs.logits
        
        # Shift for causal LM: predict next token
        shift_student = student_logits[..., :-1, :].contiguous()
        shift_teacher = teacher_logits[..., :-1, :].contiguous()
        shift_labels = input_ids[..., 1:].contiguous()
        
        # Mask padding
        mask = (shift_labels != tokenizer.pad_token_id).float()
        num_valid = mask.sum()
        
        # Hard loss (cross entropy on labels)
        hard_loss = F.cross_entropy(
            shift_student.view(-1, shift_student.size(-1)),
            shift_labels.view(-1),
            ignore_index=tokenizer.pad_token_id,
            reduction='none'
        )
        hard_loss = (hard_loss * mask.view(-1)).sum() / num_valid
        
        # Soft loss (KL divergence)
        soft_teacher = F.softmax(shift_teacher / self.temperature, dim=-1)
        soft_student = F.log_softmax(shift_student / self.temperature, dim=-1)
        soft_loss = F.kl_div(soft_student, soft_teacher, reduction='none').sum(-1)
        soft_loss = (soft_loss * mask).sum() / num_valid
        soft_loss *= self.temperature ** 2
        
        loss = self.alpha * hard_loss + (1 - self.alpha) * soft_loss
        
        return (loss, student_outputs) if return_outputs else loss

# Training data: your codebase + public permissive code
train_texts = load_your_code_corpus()  # implement this
train_dataset = DistillationDataset(tokenizer, train_texts)

training_args = TrainingArguments(
    output_dir="./student-distilled",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=2e-5,
    warmup_ratio=0.05,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    remove_unused_columns=False,
)

trainer = DistillationTrainer(
    teacher_model=teacher,
    temperature=2.0,
    alpha=0.5,
    model=student,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()
student.save_pretrained("./student-distilled-final")

Key practical notes for this setup:

• Use bf16 on Ampere+ GPUs; fp16 with gradient scaling on older hardware
• Gradient accumulation simulates larger batch sizes without OOM
• The ignore_index in cross-entropy handles padding correctly
• Temperature 2.0 works well for causal LM; classification tasks often use 4–10
• Log both loss components separately to monitor balance

Common misconceptions

“Distillation is just label smoothing”

Label smoothing replaces hard targets with a uniform distribution over incorrect classes: $y_{smooth} = (1-\epsilon)y_{hard} + \epsilon/K$. Distillation uses the teacher’s actual predictions as soft targets, which are highly non-uniform and encode semantic similarity. The teacher says “this function returns a list, not a dict” — label smoothing says “it’s probably not a float.” These are not the same signal.

“You need the teacher’s training data”

You need representative data, not the teacher’s exact training set. The student learns the teacher’s function approximation on your domain. If you distill a general-purpose 70B model for SQL generation, feed it your schema and query logs — not The Pile. Domain-specific distillation often outperforms general distillation on the target task because the teacher’s capacity focuses where you need it.

“The student architecture must match the teacher”

Architecture mismatch is common and often beneficial. Distilling a dense transformer into a MoE student, or a decoder-only into an encoder-decoder, works if the student has sufficient capacity. The loss operates on logits, not hidden states. What matters is output-space alignment, not architectural symmetry.

“Higher temperature is always better”

Temperature controls the signal-to-noise ratio in soft targets. Too high ($\tau > 20$) flattens everything toward uniform, losing the relative probabilities that carry the teaching signal. Too low ($\tau \approx 1$) collapses to near-one-hot, reverting to hard-label behavior. Sweep $\tau \in [1, 10]$ on a validation set; the optimum is task-dependent.

“Distillation preserves all capabilities”

It doesn’t. The student inherits the teacher’s behavior on the distillation data distribution. Capabilities requiring reasoning chains the teacher never demonstrates (or demonstrates poorly) won’t transfer. Emergent abilities like in-context learning or multi-step reasoning degrade disproportionately. Evaluate the student on your actual task distribution, not just benchmarks.

“One distillation run is enough”

Iterative distillation — using the previous student as the next teacher — can close the capacity gap further. Each round compresses knowledge into a smaller model, but diminishing returns hit fast. Two rounds is typical; three rarely justifies the compute. Self-distillation (student teaches itself via ensembles or dropout noise) is a separate technique with different trade-offs.

When to reach for distillation

Use distillation when:

• You have a capable teacher (your own trained model or API-accessible) and need a smaller deployment artifact
• Latency or cost constraints rule out the teacher at serving scale
• You have domain data to align the student to your workload
• You can afford the one-time teacher inference cost over your training corpus

Skip distillation when:

• The teacher is only accessible via a closed API with no logit access (you need logits for KL loss; prompt-based distillation is a weaker alternative)
• Your task is so narrow that fine-tuning a small base model from scratch matches or beats distillation
• The capacity gap exceeds ~20× — the student simply cannot represent the teacher’s function
• You need the teacher’s full generality; distillation specializes

Closing note

Knowledge distillation is the most reliable path from “model that works in a notebook” to “model that serves production traffic at reasonable cost.” It’s not magic — it’s a second training run with a better supervision signal. The teacher’s soft targets are a denser gradient field than hard labels, and the student follows it into a smaller, faster basin of attraction. Treat it as a standard step in your model lifecycle, not a last-resort optimization.