Distillation vs pruning: two ways to shrink a model

Distillation vs pruning represents the fundamental fork in model compression: one rewrites knowledge into a smaller architecture, the other surgically removes weights from the original. Both shrink models, but they optimize for different constraints and fail in different ways. Understanding which lever to pull — or whether to combine them — separates teams that ship efficient inference from teams that chase benchmarks.

What distillation actually does

Distillation trains a smaller student model to mimic a larger teacher’s output distribution. The student learns from soft targets — the teacher’s full probability vector over the vocabulary — rather than hard labels. This preserves the teacher’s “dark knowledge”: the relative probabilities among incorrect tokens that encode structural relationships in the data.

# Simplified distillation loss
def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Soft target loss (KL divergence between softened distributions)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    loss_kl = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (temperature ** 2)
    
    # Hard label loss (standard cross-entropy)
    loss_ce = F.cross_entropy(student_logits, labels)
    
    return alpha * loss_kl + (1 - alpha) * loss_ce

The student architecture is typically smaller — fewer layers, narrower hidden dimensions, fewer attention heads. You choose the architecture upfront, then train from scratch (or from a pretrained checkpoint) using the teacher’s outputs as supervision. The result is a standalone model that runs independently; the teacher is discarded after training.

Key implication: distillation changes the model architecture. The student has its own weight tensors, its own inference graph, and its own deployment footprint. You get a genuinely smaller model, not a sparse version of the original.

What pruning actually does

Pruning removes weights from an existing model based on importance criteria, then optionally fine-tunes to recover accuracy. It operates on a trained model post-hoc. The architecture stays the same on paper — same layer count, same hidden dimensions — but a fraction of weights become zero.

# Magnitude pruning (simplest criterion)
def magnitude_prune(model, sparsity=0.5):
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:  # Skip embeddings, biases, LayerNorm
            threshold = torch.quantile(param.abs().flatten(), sparsity)
            mask = (param.abs() > threshold).float()
            param.data.mul_(mask)
    return model

Structured pruning removes entire neurons, heads, or layers — yielding dense smaller tensors that standard kernels can execute efficiently. Unstructured pruning zeros individual weights, creating sparse matrices that require specialized kernels or hardware support to accelerate. Most production systems target structured pruning or semi-structured (2:4) patterns that NVIDIA Hopper and Ampere GPUs accelerate natively.

Key implication: pruning preserves the original architecture’s “shape.” You can often prune a model, fine-tune for a few epochs, and drop it into the same inference pipeline with minimal code changes — provided your runtime handles sparsity.

Comparison at a glance

Capabilities and quality retention

Distillation’s quality ceiling is higher. A well-distilled 7B model can match or exceed a pruned 13B model on downstream tasks because the student learns a more efficient representation from the teacher’s soft targets. The temperature parameter controls how much probability mass the teacher assigns to non-top tokens — higher temperature reveals more relational structure.

# Temperature effect on soft targets
# Teacher logits: [2.0, 1.5, 0.5, -0.2, ...]
# T=1.0:  [0.58, 0.24, 0.08, 0.04, ...]  -- peaked
# T=4.0:  [0.32, 0.26, 0.18, 0.14, ...]  -- flattened, more signal

Pruning hits a hard wall. Magnitude pruning typically maintains quality up to 30-50% sparsity (unstructured) or 20-30% (structured) before perplexity degrades non-linearly. Gradient-based criteria (movement pruning, SNIP, GraSP) push this further but add complexity. Beyond the wall, no amount of fine-tuning recovers the lost capacity — the model simply lacks the parameters to represent the function.

Distillation fails differently: if the student architecture is too small for the task complexity, it underfits regardless of training data. The failure mode is smooth degradation, not a cliff. You can always train a larger student.

Latency and throughput

Distillation delivers predictable latency gains. A 7B dense model runs at roughly half the latency of a 14B dense model on the same hardware — linear scaling with parameter count and FLOPs. Batch throughput scales similarly. No kernel surprises.

Pruning’s latency story depends entirely on sparsity pattern and runtime:

• Unstructured 50% sparsity: Zero speedup on standard dense kernels. Requires cuSPARSE, CUTLASS sparse GEMM, or custom kernels. Real-world speedup on H100 with 2:4 semi-structured: 1.3-1.6x for memory-bound layers, less for compute-bound.
• Structured pruning (remove 30% of heads/neurons): Dense tensors shrink. Standard kernels run faster. 1.3-1.5x speedup typical, but quality drops faster than unstructured at same parameter reduction.
• Block pruning (e.g., 32x32 blocks): Compromise between structure and granularity. Needs block-sparse kernels.

# Rough latency comparison on H100 (FP16, batch=1, seq=2048)
# Dense 7B:        ~45 ms/token
# Dense 13B:       ~85 ms/token
# Distilled 7B:    ~45 ms/token (same as dense 7B)
# Pruned 13B 50%:  ~70 ms/token (unstructured, no sparse kernel)
# Pruned 13B 50%:  ~55 ms/token (2:4 semi-structured on H100)
# Structured 30%:  ~60 ms/token (dense 9B equivalent)

If you need guaranteed latency reduction without kernel engineering, distillation wins. If you have sparse kernel infrastructure and need to compress an existing deployed model without retraining from scratch, pruning fits.

Training and deployment ergonomics

Distillation demands a full training pipeline: data curation, teacher inference (or cached logits), student training, evaluation, iteration. Compute cost is significant — comparable to pretraining the student from scratch, though shorter because the teacher provides dense supervision. You need GPU-hours, not GPU-minutes.

# Typical distillation compute (rough estimates)
# Teacher: 70B, Student: 7B, Data: 50B tokens
# Teacher inference (once): ~2,000 A100-hours
# Student training: ~15,000 A100-hours
# Total: ~17,000 A100-hours vs ~50,000+ for 7B from scratch

Pruning is a one-shot (or few-shot) operation on a checkpoint. Magnitude pruning takes minutes on CPU. Gradient-based criteria need a forward/backward pass on a calibration batch — hours on a single GPU. Fine-tuning after pruning: 1-10% of original training compute.

Deployment ergonomics flip the script. Distilled models are standard Hugging Face / vLLM / TensorRT-LLM checkpoints. AutoModelForCausalLM.from_pretrained() works. Pruned models need:

• Sparsity metadata (masks, block maps)
• Runtime that respects sparsity (TensorRT-LLM with sparse plugins, DeepSparse, custom kernels)
• Quantization compatibility checks (sparse + INT4/INT8 interactions are tricky)

Ecosystem and tooling

Distillation tooling centers on training frameworks:

• Hugging Face Transformers: Trainer with custom loss, DistillationTrainingArguments in TRL
• Megatron-LM / NeMo: Distributed distillation at scale
• DistilKit, MiniLLM: Specialized libraries with curriculum scheduling
• Logit caching: Write teacher outputs to disk once, stream during student training

Pruning tooling centers on post-training optimization:

• Torch Prune / TorchAO: Magnitude, structured, semi-structured; PyTorch-native
• NNI (Neural Network Intelligence): Microsoft’s auto-compression with search
• SparseML / DeepSparse: Neural Magic’s pipeline — prune, quantize, deploy on CPU/GPU
• TensorRT Model Optimizer: NVIDIA’s unified pruning + quantization for TRT-LLM
• LLM-Pruner, LoRA-Pruner: Structured pruning with low-rank adaptation recovery

The pruning ecosystem is more fragmented because sparsity patterns don’t standardize as cleanly as dense architectures. A pruned model from SparseML won’t run on TensorRT-LLM without conversion. Distilled models are just models.

Limits and failure modes

Distillation fails when:

• Teacher is weak on target domain: Garbage in, garbage out. Distilling a general-purpose 70B into a 7B for legal reasoning works poorly if the 70B doesn’t know law.
• Student capacity too low: A 1B student cannot absorb a 70B teacher’s knowledge, regardless of data. The bottleneck is student parameters, not training compute.
• Distribution shift: Teacher logits on out-of-distribution data mislead the student. Mix ground-truth labels (hard targets) via the alpha parameter.

Pruning fails when:

• Sparsity exceeds critical threshold: The “lottery ticket” hypothesis has limits. Past ~60% unstructured or ~30% structured, fine-tuning cannot rewire the remaining weights sufficiently.
• Importance metric misaligned: Magnitude pruning assumes small weights are unimportant. False for weights that are small but critically positioned (e.g., attention sink tokens). Gradient-aware criteria help but need calibration data representative of deployment distribution.
• Hardware doesn’t match pattern: 2:4 sparsity on A100 (no sparse tensor cores) is slower than dense. Unstructured sparsity on any GPU without sparse kernels is dead weight.

Which to choose

Choose distillation when:

• Building a new model family for a specific latency/cost budget (e.g., “we need a 3B model that runs at 20ms/token on T4”)
• You have compute budget for training and access to a strong teacher
• Target hardware lacks sparse acceleration (most CPUs, older GPUs, edge NPUs)
• You need the compressed model to run on arbitrary inference engines without custom kernels
• Quality at target size is non-negotiable — distillation finds better parameterizations than pruning can discover

Choose pruning when:

• Compressing an already-deployed model without retraining from scratch
• You have sparse kernel infrastructure (H100 2:4, CPU with DeepSparse, TensorRT-LLM sparse plugins)
• Target sparsity is moderate (<40% unstructured, <25% structured) and quality loss is acceptable
• Iteration speed matters — you need a smaller model this week, not next month
• The model is quantized (INT4/INT8) and you want to stack compression techniques

Combine them when:

• Distill a large teacher to a medium student (70B → 13B), then apply structured pruning to the student (13B → 9B dense equivalent)
• Use pruning to initialize a smaller architecture for distillation (prune 70B to 13B structure, then distill into it)
• The combination often beats either alone: distillation finds better dense representations; pruning removes remaining redundancy

# Combined pipeline sketch
def compress_pipeline(teacher, target_params, data):
    # Stage 1: Distill to intermediate size
    student_arch = design_architecture(target_params * 1.5)  # Overshoot
    student = distill(teacher, student_arch, data, epochs=3)
    
    # Stage 2: Structured prune to exact target
    student = structured_prune(student, target_sparsity=0.3)
    student = finetune(student, data, steps=1000)
    
    # Stage 3: Quantize
    student = quantize_int4(student, calibration_data)
    
    return student

The distillation vs pruning decision isn’t religious — it’s a constraint optimization problem. Map your constraints (latency budget, compute budget, hardware target, quality floor, timeline) to the comparison dimensions above. The right choice is usually obvious once constraints are explicit.