Remove unnecessary weights and transfer knowledge to smaller models.

Model Pruning

Unstructured Pruning

import torch.nn.utils.prune as prune

# Remove 30% of smallest weights
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.l1_unstructured(module, name='weight', amount=0.3)

# Make pruning permanent
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.remove(module, 'weight')

Structured Pruning (Hardware-Friendly)

# Remove entire attention heads
def prune_attention_heads(model, heads_to_prune):
    for layer_idx, head_indices in heads_to_prune.items():
        layer = model.layers[layer_idx]
        layer.self_attn.prune_heads(head_indices)

Knowledge Distillation

Teacher-Student Training

class DistillationTrainer:
    def __init__(self, teacher, student, temperature=2.0, alpha=0.5):
        self.teacher = teacher
        self.student = student
        self.temperature = temperature
        self.alpha = alpha
    
    def distillation_loss(self, student_logits, teacher_logits, labels):
        # Soft targets from teacher
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.temperature, dim=-1),
            F.softmax(teacher_logits / self.temperature, dim=-1),
            reduction='batchmean'
        ) * (self.temperature ** 2)
        
        # Hard targets from labels
        hard_loss = F.cross_entropy(student_logits, labels)
        
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

Popular Distilled Models