Fine-tuning, iterative pruning and regularization

Take the AlexNet that Han, Pool, Tran and Dally started from in 2015: 61 million parameters, five convolution layers and three fully connected ones, 57.2% top-1 accuracy. They kept roughly one weight in nine and ended with a network that was not measurably worse. That is the headline of this part, and the whole trick behind it is what happens after the cut.

Now look at what those hours say. Retraining took more than twice as long as training. The authors did not treat it as optional polish; section 3 of the paper calls the final retraining step critical and states that if the pruned network is used without retraining, accuracy is significantly impacted. So the answer to the outline's fifth question, how to fine-tune a pruned network, is a method, not a footnote: the Pruning mask decides which weights exist, and Fine-tuning decides whether the survivors can still do the job.

Why the survivors are wounded, and why the wound heals

A trained layer is a set of weights that were optimized together. Zeroing the smallest of them, the magnitude criterion from lecture 04-1, is a small perturbation to any one output, but a large Pruning ratio removes so many small contributions that the outputs shift and the accuracy drops, more steeply the higher the ratio (this is what the chart on the slide shows for the dashed curve). The survivors are still near a good solution, though. Nothing about them was wrong; the function they were tuned to complete simply lost some of its terms. Fine-tuning continues gradient descent on exactly those survivors, with the removed weights held at zero, so they can absorb the work of the missing ones.

This is why the learning rate matters so much. The slide states the rule: the fine-tuning learning rate is usually 1/100 or 1/10 of the original learning rate. In the paper, LeNet was retrained at 1/10 of the original rate (section 4.1) and AlexNet at 1/100 of the initial rate (section 4.2). A large step would throw the weights out of the basin in which the magnitude ranking was measured, undoing the very judgement the pruning made. Small steps walk the network back down the loss surface from where the cut left it.

Slide 41 makes the same point as an optimization statement. Pruning poses the problem as minimizing the loss over the pruned weights W_P subject to a budget N on the L0 norm, the count of nonzero weights. The mask chooses which entries of W_P may be nonzero at all (the constraint), and fine-tuning is the argmin over the survivors with that mask fixed: the best values the kept weights can take.

The paper adds a second reason to keep the survivors instead of re-initializing them. Section 3.3 observes that networks contain fragile co-adapted features: gradient descent finds a good solution when the network is first trained, but not after re-initializing some layers and retraining them. Fine-tuning starts from the co-adapted weights and preserves that structure. The paper also notes that during retraining it can help to fix the convolution parameters while the fully connected ones retrain, and vice versa, to avoid vanishing gradients through the sparse layers.

The three-box pipeline

The figure on the slide is Figure 2 of the paper. The first box, Train Connectivity, is ordinary training, but its purpose is reframed: it learns which connections are important, not their final values. The second box, Prune Connections, applies the magnitude threshold. The third box, Train Weights, retrains the sparse network. In Han et al.'s implementation the threshold was a quality parameter multiplied by the standard deviation of each layer's weights, and Caffe was modified to add a mask that disregards pruned parameters during network operation for each weight tensor.

Reading the chart on slide 43

The chart plots accuracy loss against the fraction of parameters pruned away for AlexNet. The dashed purple curve is pruning with no retraining; the green curve adds fine-tuning. Reading the points off the figure gives the following.

The paper summarizes the same curves in section 5: without retraining, accuracy begins dropping much sooner, at one third of the original connections rather than one tenth. It describes a free lunch of removing half the connections with no loss even without retraining, while with retraining the connections can be reduced by 9x. Read that 9x carefully: it is reached with the iterative scheme of the next two concepts. A single prune and retrain, the green curve, is already at about -1.7% by 90% pruned. Two green points (at about 67% and 75% pruned) sit slightly above zero, which the authors attribute to pruning finding the right capacity of the network and thereby reducing overfitting. Because sparsity and compression ratio are two ways of saying the same thing, keep the conversion at hand: ratio = 1 / (1 - s).

Once retraining is part of the recipe, a natural question follows: why cut to the final sparsity in one go? The slides tell the story as a sequence. Prune to 30%, retrain the weights. Prune to 50%, retrain. Prune to 70%, retrain. The only change to the pipeline figure is a single arrow from Train Weights back up to Prune Connections, and that arrow is the whole idea of Iterative pruning.

Define the unit first: one iteration is a prune followed by a fine-tune. Iterative pruning repeats that unit while gradually increasing the target sparsity in each iteration, instead of jumping to the final sparsity in one step. Han et al. put it plainly in section 3.4: learning the right connections is an iterative process, pruning followed by retraining is one iteration, and after many such iterations the minimum number of connections can be found. Each iteration is a greedy search for the best connections. They also tried pruning parameters probabilistically by absolute value and found it gave worse results.

Why gradual beats one-shot

The magnitude criterion ranks weights on the network as it is right now. After a 30% cut and a fine-tune, the survivors have re-settled: some grew to cover the removed terms, some shrank because their job was taken over. The ranking for the next cut is therefore measured on a network that has adapted, and the weights it marks as small really are the least needed. A one-shot cut to 90% ranks all 90% on a network that has never adapted to any loss, so it discards weights that would have become important once their neighbors were gone. Zhu and Gupta describe the same principle in the general form of a sparsity schedule: the binary mask is updated every few hundred steps to gradually increase the sparsity of the network while allowing the training steps to recover from any pruning-induced loss in accuracy. Their cubic schedule prunes quickly at first, when redundant connections are abundant, and slowly near the end.

The schedule in this part uses five hand-picked targets rather than a closed form, but the shape is the same: larger cuts early, smaller ones late. On a 6 x 6 matrix of 36 weights, rounding 36 s gives the following zero counts, which are also what the simulator below shows.

The loop also connects backwards to part 05. NetAdapt runs a short-term fine-tune after every layer it thins and a long-term fine-tune at the end, which is the same prune-then-fine-tune loop driven by a latency budget instead of a sparsity schedule. Whatever criterion or granularity you choose, the outer loop looks like this figure.

Slide 51 adds one bullet and one curve, and together they carry the number you will be asked for: Iterative pruning boosts the Pruning ratio from 5x to 9x on AlexNet compared to single-step aggressive pruning. The red curve is what that sentence looks like.

Notice where the curve begins. It does not start at 40% like the others. The paper explains why in section 5: the biggest gain comes from iterative pruning, where the pruned and retrained network of the green curve is pruned and retrained again. The leftmost red dot corresponds to the point on the green line at 80% (5x) pruned further to 8x. There is no accuracy loss at 9x, and not until 10x does the accuracy begin to drop sharply.

The same paper reaches 13x on VGG-16 with five iterations of pruning and retraining, shrinking it to 7.5% of its original size with fc6 and fc7 each pruned to less than 4% of their original size. These pruned networks are the starting point of Deep Compression, which quantizes the surviving weights to 8 bits in convolution layers and 5 bits in fully connected layers and Huffman codes the result to reach 35x on AlexNet (240 MB to 6.9 MB) and 49x on VGG-16. The roughly 90% Weight sparsity that EIE assumes in the next part is this number.

Fine-tuning repairs the network after the cut. Regularization (for pruning) prepares it before the cut. Follow one weight w = 0.05 through training with a penalty coefficient lambda = 0.01 and a learning rate eta = 0.1, ignoring the data loss for the moment so only the penalty acts.

The two regularized losses

The slide states the general rule. During training, or during the fine-tuning of a pruned network, a penalty is added to the loss to penalize nonzero parameters and to encourage smaller ones. The two common choices are the L1 and L2 penalties on the weights W, with lambda setting how strongly the penalty competes with the data loss L(x; W).

Goodfellow, Bengio and Courville make both halves precise. L2 regularization, commonly known as weight decay, drives the weights closer to the origin; its update multiplicatively shrinks the weight vector by a constant factor on each step (their equation 7.5). L1 regularization, in comparison, results in a solution that is more sparse, where sparsity means that some parameters have an optimal value of exactly zero, and this property is why L1 has been used extensively for feature selection (LASSO). The gradient forms above are the reason. Near zero the L1 gradient lambda sign(w) keeps its full size, so the weight is pushed onto zero and pinned there. The L2 gradient 2 lambda w shrinks with the weight, so the push fades before it arrives.

Why both help pruning, and why L2 won anyway

Either penalty shrinks the weights the data loss does not defend. That does two things for magnitude pruning. The ranking becomes a cleaner signal, because a weight that stayed large did so against a constant pull, and the mass of the distribution near zero grows, so more weights fall under the threshold at a given accuracy. L1 goes further and parks weights on exactly zero, so removing them changes nothing at all.

It is tempting to conclude that L1 must be the better choice for pruning, and this is the classic trap. Han et al. (section 3.1) report that L1 regularization does give better accuracy after pruning but before retraining, since more parameters sit near zero. However, the remaining connections are not as good as with L2, resulting in lower accuracy after retraining, and overall L2 regularization gives the best pruning results. They also tried L1 for the pruning phase followed by L2 for retraining and found it did not beat using L2 for both, because parameters from one mode do not adapt well to the other. That is why the slide says magnitude-based fine-grained pruning applies L2 regularization on weights.

Network Slimming: the same penalty on channel scaling factors

The second example on the slide moves the penalty from weights to channels, which connects back to the scaling-based criterion of lecture 04-1. Network Slimming (Liu et al., ICCV 2017) notes that every batch normalization layer already computes z_out = gamma z_hat + beta per channel, so gamma is a free scaling factor for that channel. Their training objective adds a sparsity penalty on those factors.

Channels whose gamma has been pushed toward zero contribute almost nothing and are removed with a global percentile threshold (pruning 70% of channels means choosing the 70th percentile of all gamma values), followed by fine-tuning. Because a whole channel goes, this is Channel pruning, the coarse end of Pruning granularity, and the resulting network stays dense. On VGGNet the method reports a 20x smaller model and 5x fewer computing operations. The paper also describes a multi-pass scheme that repeats train, prune and fine-tune, the same loop as the previous concepts, now driven by gamma instead of |w|.