Pruning: Removing Unnecessary Weights

Why It Matters for Modern Models

Pruning removes weights from a trained network, often followed by retraining. By pruning and retraining, Han et al. (2015) reduced AlexNet from 61 million to 6.7 million parameters with no loss of ImageNet accuracy.
SparseGPT prunes the 175- and 176-billion-parameter OPT and BLOOM language models in one shot, without retraining, reaching 60% unstructured sparsity with a negligible increase in perplexity (Frantar & Alistarh, 2023).
The lottery ticket hypothesis asks whether the small network could have been trained from the start: the winning tickets Frankle & Carbin (2019) found were less than 10% to 20% of the size of the original fully connected and convolutional networks for MNIST and CIFAR-10.
Speed-ups depend on the pattern of zeros: removing whole neurons or heads shrinks dense matrices, while scattered zeros need sparse kernels or hardware support, such as the 2:4 pattern on NVIDIA Ampere GPUs (Mishra et al., 2021).

Questions to Think About

Why does removing 90% of the weights at scattered positions often give little speed-up on a GPU, while removing whole neurons, heads or 2:4 blocks does?
Why does pruning a little at a time, with retraining in between, usually reach higher sparsity than pruning once (Han et al., 2015), and how does SparseGPT manage without retraining?

Interactive Demo

The interactive demo for this concept is in the newer notebook.

Open the notebook demo

Core Math

If you prefer intuition first, try the demo in the newer notebook, then come back for the math.

Key Equation

W_{\text{pruned}} = W \odot M, \quad M_{ij} = \mathbf{1}[|W_{ij}| > \theta]

Magnitude pruning removes the weights with the smallest absolute values, using a mask $M$ :

W_{\text{pruned}} = W \odot M, \quad M_{ij} = \mathbf{1}\big[|W_{ij}| > \theta\big]

Structured pruning removes whole neurons, channels or attention heads, which leaves smaller dense matrices:

W_{\text{pruned}} = W[:, \text{keep}]

Second-order criteria estimate how much the loss rises when weight $q$ is removed, using the Hessian $H$ of the loss at a trained minimum. Optimal Brain Damage uses only the diagonal of $H$ (LeCun, Denker & Solla, 1989):

\delta L_q \approx \tfrac{1}{2}\, H_{qq}\, w_q^2

When $H$ is positive definite, Optimal Brain Surgeon uses its inverse and also adjusts the remaining weights to compensate (Hassibi & Stork, 1992):

\delta L_q \approx \frac{w_q^2}{2\, [H^{-1}]_{qq}}

Lottery tickets. Frankle & Carbin (2019) found that a dense, randomly initialized network contains sparse subnetworks that, trained in isolation from the same initial weights $W_0$ , reach test accuracy comparable to the full network in a similar number of steps: