Language Model Distillation

Knowledge distillation is a technique for training smaller neural networks to perform like larger ones. The basic idea is simple: train a small "student" model to copy the behavior of a large "teacher" model. This lets you compress years of training and billions of parameters into something you can actually deploy. But making this work well requires understanding probability matching, intermediate representations, and training dynamics.

How Distillation Works

Hinton and his colleagues created knowledge distillation after noticing something useful: when a trained model makes predictions, it outputs a probability distribution that contains more information than just the final answer. Say a teacher model assigns 80% probability to the correct class, 15% to a similar class, and 5% to an unrelated class. This distribution tells you about relationships between classes that a simple correct/incorrect label throws away. The student learns not just what to predict, but how the teacher thinks about alternatives.

The distillation loss function combines two parts. First, it measures how well the student matches the teacher's soft predictions using KL divergence. Second, it makes sure the student still learns from the actual labels using cross-entropy loss. You control the balance between these with a parameter called alpha.

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

class DistillationLoss(nn.Module):
    def __init__(self, temperature=3.0, alpha=0.7):
        """
        Args:
            temperature: Controls softness of probability distributions
            alpha: Weight balancing distillation vs hard label loss
        """
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.kl_div = nn.KLDivLoss(reduction='batchmean')
        self.ce_loss = nn.CrossEntropyLoss()
    
    def forward(self, student_logits, teacher_logits, labels):
        # Soften distributions with temperature scaling
        student_soft = F.log_softmax(student_logits / self.temperature, dim=1)
        teacher_soft = F.softmax(teacher_logits / self.temperature, dim=1)
        
        # Distillation loss (KL divergence between soft distributions)
        distillation_loss = self.kl_div(student_soft, teacher_soft) * (self.temperature ** 2)
        
        # Standard cross-entropy with hard labels
        student_loss = self.ce_loss(student_logits, labels)
        
        # Combined objective
        return self.alpha * distillation_loss + (1 - self.alpha) * student_loss

Temperature is important here. Higher temperature makes the probability distribution "softer," revealing more about class relationships. At temperature 1, you get normal softmax. As temperature goes up, the distribution gets more uniform. The best temperature is usually between 2 and 5. The temperature squared term in the loss keeps gradients consistent across different temperature settings.

Dealing with Size Differences

The gap between teacher and student size fundamentally limits how well distillation works. If the student is too small, it can't capture what the teacher knows. If it's too large, you're not getting much compression. You need to find the point where the student is small enough to be useful but large enough to learn the important patterns.

Different approaches handle this gap in different ways. Basic distillation just matches output distributions. Intermediate distillation also aligns hidden layer representations. This works better for transformers where intermediate attention patterns encode important linguistic knowledge. Patient Knowledge Distillation goes further by matching relationships between layers, not just individual layers.

class IntermediateDistillation(nn.Module):
    def __init__(self, student_dim, teacher_dim, num_student_layers, num_teacher_layers):
        super().__init__()
        self.num_student_layers = num_student_layers
        self.num_teacher_layers = num_teacher_layers
        
        # Layer mapping strategy: map student layers to teacher layers
        self.layer_mapping = self._create_layer_mapping()
        
        # Projection layers if dimensions don't match
        if student_dim != teacher_dim:
            self.projections = nn.ModuleList([
                nn.Linear(student_dim, teacher_dim) 
                for _ in range(num_student_layers)
            ])
        else:
            self.projections = None
    
    def _create_layer_mapping(self):
        # Map student layers uniformly across teacher layers
        # For 6 student, 12 teacher: [1, 3, 5, 7, 9, 11]
        step = self.num_teacher_layers / self.num_student_layers
        return [int(i * step) for i in range(self.num_student_layers)]
    
    def forward(self, student_hidden_states, teacher_hidden_states):
        """
        Args:
            student_hidden_states: List of tensors [batch, seq_len, student_dim]
            teacher_hidden_states: List of tensors [batch, seq_len, teacher_dim]
        """
        total_loss = 0
        
        for student_idx, teacher_idx in enumerate(self.layer_mapping):
            student_hidden = student_hidden_states[student_idx]
            teacher_hidden = teacher_hidden_states[teacher_idx]
            
            # Project student to teacher dimension if needed
            if self.projections is not None:
                student_hidden = self.projections[student_idx](student_hidden)
            
            # MSE loss between intermediate representations
            layer_loss = F.mse_loss(student_hidden, teacher_hidden.detach())
            total_loss += layer_loss
        
        return total_loss / len(self.layer_mapping)

Transferring Attention Patterns

Attention mechanisms in transformers capture complex dependencies between input tokens, encoding structure, relationships, and context. Distilling these attention patterns is harder than distilling outputs because attention matrices are large and sensitive to architectural differences. Attention transfer methods usually focus on preserving the structure of attention rather than exact weights.

The attention transfer loss measures the distance between teacher and student attention distributions, typically using mean squared error or KL divergence. You need to handle different numbers of attention heads. Some approaches average attention across heads before computing loss. Others maintain head-specific alignments, which preserves more detail but requires careful mapping.

class AttentionTransfer(nn.Module):
    def __init__(self, student_heads, teacher_heads, use_head_mapping=True):
        super().__init__()
        self.student_heads = student_heads
        self.teacher_heads = teacher_heads
        self.use_head_mapping = use_head_mapping
        
        if use_head_mapping and student_heads != teacher_heads:
            # Learn which teacher heads to map to which student heads
            self.head_mapping = nn.Parameter(
                torch.randn(student_heads, teacher_heads)
            )
    
    def forward(self, student_attentions, teacher_attentions):
        """
        Args:
            student_attentions: [batch, num_heads, seq_len, seq_len]
            teacher_attentions: [batch, num_heads, seq_len, seq_len]
        """
        batch_size, _, seq_len, _ = student_attentions.shape
        
        if self.use_head_mapping and hasattr(self, 'head_mapping'):
            # Apply learned head mapping
            mapping_weights = F.softmax(self.head_mapping, dim=1)
            # [student_heads, teacher_heads] × [batch, teacher_heads, seq, seq]
            teacher_mapped = torch.einsum(
                'st,bthw->bshw', 
                mapping_weights, 
                teacher_attentions
            )
        else:
            # Simple averaging if heads match or no mapping desired
            if self.student_heads == self.teacher_heads:
                teacher_mapped = teacher_attentions
            else:
                # Average teacher heads to match student count
                teacher_mapped = teacher_attentions.reshape(
                    batch_size, self.student_heads, -1, seq_len, seq_len
                ).mean(dim=2)
        
        # MSE between attention distributions
        attention_loss = F.mse_loss(student_attentions, teacher_mapped.detach())
        
        return attention_loss

Progressive Distillation

Progressive distillation solves the problem of distilling very large teachers into very small students by using intermediate teachers. Instead of going straight from a GPT-3 scale model to a mobile-friendly size, you create a series of progressively smaller teachers, each distilled from the previous one. This staged approach lets each student learn from a teacher closer to its own capacity, which reduces the knowledge gap and improves final performance.

Curriculum learning in distillation means carefully ordering training examples. Early on, the student learns from easier examples where the teacher is more confident. As training goes on, you introduce more ambiguous cases where the teacher's soft labels provide maximum information. You can define this curriculum based on prediction entropy, loss magnitude, or example complexity.

class ProgressiveDistillationTrainer:
    def __init__(self, teachers, student, device='cuda'):
        """
        Args:
            teachers: List of teacher models ordered from largest to smallest
            student: Student model to train
        """
        self.teachers = teachers
        self.student = student
        self.device = device
        
        # Move all models to device and set teachers to eval
        for teacher in self.teachers:
            teacher.to(device)
            teacher.eval()
        self.student.to(device)
    
    def get_curriculum_weight(self, epoch, total_epochs):
        # Linearly increase difficulty over training
        return min(1.0, epoch / (total_epochs * 0.7))
    
    def compute_example_difficulty(self, teacher_logits):
        # Use entropy of teacher predictions as difficulty measure
        probs = F.softmax(teacher_logits, dim=-1)
        entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1)
        return entropy
    
    def progressive_distill(self, dataloader, stage, optimizer, 
                           temperature=3.0, epochs=10):
        """
        Distill from teachers[stage] into student or next teacher
        """
        current_teacher = self.teachers[stage]
        criterion = DistillationLoss(temperature=temperature)
        
        for epoch in range(epochs):
            curriculum_weight = self.get_curriculum_weight(epoch, epochs)
            
            for batch in dataloader:
                inputs, labels = batch
                inputs = inputs.to(self.device)
                labels = labels.to(self.device)
                
                # Get teacher predictions
                with torch.no_grad():
                    teacher_logits = current_teacher(inputs)
                    difficulties = self.compute_example_difficulty(teacher_logits)
                
                # Filter or weight examples based on curriculum
                difficulty_threshold = torch.quantile(
                    difficulties, curriculum_weight
                )
                example_weights = (difficulties <= difficulty_threshold).float()
                
                # Student forward pass
                student_logits = self.student(inputs)
                
                # Compute weighted distillation loss
                loss = criterion(student_logits, teacher_logits, labels)
                weighted_loss = (loss * example_weights.mean())
                
                # Optimization step
                optimizer.zero_grad()
                weighted_loss.backward()
                optimizer.step()
    
    def train_all_stages(self, dataloader, optimizer, epochs_per_stage=10):
        """
        Execute progressive distillation through all teacher stages
        """
        for stage in range(len(self.teachers)):
            print(f"Stage {stage}: Distilling from teacher {stage}")
            self.progressive_distill(
                dataloader, stage, optimizer, epochs=epochs_per_stage
            )
        
        return self.student

Task-Specific and Multi-Task Distillation

General distillation trains students to match teacher behavior across all tasks. Task-specific distillation optimizes for particular applications. This lets you compress more aggressively because the student only needs knowledge relevant to the target task. For example, distilling a general language model into a sentiment classifier can achieve much higher compression while maintaining or exceeding task performance.

Multi-task distillation extends this by training the student on multiple related tasks at once. The teacher might be an ensemble of task-specific expert models, and the student learns to handle all tasks in one architecture. This works well when tasks share underlying patterns, letting the student develop shared representations that generalize across tasks.

class MultiTaskDistillation(nn.Module):
    def __init__(self, task_weights=None):
        super().__init__()
        self.task_weights = task_weights or {}
    
    def forward(self, student_outputs, teacher_outputs, task_names, labels):
        """
        Args:
            student_outputs: Dict mapping task names to student logits
            teacher_outputs: Dict mapping task names to teacher logits
            task_names: List of tasks in current batch
            labels: Dict mapping task names to ground truth labels
        """
        total_loss = 0
        task_losses = {}
        
        for task in task_names:
            # Task-specific distillation loss
            criterion = DistillationLoss(
                temperature=self.get_task_temperature(task),
                alpha=self.get_task_alpha(task)
            )
            
            task_loss = criterion(
                student_outputs[task],
                teacher_outputs[task],
                labels[task]
            )
            
            # Weight by task importance
            weight = self.task_weights.get(task, 1.0)
            total_loss += weight * task_loss
            task_losses[task] = task_loss.item()
        
        return total_loss, task_losses
    
    def get_task_temperature(self, task):
        # Different tasks may benefit from different temperatures
        temperature_map = {
            'sentiment': 2.0,      # Lower for classification
            'nli': 3.0,            # Higher for complex reasoning
            'qa': 4.0,             # Highest for generation tasks
        }
        return temperature_map.get(task, 3.0)
    
    def get_task_alpha(self, task):
        # Balance between distillation and hard labels per task
        alpha_map = {
            'sentiment': 0.5,      # More weight on hard labels
            'nli': 0.7,            # Balanced
            'qa': 0.9,             # Heavy distillation weight
        }
        return alpha_map.get(task, 0.7)

Data Augmentation and Synthetic Data

How well distillation works depends heavily on the diversity and quality of training data. Basic distillation uses the same dataset that trained the teacher. Augmented distillation generates synthetic examples to expose the student to more teacher behaviors. The teacher generates labels for unlabeled data, greatly expanding the training set. This works especially well with task-specific augmentation strategies that target challenging cases or underrepresented patterns.

class DataAugmentedDistillation:
    def __init__(self, teacher, student, base_dataset):
        self.teacher = teacher
        self.student = student
        self.base_dataset = base_dataset
        
    def generate_synthetic_examples(self, num_examples, augmentation_fn):
        """
        Generate synthetic training examples using the teacher
        """
        synthetic_data = []
        self.teacher.eval()
        
        with torch.no_grad():
            for _ in range(num_examples):
                # Sample from base dataset and augment
                base_example = self.base_dataset[
                    torch.randint(len(self.base_dataset), (1,)).item()
                ]
                augmented_input = augmentation_fn(base_example)
                
                # Generate teacher predictions
                teacher_logits = self.teacher(augmented_input)
                
                synthetic_data.append({
                    'input': augmented_input,
                    'teacher_logits': teacher_logits.cpu(),
                    'source': 'synthetic'
                })
        
        return synthetic_data
    
    def hard_example_mining(self, dataloader, percentile=90):
        """
        Identify examples where student struggles most
        """
        self.student.eval()
        self.teacher.eval()
        
        example_difficulties = []
        
        with torch.no_grad():
            for batch in dataloader:
                inputs, labels = batch
                
                student_logits = self.student(inputs)
                teacher_logits = self.teacher(inputs)
                
                # Measure disagreement as difficulty proxy
                disagreement = F.kl_div(
                    F.log_softmax(student_logits, dim=-1),
                    F.softmax(teacher_logits, dim=-1),
                    reduction='none'
                ).sum(dim=-1)
                
                example_difficulties.extend(disagreement.cpu().numpy())
        
        # Return indices of hardest examples
        threshold = np.percentile(example_difficulties, percentile)
        hard_indices = np.where(
            np.array(example_difficulties) >= threshold
        )[0]
        
        return hard_indices

Training Stability and Optimization

The optimization landscape for distillation differs from standard supervised learning. The teacher's soft labels provide a smoother training signal than one-hot labels, which can speed up convergence but also cause instability if not managed carefully. Temperature directly affects gradient magnitudes, and improper tuning can lead to gradient explosion or vanishing gradients.

Learning rate scheduling is critical for success. A common strategy uses a warmup phase where the learning rate gradually increases, letting the student stabilize before full distillation. During main training, maintain a moderate learning rate with high temperature. Finally, a fine-tuning phase with reduced temperature and learning rate polishes performance.

class DistillationOptimizer:
    def __init__(self, student, initial_lr=1e-4, warmup_steps=1000):
        self.student = student
        self.initial_lr = initial_lr
        self.warmup_steps = warmup_steps
        self.global_step = 0
        
        # Use AdamW with weight decay for better generalization
        self.optimizer = torch.optim.AdamW(
            student.parameters(),
            lr=initial_lr,
            betas=(0.9, 0.999),
            weight_decay=0.01
        )
        
        self.scheduler = self._create_scheduler()
    
    def _create_scheduler(self):
        # Cosine schedule with warmup
        from torch.optim.lr_scheduler import LambdaLR
        
        def lr_lambda(step):
            if step < self.warmup_steps:
                # Linear warmup
                return step / self.warmup_steps
            else:
                # Cosine decay
                progress = (step - self.warmup_steps) / (10000 - self.warmup_steps)
                return 0.5 * (1 + np.cos(np.pi * progress))
        
        return LambdaLR(self.optimizer, lr_lambda)
    
    def step(self, loss):
        # Gradient clipping for stability
        torch.nn.utils.clip_grad_norm_(self.student.parameters(), max_norm=1.0)
        
        self.optimizer.step()
        self.scheduler.step()
        self.global_step += 1
        
        return self.scheduler.get_last_lr()[0]
    
    def get_temperature_schedule(self, max_steps):
        """
        Dynamic temperature scheduling during training
        """
        if self.global_step < self.warmup_steps:
            # Start with lower temperature during warmup
            return 2.0
        elif self.global_step < max_steps * 0.8:
            # Higher temperature for main distillation
            return 4.0
        else:
            # Reduce temperature for fine-tuning
            return 2.0

Evaluation Metrics

Evaluating distilled models requires more than just accuracy. You need to assess compression ratio, inference latency, memory footprint, and energy consumption. The distillation efficiency metric captures the trade-off between size reduction and performance retention, typically computed as the ratio of accuracy preservation to compression ratio.

class DistillationEvaluator:
    def __init__(self, teacher, student, test_loader, device='cuda'):
        self.teacher = teacher
        self.student = student
        self.test_loader = test_loader
        self.device = device
    
    def compute_compression_metrics(self):
        teacher_params = sum(p.numel() for p in self.teacher.parameters())
        student_params = sum(p.numel() for p in self.student.parameters())
        
        compression_ratio = teacher_params / student_params
        
        return {
            'teacher_parameters': teacher_params,
            'student_parameters': student_params,
            'compression_ratio': compression_ratio
        }
    
    def measure_inference_speed(self, num_samples=100):
        import time
        
        self.teacher.eval()
        self.student.eval()
        
        # Sample random inputs
        sample_inputs = []
        for batch in self.test_loader:
            sample_inputs.append(batch[0][:1].to(self.device))
            if len(sample_inputs) >= num_samples:
                break
        
        # Teacher inference time
        teacher_times = []
        with torch.no_grad():
            for inputs in sample_inputs:
                start = time.perf_counter()
                _ = self.teacher(inputs)
                teacher_times.append(time.perf_counter() - start)
        
        # Student inference time
        student_times = []
        with torch.no_grad():
            for inputs in sample_inputs:
                start = time.perf_counter()
                _ = self.student(inputs)
                student_times.append(time.perf_counter() - start)
        
        speedup = np.mean(teacher_times) / np.mean(student_times)
        
        return {
            'teacher_latency_ms': np.mean(teacher_times) * 1000,
            'student_latency_ms': np.mean(student_times) * 1000,
            'speedup_factor': speedup
        }
    
    def compute_agreement_metrics(self):
        """
        Measure how well student predictions agree with teacher
        """
        self.teacher.eval()
        self.student.eval()
        
        total_kl = 0
        total_top1_agreement = 0
        total_samples = 0
        
        with torch.no_grad():
            for inputs, labels in self.test_loader:
                inputs = inputs.to(self.device)
                
                teacher_logits = self.teacher(inputs)
                student_logits = self.student(inputs)
                
                # KL divergence
                kl = F.kl_div(
                    F.log_softmax(student_logits, dim=-1),
                    F.softmax(teacher_logits, dim=-1),
                    reduction='batchmean'
                )
                total_kl += kl.item() * inputs.size(0)
                
                # Top-1 agreement
                teacher_preds = teacher_logits.argmax(dim=-1)
                student_preds = student_logits.argmax(dim=-1)
                agreement = (teacher_preds == student_preds).float().mean()
                total_top1_agreement += agreement.item() * inputs.size(0)
                
                total_samples += inputs.size(0)
        
        return {
            'average_kl_divergence': total_kl / total_samples,
            'top1_agreement': total_top1_agreement / total_samples
        }
    
    def full_evaluation(self):
        """
        Comprehensive evaluation of distillation quality
        """
        metrics = {}
        
        # Compression metrics
        metrics.update(self.compute_compression_metrics())
        
        # Speed metrics
        metrics.update(self.measure_inference_speed())
        
        # Agreement metrics
        metrics.update(self.compute_agreement_metrics())
        
        # Efficiency score: accuracy preservation per unit compression
        metrics['efficiency_score'] = (
            metrics['top1_agreement'] * metrics['compression_ratio']
        )
        
        return metrics

Advanced Techniques

Recent advances have introduced several techniques that go beyond traditional knowledge transfer. Online distillation trains teacher and student simultaneously, with the teacher continuously updating rather than staying frozen. This co-evolution can lead to mutually beneficial learning where the student's progress informs teacher updates. Self-distillation applies distillation to the same architecture, using ensemble predictions or differently initialized models as teachers, which can improve performance even without compression.

Born-again networks are an extreme form of self-distillation where a student with the same architecture as the teacher often beats the teacher's performance. This suggests distillation provides more than just compression—it offers an improved optimization landscape and implicit regularization. Applying born-again distillation iteratively, where each generation serves as the teacher for the next, can progressively improve performance until it converges.

class OnlineDistillation(nn.Module):
    def __init__(self, teacher, student, teacher_update_freq=10):
        super().__init__()
        self.teacher = teacher
        self.student = student
        self.teacher_update_freq = teacher_update_freq
        self.step_count = 0
        
        # Initialize teacher with student parameters
        self.teacher.load_state_dict(student.state_dict())
        
        # Separate optimizers for teacher and student
        self.teacher_optimizer = torch.optim.AdamW(
            teacher.parameters(), lr=1e-5
        )
        self.student_optimizer = torch.optim.AdamW(
            student.parameters(), lr=1e-4
        )
    
    def train_step(self, inputs, labels, temperature=3.0):
        # Student learning from current teacher
        with torch.no_grad():
            teacher_logits = self.teacher(inputs)
        
        student_logits = self.student(inputs)
        
        criterion = DistillationLoss(temperature=temperature)
        student_loss = criterion(student_logits, teacher_logits, labels)
        
        self.student_optimizer.zero_grad()
        student_loss.backward()
        self.student_optimizer.step()
        
        # Periodically update teacher
        self.step_count += 1
        if self.step_count % self.teacher_update_freq == 0:
            # Teacher learns from student's predictions
            with torch.no_grad():
                student_logits_detached = self.student(inputs)
            
            teacher_logits = self.teacher(inputs)
            teacher_loss = criterion(teacher_logits, student_logits_detached, labels)
            
            self.teacher_optimizer.zero_grad()
            teacher_loss.backward()
            self.teacher_optimizer.step()
        
        return student_loss.item()

Summary

Language model distillation enables knowledge transfer from large, expensive teachers to efficient students. The techniques include output distribution matching, intermediate representation alignment, attention transfer, and progressive multi-stage distillation. Success depends on temperature scaling, curriculum learning, optimization dynamics, and architecture mapping.

The field continues to develop with innovations in online distillation, cross-modal knowledge transfer, and task-specific compression. As language models grow larger, distillation becomes necessary for making state-of-the-art natural language understanding accessible to more people. The core insight is straightforward: knowledge encoded in billions of parameters can be compressed into millions while preserving the essential patterns that drive intelligent behavior.