DL S068 · Pruning and Distillation — When Smaller Beats Quantized

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 68 / 80 · Module M11 · ~2 hours

The story

Quantization (S067) shrinks by making each weight cheaper. Pruning shrinks by making some weights zero. Distillation shrinks by training a smaller model to imitate a bigger one. All three are compression, all three have decades of history, and all three interact badly with LLMs in ways the CNN literature didn't prepare us for.

Here's the awkward truth: for LLMs, pruning is disappointing. The Lottery Ticket Hypothesis papers dazzled everyone in 2019 with 90% sparsity on ResNet. On GPT-scale models the same techniques get you 20% sparsity before quality collapses. Meanwhile distillation is phenomenal: Llama-3.2-1B (distilled from 8B/70B) is basically usable, DistilBERT retains 97% of BERT quality at 40% size, and Alpaca kicked off the entire "instruction-tuned tiny model" era.

So today: pruning honestly (with the caveats), the one form of pruning that does work on GPUs (2:4 sparsity on Ampere+), and distillation deeply — because distillation is where the real wins live.

Prerequisites

• Session 067 — compression as an axis; we're adding two more axes.
• Session 058 — you'll need to run SFT on the student.
• Session 022 — L1/L2 penalties feed straight into pruning intuition.

1 · Magnitude pruning

The oldest, dumbest, still-competitive method: rank weights by absolute value, set the smallest $k\%$ to zero, fine-tune the rest.

import torch, torch.nn as nn
 
def magnitude_prune(model: nn.Module, sparsity: float):
    """Global magnitude pruning across all Linear layers."""
    all_weights = torch.cat([
        p.detach().abs().flatten()
        for name, p in model.named_parameters()
        if "weight" in name and p.dim() == 2
    ])
    threshold = torch.quantile(all_weights, sparsity)
    masks = {}
    for name, p in model.named_parameters():
        if "weight" in name and p.dim() == 2:
            mask = p.detach().abs() >= threshold
            p.data.mul_(mask)
            masks[name] = mask
    return masks

Then during fine-tuning, re-apply masks after each optimizer step so pruned weights stay zero:

def apply_masks(model, masks):
    for name, p in model.named_parameters():
        if name in masks:
            p.data.mul_(masks[name])

This is called iterative magnitude pruning (IMP): prune 20%, fine-tune, prune another 20%, fine-tune. Better than one-shot.

The disappointing curve

On ResNet-50/ImageNet: 90% sparsity, <1% accuracy loss.
On Llama-2-7B/WikiText: 50% sparsity, ~10× perplexity increase. Even 20% hurts.

Why the gap? LLMs are trained closer to their capacity — every weight matters more. And unstructured sparsity doesn't speed up inference on a GPU (dense matmul kernels don't skip zeros; you're just wasting the same FLOPs).

2 · N:M structured sparsity — the pruning GPUs actually like

Ampere (A100) and newer support 2:4 sparsity: in every group of 4 consecutive weights along a specific axis, at most 2 are non-zero. NVIDIA's Sparse Tensor Cores literally skip the zeros and give you a 2× throughput uplift on the matmul.

def enforce_2_4(w: torch.Tensor) -> torch.Tensor:
    """w: (out, in). Keep top-2 magnitudes in every group of 4 along `in`."""
    out_f, in_f = w.shape
    assert in_f % 4 == 0
    w4 = w.view(out_f, in_f // 4, 4)
    _, topk_idx = w4.abs().topk(k=2, dim=-1)  # (out, in/4, 2)
    mask = torch.zeros_like(w4).scatter_(-1, topk_idx, 1.0).bool()
    return (w4 * mask).view(out_f, in_f)

Enable it at inference with torch.sparse.SparseSemiStructuredTensor or via TensorRT-LLM's --use_sparse_gemm. On Llama-2-7B with SparseGPT + 2:4, real papers report ~5% perplexity increase for ~1.8× decode speedup. Not free, but real.

Movement pruning (Sanh et al.) is a smarter variant: rank not by magnitude but by how much a weight moved during fine-tuning. Weights growing are important; weights drifting to zero can go. Better than magnitude on downstream tasks with small fine-tuning budgets.

3 · SparseGPT — the GPTQ of pruning

Frantar & Alistarh, 2023: same Hessian-compensation trick as GPTQ but for pruning masks. Given a target sparsity per layer, solve for the optimal mask + weight update jointly. One-shot (no fine-tuning), works at 50% for OPT-175B with <5% perplexity drop.

# pip install sparsegpt (Frantar's ref implementation)
from sparsegpt import SparseGPT
# hook each Linear layer, feed calibration data:
#   sparsegpt.fasterprune(sparsity=0.5, prunen=2, prunem=4)

prunen=2, prunem=4 gets you N:M sparsity in one shot without fine-tuning. This is what you'd actually run if you were serious about pruning today.

4 · Knowledge distillation — the wins are here

Hinton et al. (2015): a student model is trained not on the hard one-hot labels but on the soft probabilities from a teacher.

$\mathcal{L} = (1-\alpha) \cdot \mathcal{L}_{CE}(y_{\text{true}}, p_s) + \alpha \cdot T^2 \cdot \mathcal{L}_{KL}(p_t^T \,\|\, p_s^T)$

Where $p^T = \text{softmax}(z/T)$ is the softened distribution at temperature $T$. Typical $T \in [2, 5]$, $\alpha \in [0.5, 0.9]$.

Why it works: soft labels leak the teacher's uncertainty structure — "this looks 60% cat, 30% dog, 10% fox" is more informative than "cat". The student learns not just the answer but the manifold of related answers.

Response distillation (the basic recipe)

import torch
import torch.nn.functional as F
 
def distill_step(student, teacher, x, y_true, T=2.0, alpha=0.7):
    with torch.no_grad():
        z_t = teacher(x).logits           # (B, seq, V)
    z_s = student(x).logits
    # Hard-label loss
    ce = F.cross_entropy(z_s.view(-1, z_s.size(-1)), y_true.view(-1))
    # Soft-label KL
    log_p_s = F.log_softmax(z_s / T, dim=-1)
    p_t = F.softmax(z_t / T, dim=-1)
    kl = F.kl_div(log_p_s, p_t, reduction="batchmean") * (T * T)
    return (1 - alpha) * ce + alpha * kl

That's ~90% of what DistilBERT / TinyLlama / MiniLM did.

Feature distillation

Match hidden states, not just logits: $\|h_s^{\ell} - W \cdot h_t^{\ell'}\|^2$ for layer pairs $(\ell, \ell')$. Useful when the student is much shallower than the teacher — logits alone give too little signal.

Relation distillation (RKD, TinyBERT)

Match pairwise relationships between samples: $\|d(h_s^i, h_s^j) - d(h_t^i, h_t^j)\|$. Preserves the geometry of the teacher's representation space.

Modern LLM distillation

The dominant recipe now is even simpler: sample from the teacher, SFT the student. Alpaca (Stanford, 2023) took GPT-3.5 outputs on 52k prompts and fine-tuned Llama-7B. The student never sees teacher probabilities, just teacher completions. It works because the completions themselves carry the teacher's implicit reasoning.

5 · When does distillation beat quantization?

The killer combo: distill 70B → 7B, then AWQ-4bit that 7B. Effective ~35× compression with maybe 3 points of MMLU loss. This is roughly what Meta did for on-device Llama-3.2-1B.

6 · A distillation run in ~40 lines

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch, torch.nn.functional as F
 
teacher = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B-Instruct", torch_dtype=torch.bfloat16, device_map="cuda:0").eval()
student = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16, device_map="cuda:1")
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B-Instruct")
 
opt = torch.optim.AdamW(student.parameters(), lr=2e-5)
 
def loss_fn(z_s, z_t, y, T=2.0, alpha=0.7):
    ce = F.cross_entropy(z_s.view(-1, z_s.size(-1)), y.view(-1), ignore_index=-100)
    kl = F.kl_div(
        F.log_softmax(z_s / T, dim=-1),
        F.softmax(z_t / T, dim=-1),
        reduction="batchmean",
    ) * (T * T)
    return (1 - alpha) * ce + alpha * kl
 
for batch in loader:
    input_ids = batch["input_ids"].cuda(0)
    labels = batch["labels"].cuda(1)
    with torch.no_grad():
        z_t = teacher(input_ids).logits.to("cuda:1")
    z_s = student(input_ids.to("cuda:1")).logits
    loss = loss_fn(z_s, z_t, labels)
    loss.backward()
    opt.step(); opt.zero_grad()

Two things to notice:

1. Teacher on GPU 0, student on GPU 1. You need both in memory — the teacher forward dwarfs the student's memory anyway.
2. ignore_index=-100 — standard HF convention for masking prompt tokens from the loss. Only distill on the response tokens.

7 · Pitfalls

9 · The 2024–2025 distillation renaissance

Distillation went from a niche technique to the default way small models get made. Every open-weights release under 3B in 2024–2025 was distilled from a bigger sibling. Let me lay out the taxonomy that emerged.

9.1 White-box vs black-box distillation

• White-box (classical): you have the teacher's logits per token. KL loss against them. This is the code above; requires access to teacher weights.
• Black-box (a.k.a. "distillation from API"): you only have generated text from the teacher — no logits. Train the student on (prompt, teacher_response) pairs as SFT. Alpaca, Vicuna, WizardLM all did this from GPT-3.5/GPT-4. Legal grey area (violates most APIs' terms) but wildly effective.

Black-box distillation seeded the entire open-source instruction-tuned model ecosystem in 2023. In 2024–2025 it's less common because Llama-3/Qwen2.5 have good enough instruction data of their own, but any time you see a "trained on synthetic data" tag on HuggingFace, that's black-box distillation.

9.2 On-policy vs off-policy distillation

Agarwal et al., "On-policy distillation" (Google 2024) argue: don't just distill on a fixed dataset. Instead, at each step, let the student generate a response, then use the teacher to score/relabel it. This closes the train-test distribution gap (student never sees its own outputs during training). Gemma-2 was distilled this way; Google claim ~2× sample efficiency.

9.3 MiniLLM and reverse-KL

Gu et al., MiniLLM (ICLR 2024) noticed that standard forward-KL (student matches teacher's distribution) forces the student to spread probability across all modes the teacher covers — including modes the small student can't do well. Switching to reverse-KL (teacher matches student's distribution) makes the student concentrate on modes it can fit. On smaller students (≤ 1B) this is a 3–5 point win on MT-Bench.

9.4 Distilling reasoning: the DeepSeek-R1 story

January 2025. DeepSeek releases DeepSeek-R1, a reasoning model trained with pure RL. Then they do something remarkable: they distill R1's reasoning traces into Qwen and Llama students of various sizes and release them.

DeepSeek-R1-Distill-Qwen-7B scores 55.5% on AIME 2024 vs Qwen-2.5-7B's ~12%. The student inherits chain-of-thought reasoning from the teacher, essentially for free, at 7B scale. This kicked off the entire "reasoning distillation" wave of 2025 — within 4 months there were 30+ open reasoning models all trained this way.

The lesson: distillation transfers not just knowledge but capabilities, if you can generate teacher outputs that exercise the capability. Alpaca did this for instruction following in 2023; R1 did it for reasoning in 2025.

Further reading — the modern distillation stack:

• Agarwal et al., Generalized Knowledge Distillation (2024).
• Gu et al., MiniLLM (ICLR 2024).
• DeepSeek-AI, DeepSeek-R1 (2025) — §2.4 on distillation to small dense models.
• Meta, Llama 3.2 blog — layer-drop + distill recipe.
• Xia et al., Sheared LLaMA (ICLR 2024) — structured pruning + continued pretraining as a distillation alternative.

10 · Try it yourself

1. Prune → measure. Take a small model (gpt2 or Qwen2.5-0.5B), magnitude-prune to 30%/50%/70% sparsity without fine-tuning, measure WikiText perplexity at each. Plot the curve. Then re-run with a short healing fine-tune. Notice how much fine-tuning recovers.
2. 2:4 sparsify one Linear layer with torch.sparse.to_sparse_semi_structured and time the matmul on an A100 vs a dense reference. Confirm the ~1.8× speedup.
3. Reasoning distillation lite. Generate 200 chain-of-thought traces from DeepSeek-R1-Distill-Qwen-7B on GSM8K problems. SFT a Qwen-0.5B student on those traces. Measure GSM8K accuracy delta vs the un-distilled baseline. Even 200 examples move the needle noticeably.

8 · Mermaid: the compression landscape

Recall

Stretch: run a distillation of a 3B student from an 8B teacher on Alpaca-style data, compare against SFT-only baseline. Report MMLU delta. Predict what happens at temperature 1 vs 3.

In your own words: the reason pruning underperforms on LLMs vs CNNs is __________________________.

Spaced review: S022 regularization, S058 SFT, S067 quantization.

Next session (S069): speculative decoding — a tiny draft model proposes tokens the big model verifies in parallel. 2–3× decode speedup for zero quality loss.

Bring back tomorrow:

• The N:M sparsity idea (only pruning that maps to hardware).
• The soft-label + hard-label distillation loss.
• The distill-then-quantize combo.

import torch, time
W = torch.randn(4096, 4096, device='cuda', dtype=torch.float16)
x = torch.randn(4096, 4096, device='cuda', dtype=torch.float16)
mask = (W.abs() > W.abs().flatten().kthvalue(W.numel()//2).values)
Ws = W * mask                      # 50% zeros, same shape, same dtype

def bench(A):
    for _ in range(5): A @ x
    torch.cuda.synchronize(); t = time.time()
    for _ in range(50): A @ x
    torch.cuda.synchronize(); return (time.time()-t)/50

print('dense ', bench(W))
print('sparse', bench(Ws), '<- same bytes, same shape')
print('narrow', bench(W[:2048]), '<- structurally smaller')

Previous: ← DL S067 · Next: DL S069 →