Model Pruning and Distillation

Quantization changes how weights are STORED. Pruning and distillation instead change how many weights (or how big a model) you have in the first place — two different routes to a smaller, cheaper model.

Pruning removes parameters that contribute little to the output. The intuition: trained networks are over-parameterized, like a hedge that's grown far bushier than it needs to be — you can trim a lot without changing its shape. There are two flavors. Unstructured pruning zeroes out individual weights scattered throughout the matrices (typically the ones with smallest magnitude). It can remove a large fraction of weights with little quality loss, BUT the result is a sparse matrix with holes everywhere, and standard GPUs don't run irregular sparsity any faster — so you save storage but often not wall-clock time (unless you have hardware with structured-sparsity support). Structured pruning removes whole coherent pieces — entire attention heads, neurons, or even layers. This yields a genuinely smaller, still-dense model that runs faster on ordinary hardware, at the cost of being harder to do without hurting quality. Pruning is usually followed by a short fine-tune to let the remaining weights recover.

Distillation takes a different tack: rather than trimming one model, you train a brand-new SMALL model (the 'student') to reproduce the behavior of a large, capable model (the 'teacher'). Crucially, the student learns not just from the correct answer but from the teacher's full output distribution — the 'soft labels' or logits that reveal HOW confident the teacher was across all options. Those soft targets carry far more information than a single hard label ('the teacher thought this was 70% cat, 25% dog, 5% fox' teaches more than just 'cat'). This is how models like DistilBERT reached ~97% of BERT's performance at ~40% of the size, and how many small open models are trained on outputs from larger frontier models.

When to use which: reach for structured pruning when you have a good model and want a faster same-family version with modest effort; reach for distillation when you want the biggest size reduction and can afford a real training run, or when you want to compress a proprietary teacher's behavior into a model you control. In practice these compose — you might distill, then quantize the student, then serve it with continuous batching.

# (1) Magnitude pruning: zero out the smallest-magnitude weights.
# (2) Distillation loss: student mimics the teacher's SOFT targets.
import torch
import torch.nn.functional as F

def magnitude_prune(weight: torch.Tensor, sparsity: float) -> torch.Tensor:
    """Unstructured pruning: keep the largest |weights|, zero the rest."""
    k = int(weight.numel() * sparsity)
    if k == 0:
        return weight
    threshold = weight.abs().flatten().kthvalue(k).values
    return torch.where(weight.abs() > threshold, weight, torch.zeros_like(weight))

w = torch.randn(4, 4)
print("pruned 50% of weights:\n", magnitude_prune(w, sparsity=0.5))

def distillation_loss(student_logits, teacher_logits, temperature=2.0):
    """KL divergence between softened student and teacher distributions —
    the student learns the teacher's full 'soft label' distribution."""
    t = temperature
    soft_teacher = F.softmax(teacher_logits / t, dim=-1)
    soft_student = F.log_softmax(student_logits / t, dim=-1)
    return F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (t * t)

teacher = torch.tensor([[2.0, 0.5, 0.1]])  # confident-but-informative
student = torch.tensor([[1.0, 0.8, 0.3]])
print("distillation loss:", distillation_loss(student, teacher).item())