Efficient ML #3 — Neural Net Pruning

This is the third post (#3) in the Efficient ML series. If data types → quantization were about “representing numbers with fewer bits,” this time we go further and cut the connections themselves — pruning. The code lives in the pruning.ipynb notebook in github.com/warpspaceinc/efficient-ml-practice, and every plot below is our own numbers, measured by running that code on Colab (A100) ourselves.

The “brain damage” of 1989

When you talk about pruning, there’s one paper you can’t skip. Its title is downright gruesome — Optimal Brain Damage. The authors are Yann LeCun, John Denker, and Sara Solla of AT&T Bell Labs, and it was presented at NeurIPS 1989.

Stop and think about it for a second and it’s astonishing. 1989. That’s 36 years ago. Before the World Wide Web was released to the world, before LeCun became famous for CNNs (LeNet), back when the term “deep learning” didn’t even exist. That same LeCun is now Meta’s Chief AI Scientist and a Turing Award winner. The core idea of pruning — “if you cut away the less important connections in a neural network, you get a smaller, better-generalizing network” — had already been formalized that long ago.

The paper’s premise also touches on a biological metaphor. The human brain doesn’t just keep piling on synapses either. The number of synapses per neuron explodes from about 2,500 in infancy → up to 15,000 around age 2–4 → and then falls back to about 7,000 in adulthood. This “pruning” — snipping away the useless connections — is what actually produces an efficient brain.

This article follows how that 36-year-old insight becomes reality today, as 2:4 sparsity on NVIDIA GPUs, through three questions.

1. What to cut (criterion)
2. How much you can cut (retraining and the surprise of compression)
3. How to cut so the hardware actually speeds up (granularity and 2:4)

What pruning is & the 5-step framework

Pruning means deleting a connection by setting its weight to 0. You take a dense network and make it sparse, saving on storage, computation, and power. Formally, it’s the problem of “minimizing the loss while keeping the number of nonzero weights at or below $N$.”

Since Han et al. (NeurIPS 2015), pruning is usually organized into 5 steps: ① train → ② granularity (how) → ③ criterion (what) → ④ ratio (how much) → ⑤ fine-tune (recovery). This article walks through the core of it in the order criterion → ratio/recovery → granularity.

The motivation is the same as in the quantization article. Weights are mostly clustered near 0, and memory access is hundreds of times more expensive than computation (Horowitz: DRAM 640pJ vs int ADD 0.1pJ). So it’s natural to ask: “so many weights are almost 0 — do we really need all of them?”

1. What to cut — Criterion

The first question to answer is “which weights do we erase?” There are three candidates.

• Magnitude: importance $= |w|$. If it’s small, erase it. Cheap and simple.
• OBD saliency (second-order / Hessian): this is OBD’s answer. It uses a Taylor expansion to approximate how much the loss would grow if you erased each weight.
• Random: at random. A lower bound for comparison.

OBD’s derivation goes like this. Expand the loss $L$ to second order in the weight change $\delta w$:

Here we make three assumptions — convergence complete (the first-order term $g_i \approx 0$), diagonal approximation (ignore the cross terms $h_{ij}$), and second-order approximation (ignore everything third-order and up). Then the loss increase from setting a single weight to 0 ($\delta w_i = -w_i$) — the saliency — comes out cleanly.

It says the larger the weight ($w_i^2$), and the steeper the loss curvature ($h_{ii}$), the more important it is. The catch is that the Hessian $h_{ii}$ is expensive — so in practice we approximate the diagonal Hessian with the empirical Fisher (the mean of squared gradients). Our experiment did the same.

Measured directly: comparing the three criteria

After training an MNIST MLP (784→512→512→10), we raised the sparsity under each of the three criteria and measured accuracy (without retraining).

The key things to read:

• Random collapses the moment it passes 50%. Proof that what you cut is decisive.
• Magnitude and OBD nearly overlap. Both hold their accuracy up to 90%. At the extreme (98%) OBD’s second-order insight edges ahead (44.6% vs 29.3%), but right up until that point magnitude is actually marginally better.

Try it yourself below. As you raise sparsity with the slider, cut synapses disappear from the network on the right, and neurons left with no connection are removed. The compute (MACs·FLOPs) drops along with them — but that saving is only real if the hardware actually skips the zeros. Switch the criterion between magnitude ↔ random and, at the same sparsity (same MAC count), you see which connections and neurons vanish first change.

This result is exactly the answer to “why does practice use magnitude instead of OBD.” OBD is theoretically elegant, but the Hessian is expensive. Magnitude, on the other hand, is free and this powerful. That’s why, 36 years later, most pruning is still magnitude-based. The door OBD opened (“let’s principled-ly choose what to cut”) ended up being walked through, in practice, by the simplest criterion of all.

2. How much can you cut — retraining and the surprise of compression

Cut a lot at once and accuracy naturally drops. But if you repeat cut → retrain (fine-tune), the story changes completely. The retraining learning rate is usually set low, at about 1/10 to 1/100 of the original.

• prune only (gray): cut without retraining and at 95% you fall to 88%.
• iterative prune+retrain (green): cut a little at a time and retrain each round, and even cutting away 95% of the weights (= 20× fewer parameters) holds essentially the same accuracy as dense (97.3% vs 97.6%).

Keeping only one-twentieth of the weights and losing no performance is precisely why pruning is still researched today. Look at the weight distribution after pruning and you can see what happened at a glance — a giant spike appears exactly at 0.

The work that pushed this “cut and re-learn” idea to the extreme is Deep Compression (Han et al., 2016). Iterative pruning raised AlexNet’s pruning limit from 5× to 9×, and adding quantization and Huffman coding on top achieved up to dozens of times.

3. How to cut so it speeds up — Granularity and hardware

There’s a trap here that absolutely must be pointed out. Getting sparsity from pruning does not automatically make things faster.

Think about it carefully: “cutting” a weight is really no different from putting a 0 in its place. And if you just feed a 0 into a multiply-accumulate unit — $0 \times a = 0$ — the number of operations doesn’t drop at all. The accuracy experiments above actually set the weights to 0 and then ran them through the dense kernel as-is, so accuracy was measured but neither storage nor speed shrank one bit (exactly the “fake quantization” situation from the quantization article). Whether you multiply by 0 or by 0.3, from the multiplier’s point of view it’s the same single multiplication.

To reap the benefit, you have to not store the 0, not multiply it, and skip it entirely. But that turns out to be hard. Pruning so far has been the unstructured kind, where zeros scatter anywhere, so the compression ratio is best — but from the GPU’s point of view, the positions of the zeros are irregular, so you can’t skip them. The cost of finding “where is the next zero” every time is more expensive than the single multiplication you’d save, so in the end running dense is better.

That’s why this principle always tags along with pruning:

Pruning only yields a real payoff when the hardware supports that sparsity pattern. You have to know what the hardware you’ll use accelerates, and cut accordingly.

Dedicated hardware for sparse computation (EIE, ISCA 2016, etc.) has been researched over the years, and the most widely used answer today is NVIDIA’s 2:4 structured sparsity.

2:4 pruning

The rule is simple. For every 4 consecutive weights, set exactly 2 to 0 (50% sparsity). Because the positions of the zeros are regular — “2 out of every 4” — NVIDIA’s Sparse Tensor Core (from Ampere / A100 onward) accelerates this pattern in hardware. Storage is efficient too — you only need half the nonzeros plus a 2-bit index pointing to each value’s position.

Here’s how to build a mask that cuts the smaller 2 of every 4 columns by magnitude.

def two_four_mask(weight):
    """A 2:4 mask that keeps only the 2 largest |w| per group of 4 columns."""
    w = weight.detach().abs(); R, C = w.shape
    mask = torch.zeros_like(weight)
    for j in range(0, C - C % 4, 4):
        idx = w[:, j:j+4].topk(2, dim=1).indices   # the larger 2 of each group of 4
        mask[:, j:j+4].scatter_(1, idx, 1.0)
    return mask

The result of applying 2:4 to our MLP: dense 97.6% → 2:4 (no retraining) 96.8% → 2:4 + retraining 98.2%. Even cutting half regularly, retraining fully recovers it.

Using 2:4 from Python

PyTorch provides an API to store the 2:4 pattern compressed and multiply it on the Sparse Tensor Core (requires fp16, Ampere+).

from torch.sparse import to_sparse_semi_structured

mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).bool()
W = (torch.rand(64, 64) * mask).half().cuda()
Ws = to_sparse_semi_structured(W)          # compressed to nonzeros + 2-bit indices
x  = torch.rand(64, 64).half().cuda()

torch.mm(Ws, x)                            # uses the Sparse Tensor Core

Does it really speed up — benchmarked directly

On an A100 we measured dense vs 2:4 sparse matmul speed by size. The result is honest and instructive.

• For small matrices (N < 4096), 2:4 is actually slower. The compression and indexing overhead outweighs the benefit.
• The benefit only shows up on large matrices. At N=8192 it’s about 1.72× faster.

This is the concrete meaning of “you have to know your hardware and workload.” The very same 2:4 pruning only speeds up on an Ampere-or-later GPU, and only for sufficiently large matrices. On Turing (T4) or small matrices there’s no benefit — or even a loss. Pruning isn’t purely an algorithm problem; it’s a problem you have to design together with the hardware from start to finish.

4. Pruning × quantization are orthogonal — Deep Compression

We’ve looked at pruning up to here, and in the earlier quantization article we reduced the bits. The key point is that these two are orthogonal. Pruning decides “which weights to erase,” while quantization decides “how many bits to represent the surviving weights with.” The axes are different, so using them together multiplies the gains.

The best demonstration of this is the 3-stage pipeline of Deep Compression (Han et al., 2016).

1. Pruning + retraining — erase the weak connections and re-learn to recover.
2. Quantization + retraining — group the surviving weights into a handful of K-Means codebook entries (retraining the codebook) and store only the indices.
3. Huffman coding — since the index distribution is skewed to one side, compress once more, losslessly.

We ran it on our own MLP. Watching how the fc2 weight distribution changes at each stage, you can see at a glance that pruning and quantization touch different parts of the distribution. (Zeroed-out weights are hidden; only the survivors are plotted.)

• ① dense — the familiar bell-shaped normal distribution.
• ② pruning (80%) — the small weights near 0 vanish wholesale. The middle is blown away and only the two tails remain.
• ③ retraining — the surviving weights re-fit to the data, becoming a clear bimodal distribution.
• ④ quantization (4-bit) — that continuous distribution folds into 16 discrete levels (this is what the K-Means in the [quantization article] does).

Pruning empties out the middle of the horizontal axis (the small values), and quantization discretizes the remaining values into a few vertical lines. Since they touch different places, they stack without overlapping. As a result, the model size shrinks multiplicatively at each stage.

• dense 2,612 KB → pruning 604 KB (4.3×) → +quantization 147 KB (17.8×) → +Huffman 141 KB (18.6×)
• Meanwhile accuracy was actually held and improved, from 97.5% → 98.3% (thanks to retraining).

So pruning alone gives 4.3×, adding quantization gives 17.8×, and Huffman on top gives 18.6×. The gains of each technique don’t add — they multiply. The original paper used this pipeline to shrink AlexNet by 35× and VGG-16 by 49× (with no accuracy loss). Data types → quantization → pruning: the three axes this series covered are ultimately orthogonal tools you can stack onto one model at the same time.

Run it yourself

Every plot above came from running the notebook below on Colab. Runtime → Run all is all it takes (for the 2:4 speed measurement, select an Ampere GPU such as an A100).

• 📓 pruning.ipynb — OBD vs magnitude, retraining compression, 2:4 sparsity
• Repository: github.com/warpspaceinc/efficient-ml-practice

Wrap-up

• What: random collapses quickly. Magnitude (|w|) is cheap and powerful, matching OBD’s second-order saliency — which is why it became the practical standard.
• How much: with repeated prune→retrain, even cutting 95% (20×) holds accuracy.
• How: unstructured only boosts the compression ratio. It takes 2:4 structured for the Ampere Sparse Tensor Core to actually accelerate — but only on large matrices.
• Together: pruning and quantization are orthogonal. Stack them the way Deep Compression does and the gains multiply (18.6× in our own measurement, with accuracy held).

The insight Optimal Brain Damage planted in 1989 still holds. Only, for it to become a “real payoff,” the hardware 36 years later had to be ready to receive that sparsity. The history of pruning is, in a way, the history of algorithms and hardware waiting for each other.

In the next article we’ll cover how to set sparsity automatically, layer by layer (sensitivity analysis · AMC · NetAdapt).

References

• Y. LeCun, J. S. Denker, S. A. Solla. Optimal Brain Damage. Advances in Neural Information Processing Systems (NeurIPS), 1989.
• S. Han, J. Pool, J. Tran, W. J. Dally. Learning both Weights and Connections for Efficient Neural Networks. NeurIPS, 2015.
• S. Han, H. Mao, W. J. Dally. Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR, 2016.
• A. Mishra et al. Accelerating Sparse Deep Neural Networks (2:4 Structured Sparsity). arXiv:2104.08378, 2021.
• S. Han et al. EIE: Efficient Inference Engine on Compressed Deep Neural Network. ISCA, 2016.
• M. Horowitz. Computing’s Energy Problem (and what we can do about it). ISSCC, 2014.
• Course reference: MIT 6.5940 TinyML and Efficient Deep Learning Computing (Song Han), Lec 3–4 Pruning & Sparsity.