Pruning I: Making Networks Sparse

Chapter 1: The Pruning Problem (Formal)

Let's be precise about what we're actually trying to do. You have a trained network with weights W and an objective (loss) function L(x; W). You want to find a pruned set of weights W_P that minimizes the same objective while having few non-zero entries:

Here ‖W_P‖₀ counts the number of non-zero elements in W_P — the L0 "norm" (it's not actually a norm, but it counts non-zeros). N is your target sparsity budget: how many connections the pruned network is allowed to have.

This formulation is clean but the solution is not. Exactly minimizing this is an NP-hard combinatorial search — you'd have to evaluate every possible subset of weights of size N and pick the best one. For a 61M-parameter network, that's more subsets than atoms in the observable universe.

The practical approach is a three-step decomposition:

1. Score every weight (or group of weights) with an importance metric — a proxy for how much the loss would increase if we removed it.
2. Rank and threshold — remove the k% least important weights.
3. Finetune — run a few epochs of gradient descent with the surviving weights to recover accuracy.

Two design choices dominate the rest of this lesson: what you prune (granularity — individual weights? entire channels?) and why you prune it (criterion — magnitude? curvature? activation statistics?). These are independent decisions: you can use magnitude as a criterion for either fine-grained or channel pruning.

Worked numbers — the compression table from Han et al. 2015:

Note: compression ratio > MAC reduction for unstructured pruning because zero weights still occupy compute cycles on dense GPU kernels. Real speedup requires sparse hardware support or structured pruning.

Chapter 2: Granularity I — Fine-Grained & N:M Sparsity

Before you decide which weights to prune, you must decide at what scale to prune. This is granularity: how big is the unit you remove? The spectrum runs from individual scalar weights at one end, to entire filters (all weights in an output channel) at the other.

Fine-grained pruning (also called unstructured pruning) removes individual weights. For a 2D weight matrix, you get a pattern that looks random — a few surviving non-zeros scattered irregularly across the tensor. This is the most flexible option: you can remove exactly the weights with the smallest importance score, regardless of where they live. The result is a sparse matrix.

The problem: regular hardware runs on dense matrix operations (GEMM). A sparse weight tensor still occupies the same memory slots as its dense counterpart — you'd need to explicitly zero out entries and skip them in the multiply-accumulate pipeline. On a standard GPU, zeroing out 90% of a matrix doesn't make inference 10× faster; it might be the same speed or slower, because the hardware still executes dense vector operations and just multiplies by 0. To actually accelerate fine-grained sparse inference, you need:

• Sparse tensor cores (NVIDIA A100 supports 2:4 sparsity with ~2× actual throughput)
• Custom ASICs (EIE by Han et al. — pointer-jumping CSR format)
• Encoding overhead: you must store the non-zero values plus their indices

N:M sparsity is a structured-irregular middle ground. For every contiguous M elements in the weight tensor, exactly N are kept and M−N are zeroed. The classic case is 2:4 sparsity: in every group of 4 weights, exactly 2 are non-zero (50% sparsity). NVIDIA's Ampere architecture (A100) supports 2:4 sparsity natively in its Sparse Tensor Cores: weights are stored in a compressed format (non-zero values + 2-bit indices), and the hardware skips zero multiplications automatically, delivering up to 2× throughput. Accuracy tests across BERT, ResNet, and ViT show that 2:4 sparsity typically recovers to within 0.5% of the dense baseline after sparse-aware training.

Why 2:4 was chosen: The 2-bit index per non-zero (4 possible positions in a group of 4) is the minimum overhead that delivers near-unstructured flexibility. At 2:4 each compressed block stores 2 values + 2 × 2 bits = 2 values + 4 bits overhead, versus the 4 values in the dense form. Memory footprint is exactly halved. The hardware can then schedule sparse-dense matmuls with a dedicated "Sparse GEMM" kernel that reads the compressed matrix and uses the indices to select matching dense columns — achieving 2× throughput with no accuracy-to-sparsity tradeoff beyond what fine-grained 50% pruning would give.

Chapter 3: Granularity II — Structured Pruning

Structured pruning removes entire groups of weights that correspond to a unit the hardware naturally processes: a row, a column, a kernel (one 3×3 slice), or a channel (all kernels feeding one output). The result is a smaller dense network — no sparse formats needed, no special hardware, just the same GPU GEMM running on a tensor with fewer dimensions.

For a convolutional layer, the weight tensor has shape (C_out, C_in, k_h, k_w). The granularity hierarchy:

Channel pruning is the most popular structured method. If a layer has C_out = 512 output channels and you prune 50% of them, you get C_out = 256. The next layer's input channels must also shrink from 512 to 256 — you remove the corresponding input-channel slices in that layer too. The network literally becomes smaller: every tensor in the pruned region is narrower. No zero-skipping needed; plain dense convolution runs faster because the tensor is smaller.

Worked numbers — channel pruning a conv layer:

Layer spec: C_in=256, C_out=512, k=3×3, output 14×14.

After 50% channel pruning (C_out → 256):

For channel pruning at sparsity s: params and MACs both scale as (1−s), because the output tensor shrinks by (1−s) in the channel dimension and the next layer also shrinks its input side. This is the key advantage: real latency reduction without sparse-execution hardware.

Chapter 4: Criteria I — Magnitude-Based Pruning

You've decided on a granularity. Now you need to score every weight (or group of weights) with an importance value. The simplest and most widely-used criterion is magnitude: larger absolute value = more important. The intuition is direct — a weight of 0.001 contributes almost nothing to the output, while a weight of 5.0 strongly amplifies or suppresses its input. Remove the small ones.

For element-wise (fine-grained) pruning, the importance of a single weight w_i is:

You sort all weights by |w|, and remove the bottom-k%. This is the approach in Han et al. 2015 — what they called "learning connections" (find a threshold, zero out everything below it, retrain with a frozen mask).

For row-wise (structured) pruning, you need a single scalar importance for a whole group of weights. Two natural choices:

Worked example with numbers: Take a 2×4 weight matrix:

# Weight matrix (2 rows, 4 cols)
W = [[3, -2, 0, 1],
     [-5, 0, 1, -0.2]]

# L1-norm row importance
row0_L1 = |3| + |-2| + |0| + |1| = 6
row1_L1 = |-5| + |0| + |1| + |-0.2| = 6.2

# L2-norm row importance
row0_L2 = sqrt(9 + 4 + 0 + 1) = sqrt(14) ≈ 3.74
row1_L2 = sqrt(25 + 0 + 1 + 0.04) = sqrt(26.04) ≈ 5.10

# Both norms agree: row1 is more important than row0.
# At 50% row sparsity, row0 would be pruned.

Per-layer magnitude pruning in PyTorch:

import torch

def magnitude_prune(weight: torch.Tensor, sparsity: float) -> torch.Tensor:
    """Return a binary mask: 1=keep, 0=prune. Keeps top (1-sparsity) fraction."""
    flat = weight.abs().view(-1)
    k = int(sparsity * flat.numel())  # number to prune
    threshold, _ = torch.kthvalue(flat, k)  # k-th smallest
    mask = (weight.abs() > threshold).float()
    return mask

# Example: 4x4 weight matrix, 50% sparsity
W = torch.randn(4, 4)
mask = magnitude_prune(W, sparsity=0.5)
W_pruned = W * mask  # zero out pruned weights
print(f"Nonzero weights: {mask.sum().int()}/16")  # → 8

Chapter 5: Criteria II — Scaling Factors & APoZ

Magnitude-of-weights is a natural criterion when weights are all at the same scale. But two other signals carry more direct information about channel importance: the BatchNorm scaling factor (γ) and the average percentage of zero activations (APoZ). Each exploits information that is simply not available in the weight magnitude alone.

Scaling-based pruning — BN γ as channel importance

In a convolutional network with BatchNorm, each output channel j has a learned scaling factor γ_j. The BatchNorm output is:

If γ_j is very small, channel j barely contributes to the next layer's input — regardless of what the convolutional weights themselves look like. A γ near zero says: "this channel's output is being globally suppressed by training." That's a much more reliable importance signal than raw weight magnitude, because it reflects the network's learned decision about which channels matter.

The Network Slimming technique (Liu et al. ICCV 2017) adds an L1 sparsity penalty on all γ values during training:

This regularization nudges unimportant channel scaling factors toward zero. After training, pruning is trivial: sort channels by |γ_j|, remove the bottom-k%, and retrain.

APoZ — percentage of zero activations

A different approach looks not at weights but at activations. ReLU networks produce many zeros in activation maps — any negative pre-activation becomes exactly zero. If a neuron (or channel) produces zero for the vast majority of inputs, it is contributing almost nothing to downstream computation. APoZ — Average Percentage of Zero activations — quantifies this.

where ϕ is an indicator that equals 1 when the activation is exactly zero, N is the number of evaluation samples, and H×W is the spatial dimension of the feature map.

APoZ is computed on a calibration dataset (typically a few thousand training examples) by running a forward pass and recording which channels are zero. A channel with APoZ = 95% is dead for 95% of inputs — an obvious pruning candidate.

import torch

def compute_apoz(model, dataloader, layer_name):
    """Compute Average Percentage of Zeros for each channel in a layer."""
    zero_counts = None
    total = 0
    def hook(module, inp, output):
        nonlocal zero_counts, total
        # output: (N, C, H, W)
        zeros = (output == 0).float().sum(dim=[0, 2, 3])  # sum over N,H,W → (C,)
        count = output.shape[0] * output.shape[2] * output.shape[3]
        if zero_counts is None:
            zero_counts = zeros
        else:
            zero_counts += zeros
        total += count
    h = model.get_submodule(layer_name).register_forward_hook(hook)
    model.eval()
    with torch.no_grad():
        for x, _ in dataloader:
            model(x)
    h.remove()
    return zero_counts / total  # APoZ per channel

Chapter 6: Criteria III — Second-Order Pruning (OBD)

Magnitude tells you how big a weight is. But what you really want to know is: how much does the loss change if I remove this weight? For a well-trained network, removing weight w_i means setting δw_i = w_i (deleting it is equivalent to perturbing it by its current value). The change in loss is δL.

Step 1: Taylor expansion of δL. Expand L(x; W − δW) around the current weights W:

where g_i = ∂L/∂w_i (gradient) and h_ij = ∂²L/(∂w_i∂w_j) (Hessian entry).

Step 2: Three OBD assumptions.

1. Quadratic objective: Drop the cubic and higher terms → δL = ∑ g_iδw_i + ½∑ h_iiδw_i² + cross terms.
2. Training has converged: At a local minimum, gradients are approximately zero (g_i ≈ 0 for all i). Drop the first-order term → δL = ½∑ h_iiδw_i² + cross terms.
3. Independent errors: Assume pruning one weight doesn't affect the contribution of others (cross terms ≈ 0). Drop off-diagonal Hessian → δL = ½∑ h_iiδw_i².

Step 3: The OBD saliency. Now set δw_i = w_i (pruning weight i means removing it, i.e., the perturbation equals the weight value) and δw_j = 0 for all j ≠ i:

This is the OBD saliency — the estimated loss increase from removing weight i. Weights with small saliency are safe to prune. The key difference from magnitude: saliency weighs the weight by its curvature h_ii. A small weight in a high-curvature region (h_ii large) can be crucial; a large weight in a flat region (h_ii small) can be safely removed.

Worked numerical example: Two weights, w₁ = 0.1 and w₂ = 3.0.

OBD decision: Prune w₂ (saliency 0.045 < 1.0 — costs only 0.045 loss units). Magnitude decision: Prune w₁ (|w₁| = 0.1 < |w₂| = 3.0). OBD avoids the expensive mistake: w₁ sits on a steep ridge of the loss surface (h₁₁=200 means moving w₁ by 0.1 costs 1.0 loss units), while w₂=3.0 sits in a flat valley (h₂₂=0.01 means removing it barely changes the loss).

Taylor expansion pruning (1st order): A simpler compromise. Keep the gradient term and approximate saliency as |g_i × w_i|. This is equivalent to the first-order Taylor series and is much cheaper than computing second derivatives. It works well for small pruning steps (when the network isn't far from the original optimum). Used in Molchanov et al. CVPR 2019.

Chapter 7: The Prune–Finetune Loop

You've picked a granularity and a criterion. Now comes the most practical question: do you prune the entire network at once (one-shot), or do you alternate between pruning and finetuning in small increments (iterative)?

The intuition for why iterative wins: imagine you're hiking and you want to remove 90% of your gear. If you dump 90% immediately, you've made a huge random change to how you carry load — you might drop critical items and the remaining 10% might not form a coherent system. But if you remove 10% at a time, evaluating and re-packing after each step, you can carefully select what to discard and the remaining gear stays coherent.

In network terms: one-shot pruning at 90% sparsity simultaneously removes many interdependent weights. The loss surface shifts sharply; the surviving weights were optimized for a different parameter configuration and finetuning may not recover accuracy. Iterative pruning at 10% per step makes a small change, finetunes to re-adapt the surviving weights, then repeats. Each finetuning step keeps the network near a good minimum before the next prune.

The standard loop:

# Iterative magnitude pruning with PyTorch
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

model = YourModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
target_sparsity = 0.90
steps = 9            # prune 10% → 20% → ... → 90%
step_size = target_sparsity / steps  # 0.10 per step

for step in range(steps):
    current_sparsity = (step + 1) * step_size
    # --- Prune each linear/conv layer ---
    for name, module in model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            prune.l1_unstructured(module, name='weight', amount=step_size)
    # --- Finetune for a few epochs ---
    for epoch in range(finetune_epochs):
        train_one_epoch(model, optimizer, dataloader)
    print(f"Step {step+1}: sparsity={current_sparsity:.0%}")

# Remove the pruning hooks and make sparsity permanent
for name, module in model.named_modules():
    if hasattr(module, 'weight_mask'):
        prune.remove(module, 'weight')

Per-layer sensitivity

Not all layers are created equal. The first convolutional layer in a CNN processes raw pixels — its small number of learned edge detectors is hard to replace. If you prune 80% of them, the whole network sees degraded features. The last fully-connected layer, by contrast, often has massive redundancy and can tolerate 80–90% sparsity with negligible accuracy loss.

Sensitivity analysis: for each layer independently, measure the accuracy drop as a function of sparsity. Plot it as a curve (accuracy vs sparsity per layer). This tells you the maximum safe sparsity for each layer and motivates non-uniform sparsity: sensitive layers get low sparsity (10–30%); robust layers get high sparsity (70–90%). The overall compression is a weighted average.

Chapter 8: Showcase — Sparsity Explorer

This showcase puts together everything from this lesson: granularity, criterion, and the per-layer sensitivity landscape. You're the ML engineer. You have a small CNN with five layers. Each layer has a different sensitivity to pruning — the first conv is fragile, the FC layers are robust. Your tools are: choose the granularity (fine or channel), choose the criterion (magnitude or OBD proxy), choose per-layer sparsity, and see how accuracy and MAC savings evolve.

Worked MAC savings at different sparsity targets:

Channel pruning gives exact proportional scaling: 50% channel sparsity ⇒ 50% param reduction ⇒ 50% MAC reduction ⇒ ~50% wall-clock speedup (varies with memory bandwidth and other overheads).

Chapter 9: Connections & Cheat Sheet

You now have the core vocabulary of pruning. Here's the full picture in one table, followed by what comes next.

Granularity × Criteria Cheat Sheet

What this lesson covered

• The pruning formulation: argmin L(x; W_P) s.t. ‖W_P‖₀ ≤ N — NP-hard, solved by heuristics.
• Granularity spectrum: fine-grained (max compression, needs sparse HW) → channel (less compression, always accelerates).
• N:M 2:4 sparsity: NVIDIA Ampere native support, 2× throughput at 50% sparsity.
• Magnitude criterion: sort by |W|, threshold, persistent mask during finetune.
• BN γ criterion: training-time L1 reg on scaling factors, free importance signal.
• APoZ: activation-based, data-dependent, catches dead neurons magnitude misses.
• OBD derivation: δL ≈ ½ h_iiw_i² after 3 assumptions (quadratic, converged, independent).
• Iterative pruning beats one-shot by keeping network near good minima throughout.
• Per-layer sensitivity: first conv sensitive, last FC robust → non-uniform sparsity budgets.

Bridge to Pruning II (Lecture 4)

This lesson covered the foundations. Pruning II dives into three open frontiers:

• The Lottery Ticket Hypothesis (Frankle & Carlin, 2019): sparse sub-networks ("winning tickets") that, when trained in isolation from scratch with the right random initialization, can match the full network's accuracy. This reframes pruning as sub-network discovery, not just weight removal.
• AMC — Automated Model Compression (He et al., ECCV 2018): use reinforcement learning to automatically determine the per-layer sparsity budget, replacing the hand-crafted sensitivity analysis with a learned policy. Input: layer spec. Output: optimal sparsity. 2× better latency-accuracy tradeoff than manual tuning.
• System support: how to represent sparse models efficiently in memory (CSR, compressed sparse column, 2:4 compressed format), and what the hardware actually does with them (NVIDIA sparse tensor cores, EIE, SpArch).

Related lessons on this site

• TinyML L1: Why Efficiency & Metrics — the roofline model, energy costs, and why DRAM access (640 pJ) dominates computation (0.9 pJ). Pruning reduces model size and DRAM loads.
• TinyML L2: NN Building Blocks & Compute — layer-level param and MAC formulas used throughout this lesson. Channel pruning shrinks C_out; the formulas give exact savings.
• LoRA & PEFT — a complementary compression direction: instead of removing weights, add a small low-rank adapter. Often combined with pruning for maximum compression.
• Quantization (TinyML L5) — reduce bitwidth rather than count. Orthogonal to pruning: you can quantize a pruned network for compounding savings.