Knowledge Distillation for LLMs

Knowledge distillation enables deploying capable models at a fraction of the cost. A 70B parameter model distilled into a 7B model can retain 80-95% of the teacher's performance on many tasks while being 10x faster and cheaper to run.

The Distillation Framework

Temperature-Scaled Softmax

Distillation Loss

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

class DistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
    
    def forward(self, student_logits, teacher_logits, labels):
        # Soft loss (KL divergence between softened distributions)
        student_soft = F.log_softmax(student_logits / self.temperature, dim=-1)
        teacher_soft = F.softmax(teacher_logits / self.temperature, dim=-1)
        soft_loss = F.kl_div(
            student_soft, teacher_soft,
            reduction="batchmean"
        ) * (self.temperature ** 2)
        
        # Hard loss (standard cross-entropy with true labels)
        hard_loss = F.cross_entropy(student_logits, labels)
        
        return self.alpha * hard_loss + (1 - self.alpha) * soft_loss

Types of Knowledge Distillation

Response-Based Distillation

def response_distillation_step(teacher, student, batch, loss_fn):
    with torch.no_grad():
        teacher_outputs = teacher(batch["input_ids"])
    student_outputs = student(batch["input_ids"])
    loss = loss_fn(student_outputs.logits, teacher_outputs.logits, batch["labels"])
    return loss

Feature-Based Distillation

class FeatureDistillationLoss(nn.Module):
    def __init__(self, teacher_dim, student_dim):
        super().__init__()
        self.projection = nn.Linear(student_dim, teacher_dim)
    
    def forward(self, student_hidden, teacher_hidden):
        projected = self.projection(student_hidden)
        return F.mse_loss(projected, teacher_hidden)

Attention-Based Distillation

def attention_distillation_loss(teacher_attn, student_attn):
    """Match attention distributions between teacher and student."""
    teacher_attn = teacher_attn.mean(dim=1)  # Average over heads
    student_attn = student_attn.mean(dim=1)
    
    teacher_attn = F.softmax(teacher_attn, dim=-1)
    student_attn = F.softmax(student_attn, dim=-1)
    
    return F.kl_div(
        student_attn.log(), teacher_attn,
        reduction="batchmean"
    )

Chain-of-Thought Distillation

Distilling Reasoning Capabilities

def cot_distillation(teacher, student, problems, loss_fn):
    """Distill chain-of-thought reasoning from teacher to student."""
    # Teacher generates reasoning traces
    teacher_traces = []
    for problem in problems:
        trace = teacher.generate(
            problem["question"],
            max_new_tokens=1024,
            temperature=0.3
        )
        teacher_traces.append(trace)
    
    # Student learns to produce same reasoning
    total_loss = 0
    for problem, trace in zip(problems, teacher_traces):
        student_output = student(problem["question"], labels=trace)
        total_loss += student_output.loss
    
    return total_loss / len(problems)

Selective Distillation

def selective_distillation(teacher, student, dataset, confidence_threshold=0.9):
    selective_loss = 0
    count = 0
    
    for batch in dataset:
        with torch.no_grad():
            teacher_probs = F.softmax(teacher(batch["input_ids"]).logits, dim=-1)
            max_probs, _ = teacher_probs.max(dim=-1)
            confidence = max_probs.mean().item()
        
        if confidence > confidence_threshold:
            student_out = student(batch["input_ids"])
            loss = distillation_loss(student_out.logits, teacher(batch["input_ids"]).logits, batch["labels"])
            selective_loss += loss
            count += 1
    
    return selective_loss / max(count, 1)

Multi-Teacher Distillation

def multi_teacher_distillation(teachers, student, batch, weights):
    combined_teacher_loss = 0
    for teacher, weight in zip(teachers, weights):
        with torch.no_grad():
            teacher_logits = teacher(batch["input_ids"])
        student_logits = student(batch["input_ids"])
        loss = F.kl_div(
            F.log_softmax(student_logits / 4.0, dim=-1),
            F.softmax(teacher_logits / 4.0, dim=-1),
            reduction="batchmean"
        )
        combined_teacher_loss += weight * loss
    return combined_teacher_loss

Practical Distillation Pipeline

Full Training Loop

def distill(teacher, student, train_dataloader, val_dataloader, epochs=3, lr=5e-5):
    optimizer = torch.optim.AdamW(student.parameters(), lr=lr)
    loss_fn = DistillationLoss(temperature=4.0, alpha=0.3)
    
    teacher.eval()
    student.train()
    
    for epoch in range(epochs):
        total_loss = 0
        for batch in train_dataloader:
            with torch.no_grad():
                teacher_outputs = teacher(batch["input_ids"])
            
            student_outputs = student(batch["input_ids"])
            loss = loss_fn(
                student_outputs.logits,
                teacher_outputs.logits,
                batch["labels"]
            )
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch + 1}: Loss = {avg_loss:.4f}")
        
        # Evaluate
        evaluate(student, val_dataloader)

Distillation vs Other Compression Methods

Practice Exercises

1. Distillation Design: Design a distillation pipeline to compress a 70B teacher into a 7B student. What temperature, alpha, and training data would you use?

2. CoT Distillation: Implement chain-of-thought distillation for mathematical reasoning. How would you measure whether reasoning capabilities transfer effectively?

3. Multi-Teacher Ensemble: If you have a general-purpose teacher and a code-specialized teacher, how would you combine their knowledge for a student that needs both capabilities?

4. Distillation Analysis: Compare distillation from a 70B model vs. training a 7B model from scratch on the same data. What are the tradeoffs?

Key Takeaways

What to Learn Next