Train, prune, fine-tune: what pruning achieves

Take AlexNet as Han, Pool, Tran and Dally trained it for their NeurIPS 2015 paper: 57.2% top-1 and 80.3% top-5 on ImageNet. Its first fully connected layer alone holds 38 million weights, and if you histogram them you get a narrow bell centered on zero, almost all of it inside [-0.015, 0.015] by the paper's own account. Now pick a threshold that removes the smallest half of every layer. Nothing happens to accuracy. Remove three quarters, and top-5 accuracy falls by about 2.2%. Remove four fifths, and it falls by about 4%. That single experiment is the whole of slides 7 and 8.

The rule behind it is a two-step recipe. Step one is ordinary training, but Han reads it differently: the dense network is not trained to learn final weights, it is trained to learn which connections are important. Step two removes every weight whose magnitude falls below a threshold, which the paper sets as a quality parameter times the standard deviation of that layer's weights. A dense layer becomes a sparse layer. This is Pruning at its simplest, Synapse pruning in the vocabulary of Part 01, and it is the plainest form of Magnitude-based pruning: importance is |w| and nothing else. Once connections are gone, some neurons are left with no surviving input or no surviving output. Those neurons contribute nothing and are removed too, so Neuron pruning falls out of synapse pruning for free.

The right-hand inset of slide 8 shows exactly that gap: the same bell with its center cut out, a hole of width 2t around zero, and nothing else changed. The paper's discussion of its Figure 7 describes the distribution as "centered around zero with tails dropping off quickly" before pruning, then "the center of the distribution" removed. The weights that survive have not moved. They will move in the next concept, and that movement is what retraining does.

Reading the prune-only curve

The chart on slide 8 plots top-5 accuracy loss against the Pruning ratio, the fraction of parameters pruned away, from 40% to 100% on the x axis and from +0.5% to -4.5% on the y axis. The dashed purple curve is what happens with no retraining. It sits at zero up to 50%, which the paper calls a "free lunch of reducing 2x the connections without losing accuracy even without retraining". It reaches about -1% at 67%, about -2.2% at 75%, about -4% at 80%, and leaves the chart just past 82%. Without retraining, the paper says, accuracy "begins dropping much sooner", once only a third of the connections remain.

The x axis counts what is removed, so you have to translate it into a reduction factor before you can compare with a table that reports "9x". Keeping a fraction 1 - r of the weights makes the model 1 / (1 - r) times smaller.

Where the threshold comes from

Because the criterion is magnitude, the threshold and the pruning ratio are two views of the same cut. If the weights of a layer are roughly Gaussian with standard deviation sigma, removing a fraction r means removing everything inside ±t where the Gaussian mass inside ±t equals r. That is an inverse normal lookup: t = sigma × z with z = 0.674 for 50%, z = 1.282 for 80% and z = 1.645 for 90%.

Return to the AlexNet that lost 4% at 80% pruned. Keep the mask fixed, so the removed weights stay at zero, and train the survivors again with a learning rate one hundredth of the original. The loss comes back to 0.0%. That run took 173 hours on a Titan X against 75 hours for the original training, which is why the paper says pruning is "not used when iteratively prototyping the model, but rather used for model reduction when the model is ready for deployment".

This is step three of the pipeline, and the paper is blunt about its status: retraining "learns the final weights for the remaining sparse connections. This step is critical. If the pruned network is used without retraining, accuracy is significantly impacted." The lecture calls it Fine-tuning. The green curve on slide 9 shows what it buys. It stays at zero, or a hair above, all the way to 80%, then bends: about -0.25% at 84%, about -0.9% at 88%, about -1.7% at 90%, about -4% at 93%. The paper summarizes the same curve as "with retraining we are able to reduce connections by 9x".

Two of the green points sit slightly above zero. The authors do not treat that as noise: "We believe this accuracy improvement is due to pruning finding the right capacity of the network and hence reducing overfitting." In other words, Pruning behaves like a regularizer, and an over-parameterized network can lose most of its weights and generalize a little better for it.

What retraining does to the weights

The third inset on slide 9 is the most informative picture in this part. After pruning, the histogram was a bell with a hole. After retraining it is two smooth lobes, one on each side of zero, and the whole distribution is wider. The paper's text on Figure 7 states the numbers: the original weights lived inside [-0.015, 0.015]; after retraining "the parameters form a bimodal distribution and become more spread across the x-axis, between [-0.025, 0.025]". The survivors have grown in magnitude to take over the work of the connections that were removed. Nothing has been restored; the mask still holds the pruned weights at zero.

Three rules the paper gives for retraining

• Keep the surviving weights; do not reinitialize them. The paper argues that networks contain "fragile co-adapted features" that gradient descent finds when the network is first trained but cannot find again from scratch on a sparse layout.
• Shrink dropout. Pruning has already removed capacity, so the retraining dropout D_r is scaled from the original D_o by the square root of the fraction of connections kept.
• Choose the regularizer for the stage. L1 pushes more weights toward zero and gives better accuracy straight after pruning, but L2 gives better accuracy once the survivors are retrained, so the paper uses L2 for the curves on these slides.

One more observation from the paper's sensitivity study (Figure 6) matters for your own experiments: convolutional layers are more sensitive to pruning than fully connected ones, and the first convolutional layer is the most sensitive of all, because its input has only three channels and there is little redundancy to remove. That is why the per-layer keep rates in the next concepts are so uneven.

Start from the retrained network at 80% pruned, the last green point still on the zero line, a 5x model with no loss. Now prune it again, using the retrained magnitudes, and fine-tune again. The paper describes what happens: "The leftmost dot on this curve corresponds to the point on the green line at 80% (5x pruning) pruned to 8x. There's no accuracy loss at 9x. Not until 10x does the accuracy begin to drop sharply." The red curve on slide 10 is that experiment: flat through 90%, about -0.5% at 92%, -1% at 93%, -2% at 94%, and about -4% between 95% and 96%.

This is Iterative pruning, and the paper states the rule and the payoff in one breath: "Pruning followed by a retraining is one iteration, after many such iterations the minimum number connections could be found. Without loss of accuracy, this method can boost pruning rate from 5x to 9x on AlexNet compared with single-step aggressive pruning." Each iteration is a greedy search. VGG-16 used five rounds; ResNet-50 in Han's thesis used three. The left side of slide 10 draws the loop: train connectivity, prune connections, train weights, and an arrow from train weights back to prune.

Why several small cuts beat one big cut

A single cut to 90% judges every weight by the magnitude it had in the dense network. But the previous concept showed that retraining changes those magnitudes: the survivors grow and the distribution becomes bimodal. Weights that looked expendable in the dense network may become load-bearing after the first round, and weights that looked important may shrink once their neighbours take over. Iterating lets each cut use the freshest evidence. Frankle and Carbin reach the same conclusion in a different setting: iterative pruning over n rounds, each removing p^(1/n) of what remains, "finds winning tickets that match the accuracy of the original network at smaller sizes than does one-shot pruning". The price, which they also state, is that repeated retraining is expensive.

Look at AlexNet layer by layer, as the paper's Table 4 does. Its first fully connected layer, fc6, holds 38 M weights but costs only 75 M FLOPs, because each weight is used once per image. Its second convolutional layer, conv2, holds 307 K weights but costs 448 M FLOPs, because each weight is reused at every spatial position. Add the layers up and the network splits into two worlds: the three FC layers hold 58.6 M of the 61 M parameters (96.2%, which the paper rounds to 59 M) but only about 117 M of 1.45 G FLOPs (8%), while the five conv layers hold 2.3 M parameters (3.8%) but 1.33 G FLOPs (92%).

That split explains the table on slide 11 before you read it. Pruning fc6 to 9% deletes 34.6 M parameters but only 68 M FLOPs by weight count (73 M once zero activations are skipped as well, which is how Han counts). Pruning conv2 to 38% deletes only 190 K parameters but 278 M FLOPs by weight count (300 M by Han's count). Parameter reduction is a story about FC layers; MAC reduction is a story about conv layers, and the two are pruned to very different depths.

Why the over-parameterized nets compress more

AlexNet and VGG-16 get 9x and 12x; GoogleNet, ResNet-50 and SqueezeNet get about 3.2x to 3.5x. The extra column is the reason. VGG-16's fc6 alone holds 103 M of its 138 M weights, and the paper prunes fc6 and fc7 to 4% of their size while conv layers keep 22% to 58%. GoogleNet has one FC layer of about 1 M in 7 M, ResNet-50 one of about 2 M in 25.5 M, and SqueezeNet has none at all, ending in global average pooling. Han's thesis says it directly: the pruning ratio of GoogleNet is smaller than AlexNet and VGG-16 "because convolutional layers dominate GoogleNet, and convolutional layers are much more efficient than fully connected layers". The three conv-dominated networks land in the same place, about 30% of parameters nonzero (GoogleNet 29%, SqueezeNet 31%, ResNet 29%), which is the 3.4x in the table. Their FC rows, where Neuron pruning would delete whole rows of a weight matrix, are too small to matter.

Why MAC reduction is not parameter reduction

Two mechanisms separate the two columns. The first is the split above: parameters concentrate in FC layers, which are pruned hardest, and MACs concentrate in conv layers, which are pruned least. That is how AlexNet gets 9x on parameters but 3x on MACs, and VGG-16 12x against 5x. The second is how Han counts. A multiply-accumulate is only counted as saved when it is really skipped, and it is skipped when either the weight or the input activation is zero. ReLU produces plenty of zero activations, so the FLOP column in Table 4 is roughly the weight density times the density of the layer's input activations.

Even where MACs barely fall, the parameter count is worth cutting. A 6.7 M parameter AlexNet fits in on-chip SRAM, and the paper's point from Part 01 stands: at 640 pJ per DRAM access, the energy of fetching weights dominates the energy of multiplying them.

Everything so far was an image classifier with a huge fully connected tail. NeuralTalk, Karpathy and Fei-Fei's image captioner, is a different animal: a CNN feature extractor feeding an LSTM that writes a sentence one word at a time. Han froze the CNN, pruned every LSTM weight matrix except the word-embedding table to 10% nonzeros, retrained with the original weight decay and batch size, and measured BLEU-1 to BLEU-4 on Flickr-8K. The thesis reports: "Not until pruning away 90% of the parameters does the BLEU score begin to drop sharply."

The recipe did not change. Train, cut the small weights, retrain the survivors: the same three steps, the same knee near 90%, and the same dependence on Fine-tuning, which the thesis says "plays a very important role" when it compares the retrained curve with the one that skips retraining. Slide 12 shows the qualitative check that a BLEU number cannot: the captions themselves.

At 90% the pruned model "sometimes produces the same caption, sometimes produces a different word to describe the same thing", and sometimes describes the scene differently but still correctly. Replacing "ball" with "basketball" is arguably an improvement. The fourth image is the warning. At 95% pruned, a 20x model, the sentence starts to stutter: "a red shirt and black and white black shirt" repeats a phrase and loses the word "soccer". That is what crossing the knee of the Pruning ratio curve looks like in a language model: not silence, but drift.

Take fc6 of AlexNet pruned to 9% and run it with an ordinary dense matrix kernel on a GPU. The kernel does not know the zeros are zeros. It multiplies all 38 M weight-activation pairs, 34.6 M of them by zero, and finishes in exactly the time the dense layer took. Nine times fewer parameters, zero speedup. The reduction table of two concepts ago counted what could be skipped; it did not promise that anything would be. Skipping needs either a custom engine that understands sparse formats or a sparsity pattern that commodity hardware is built for. Slide 13 shows both.

The commodity route: NVIDIA A100 and 2:4

The Ampere whitepaper defines 2:4 sparsity as a structure "that allows two non-zero values in every four-entry vector". Every group of four contiguous weights along a row keeps at most two, so the matrix is exactly half zeros, stored as the nonzero values plus a 2-bit index per survivor. That compression cuts storage and bandwidth "by almost 2x", and the Sparse Tensor Core has a matrix multiply instruction that reads the indices, gathers only the matching activations, and finishes a tile in N/2 cycles instead of N: "a 2x speedup". This is the special case N = 2, M = 4 of N:M sparsity.

How does a model get into that pattern? The whitepaper's recipe is the pipeline of this part with one constraint added: the network is "first trained using dense weights, then fine-grained structured pruning is applied, and finally the remaining non-zero weights are fine-tuned". NVIDIA's developer blog reports BERT-Large keeping its SQuAD F1 of 91.9 dense and sparse. The magnitude criterion, the mask and the retraining are all Han's; the only change is that the two smallest of every four are cut instead of the smallest 50% of the layer.

The gap between 2x and 1.5x is the gap between a peak instruction rate and a whole layer. The whitepaper says "up to 2x". NVIDIA's cuSPARSELt measurements on BERT-Large GEMM layers on an A100 give 1.3x for the projection, 1.4x for QKV and FC1, and 1.6x for FC2 over dense cuBLAS, because memory traffic for activations, the non-GEMM parts of the layer and kernel launch costs do not halve. The slide's "1.5X measured BERT speedup" is Han's summary of those numbers. Treat 1.5x as the realistic middle and 2x as the ceiling.

The custom route: engines built to skip zeros

Before 2:4 existed, the way to profit from unstructured sparsity was to build the hardware yourself. The four papers cited on slide 13 are all from Han and his students, and each makes a different thing sparse. EIE, the Efficient Inference Engine, holds the compressed model in on-chip SRAM, so the 640 pJ DRAM access from Part 01 disappears, and it skips both zero weights and the zero activations that ReLU produces, which its authors say "saves another 3x". ESE moves the same idea onto an FPGA for pruned speech-recognition LSTMs, the same recipe the previous concept applied to NeuralTalk. SpArch accelerates multiplying two sparse matrices, and SpAtten prunes tokens and attention heads at run time instead of weights.