Pruning Theory and Intuition

Imagine that you want to cut down on comforts in your living space – you think that there might be more than you need to work and that keeping all the comforts is increasing your cost of living beyond what it could be. You’ve made changes in your living space with an eye toward minimally reducing the impact on your ability to work – you still keep your computer, good Wi-Fi, steady electricity. However, you’ve cut down on comforts that may aid your work but that are fundamentally auxiliary by removing or donating items like a television subscription, a nice couch, or tickets to the local symphony.

This change in living space shouldn’t theoretically impact your work, provided you are reasonably resilient – your mental facilities have not been explicitly damaged in a way that would impair your ability to perform the functions of work (perhaps it affects your comfort, but that’s not a factor in this discussion). Meanwhile, you’ve managed to cut down on general costs of living.

However, you’re a bit disoriented: you reach instinctively for a lamp in the corner that has been removed. You find yourself disappointed that you can’t watch unlimited television anymore. Making these changes to this space requires reorienting yourself to it. It takes a few hours (or even days) of exploration and orientation to acclimate toward these new changes. Once you have completed acclimating, you should be ready to work in this newly modified space just as well as you did in your prior living space.

It should be noted, though, that if the disparity in your prior living space and your current living space is too large, you may never be able to recover – for instance, removing your exercise equipment, all sources of entertainment, and other items very close to but still not directly impacting your work. If you cut further into your living costs by cutting down on running water and electricity, your work ability would be directly impaired.

Let’s rewind: you step back and take a look around your living space and decide that it has too many unnecessary comforts, and you want to cut these comforts down. You could take away all the comforts at once, but you decide instead that the immediate, absolute difference might be too stark for you to handle. Instead, you decide to embark on an iterative journey, in which you remove one or two things of the least significance every week, and stop whenever you feel that removing any more items would result in damage to your core working facilities. This way, you have time to acclimate to a series of small differences.

The logic of cutting down comforts in living spaces applies in parallel to the logic of pruning – it provides a useful intuitive model when feeling out how to perform pruning campaigns.

Pruning was initially conceived in Yann LeCun’s 1990 work in “Optimal Brain Damage” – not all parameters contribute significantly to the output, so those parameters can be pruned away in the most optimal form of “brain” (neural network) damage.

Pruning is generally performed after the network has been substantially trained, such that evaluation of parameter importance is meaningful and not based only on random initialization or values in the early stages of training. In order to determine which neural network entities (nodes, connections, layers, etc.) contribute the most or least significantly to the output, each entity must be evaluated according to some importance criterion. The least important entities are removed (Figure 4-3). In practice, removal simply means setting a parameter to zero, which is much cheaper to store.

This action of pruning away entire connections, nodes, and other neural network entities can be thought of as a form of architecture modification. For optimal performance, the model needs to be reoriented toward its new architecture via fine-tuning. This fine-tuning can be performed simply by training the new architecture on more data.

Thus, pruning follows the following general iterative process (Figure 4-4).

In this sense, you can think of pruning as a “directed” dropout, in which connections are not randomly dropped but instead pruned by the importance criterion. Because pruning is directed rather than random, a more extreme percentage of weights can be pruned with reasonable performance compared to the reasonable percentage of weights dropped in dropout.

When implementing pruning, you can specify the initial percent of parameters to be pruned and the end percent of parameters to be pruned; TensorFlow will figure out how much to remove at each step for you.

Each neural network and each task require a differently ambitious pruning campaign; some neural networks are built relatively lightweight already, and further pruning could severely damage the network’s key processing facilities. On the other hand, a large network trained on a simple task may not need the vast majority of its parameters.

Using pruning, 90% to 95% of the network’s parameters can reliably be pruned away with little damage to performance.

Pruning Implementation

To implement pruning (as well as other model compression methods), we’re going to require the help of other libraries. The TensorFlow Model Optimization library works with Keras/TensorFlow models but is installed separately (pip install tensorflow-model-optimization). It should be noted that the TensorFlow Model Optimization library is relatively recent and therefore less developed than larger libraries; you may encounter a comparatively smaller forum community to address warnings and errors. However, the TensorFlow Model Optimization documentation is well written and contains additional examples, which can be consulted if necessary. We will also need the os, zipfile, and tempfile libraries (should be included in Python by default), which allow us to understand the cost of running a deep learning model.

Although TensorFlow Model Optimization significantly aids in the code required to implement pruning, it involves several steps and needs to be approached methodically. Moreover, note that because intensive work in pruning and model compression broadly has been relatively recent, at the time of this book’s writing, TensorFlow Model Optimization does not support a wide array of pruning criterion and scheduling. However, its current offerings should satisfy the pruning needs of most compression needs.

Pruning Individual Layers

Recall that when pruning the entire model, we called tfmot.keras.sparsity.prune_low_magnitude() on the entire model. One approach to prune individual layers is to call tfmot.keras.sparsity.prune_low_magnitude() on individual layers as they are being compiled. This is compatible with layer objects in both the Functional and Sequential API.

In this example neural network, we prune all layers other than the first and last Dense layers after the Input layer and before the output layer (Listing 4-15). When choosing which layers to prune, avoid ambitiously pruning initial layers responsible for feature extraction and layers critical to the knowledge-building abilities of the model.

from tfmot.sparsity.keras import prune_low_magnitude as plm

pruned_model = keras.Sequential()

pruned_model.add(L.Input((784,)))

pruned_model.add(L.Dense(2**9))

pruned_model.add(plm(L.Dense(2**8), **pruning_params))

pruned_model.add(plm(L.Dense(2**7), **pruning_params))

pruned_model.add(plm(L.Dense(2**6), **pruning_params))

pruned_model.add(plm(L.Dense(2**5)))

pruned_model.add(L.Dense(10, activation='softmax'))

Listing 4-15

Pruning individual layers by adding wrappers around layers. Activations are left out for the purpose of brevity

The benefit of pruning layers independently is that you can use different pruning schedules for different layers, for instance, by less ambitiously pruning layers that have less parameters to begin with. The model can then be compiled and fitted with the UpdatePruningStep() callback as discussed earlier and fine-tuned afterward.

However, the disadvantage of this method of selecting layers to prune is that you can’t do any pretraining before pruning, since layers are wrapped in the pruning wrapper from they are defined. This results in a worse outcome than if the model was pretrained on data before pruning. To select specific layers for pruning on a model that has already been trained, we need to clone the model using keras.models.clone_model(model) with a cloning function. The cloning function maps each layer to another layer; in this case, we can map layers that we want to prune with a pruned version of that layer (Figure 4-5).

Let’s construct a cloning function that either maps a layer to a pruned version of itself or returns the original layer if we do not desire to perform pruning on it (Listing 4-16). There are many ways you can select which layers to quantize; you can prune certain types of layers, layers by name, their position in the network, and so on. If a layer satisfies a condition for pruning, we return the layer wrapped in a pruning wrapper. Otherwise, we return the layer as is, untouched.

def cloning_func(layer):

    # is it a Dense layer?

    if isinstance(layer, keras.layers.Dense):

        return plm(layer)

    # does it have a certain name?

    if layer.name == 'dense5':

        return plm(layer)

    # if does not meet any conditions for pruning

    return layer

Listing 4-16

Defining a cloning function to map a layer to the desired state

Using this function, we can annotate the model by passing it as a cloning function when cloning the original model (Listing 4-17).

pruned_model = keras.models.clone_model(

    model,

    clone_function = cloning_func

)

Listing 4-17

Using the cloning function with Keras’ clone_model function

Then, compile and fit (with the Update Pruning Step callback) as usual.

Pruning in Theoretical Deep Learning: The Lottery Ticket Hypothesis

Pruning is an especially important method not only for the purposes of model compression but also in advancing theoretical understandings of deep learning. Jonathan Frankle and Michael Carbin’s 2019 paper “The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks”² builds upon the empirical success of pruning to formulate the Lottery Ticket Hypothesis, a theoretical hypothesis that reframes how we look at neural network knowledge representation and learning.

Pruning has demonstrated that the number of parameters in neural networks can be decreased by upward of 90% with little damage to performance metrics. However, a prerequisite of pruning is that pruning must be performed on a large model; a small, trained network the same size as a pruned network still will not perform as well as the pruned network. A key component of pruning, it has been observed, is the element of reduction; the knowledge must first be learned by a large model, which is then iteratively reduced into fewer parameters. One cannot begin with an architecture mimicking that of a pruned model and expect to yield results comparable to the pruned model. These findings are the empirical motivation for the Lottery Ticket Hypothesis.

The primary contribution of the Lottery Ticket Hypothesis is in explaining the role of weight initialization in neural network development: weights that begin with convenient initialization values are “picked” up by optimizers and “developed” into playing meaningful roles in the final, trained network. As the optimizer determines how to update certain weights, certain subnetworks within the neural network are delegated to carry most of the information flow simply because their initialized weights were the right values to spark growth. On the other hand, weights that begin with poor initialization values are dimmed down as inconveniences and superfluous weights; these are the “losing tickets” that are pruned away in pruning. Pruning reveals the architecture containing the “winning tickets.” Neural networks are running giant lotteries; the “winners” are amplified and the “losers” are attenuated.

You can think of a neural network from this perspective as a massive delivery package with a tiny valuable product inside and lots of stuffing. The vast majority of value is in a small minority of the actual package, but you need the package in the first place to find the product inside. Once you have the product, though, there’s no need for the box anymore. Correspondingly, given that the initialization values are key to a network’s success, you can retrain the pruned model architecture with the same corresponding initialization values and obtain similar performance to that of the original network.

This hypothesis reframes how we look at the training process of neural networks. The conventional perspective of machine learning models has always been that models begin with a set of “bad” parameters (the “initial guess”) that are iteratively improved by finding updates that make the largest decrease to the loss function. With the largeness and even possible “overparameterization” of modern neural networks, though, the Lottery Ticket Hypothesis hints at a new logic of understanding training: learning is primarily a process not only of improving but also of searching. Promising subnetworks are developed via an alternating pattern of searching for promising subnetworks and improving promising subnetworks to become more promising. This new perspective toward understanding parameter updates in the large context of modern deep learning may fuel further innovation in theoretical understandings and in practical developments. For instance, we understand now that the initialization of weights plays a key role in the success of a subnetwork, which may direct further research toward understanding how weight initialization operates with respect to trained subnetwork performance.

The Lottery Ticket Hypothesis and undoubtedly further theoretical advances in our understanding of neural networks guided by empirically observed phenomena in model compression will continue to serve as stepping stones to accelerating the improvement of our model-building methods.

Quantization Theory and Intuition

Traditionally , neural networks use 32 bits to represent parameters; while this is fine for training in modern deep learning environments that have the computing power to use such a precision, it is not feasible in applications that require lower storage and faster predictions. In quantization, parameters are reduced from their 32-bit representations to 8-bit integer representations, leading to a fourfold decrease in memory requirements.

In mathematics, quantization is the mapping of a continuous set of values to a smaller set of discrete values (Figure 4-6). In deep learning, quantization refers to a wide set of methods that can be used to reduce the precision of a parameter via a similar method. Generally, this is performed by separating values into information buckets. In binary quantization, values are quantized into two buckets; in ternary quantization, values are quantized into three buckets. However, binary and ternary quantization may be too extreme, which is why most deployed models employ a multiple-bit-to-multiple-bit quantization approach. How these bins are placed, how large each bin is, and other parameters to perform this mapping are dependent on which quantization strategy is being used.

(You could view magnitude-based pruning as a selective form of quantization, in which weights smaller in magnitude than a certain threshold are “quantized to 0” and other weights are binned to themselves.)

Recall the living space analogy for pruning. Instead of outright removing certain items from your living space, imagine reducing the cost to keep each one by a little bit. You decide to downgrade your television subscription to a lower tier, decrease the electricity consumption of your light, order takeout once a week instead of twice or thrice a week, and other amends that “round” the cost of your experience down.

Post-processing quantization is a quantization procedure performed on the model after it is trained. While this method achieves a good compression rate with the advantage of decreased latency, the errors in the small approximations performed in each of the weights via quantization accumulate and lead to a significant decrease in performance.

Like the iterative approach to pruning, quantization is generally not performed on the entire network at once – this is too jarring a change, just as pruning away 95% of a network’s parameters all at once is not conducive to recovery. Rather, after a model is pretrained – ideally, pretraining develops meaningful and robust representations that can be used to help the model recover from the compression – the model undergoes Quantization Aware Training, or QAT (Figure 4-7).

Throughout Quantization Aware Training, the model itself remains unquantized, representing all of its parameters with the standard 32 bits. However, a quantization error is introduced for consideration: in the network’s feed-forward stage, the output of the network is the same as if the network had been quantized. That is, before any prediction, the network undergoes “simulated quantization” – for the purposes of prediction, its parameters are quantized. This simulated quantization output is used to update the model parameters, which are still unquantized. Thus, while the model itself remains unquantized throughout Quantization Aware Training, it learns to develop parameters that will succeed when the model becomes quantized. The model is left unquantized because it is significantly easier to update model parameters with more precise parameters.

After Quantization Aware Training, the model is formally quantized – its parameters are binned and it uses the 8-bit integer representation (or some other representation, depending on the implementation). Because of Quantization Aware Training’s preparation, the model should have developed parameters that are robust and successful when quantized (Figure 4-8).

With quantization, a model’s storage requirements and latency can be dramatically decreased with little effect on performance.

Quantization Implementation

Like in pruning, you can quantize an entire model or quantize layers independently.

Weight Clustering Theory and Intuition

Weight clustering is a combination of pruning and quantization in character – it reduces the number of unique weights by slightly adjusting each weight value. Given a user-specified quantity of clusters n, the weight clustering algorithm assigns each weight value a cluster and sets the weight value to the centroid of that weight value (Figure 4-9).

Weights that are part of one cluster all share the same value, thus allowing for more efficient means of storage. Similarly to quantization, the decrease in storage requirement is a matter of precision; each parameter’s precise value can be replaced by the index of the associated centroid value. These precise values can be stored once in an indexable list of centroid values (Figure 4-10). (Note that even if this method of centroid indexing is not used, compression algorithms will be able to take advantage of repeated values.)

The key parameter in weight clustering is determining the number of clusters. Like the percent of parameters to prune, this is a trade-off between performance and compression. If the number of clusters is very high, the change each parameter makes from its original value to the value of the centroid it was assigned is very small, allowing for more precision and easier recovery from the compression operation. However, it diminishes the compression result by increasing storage requirements both for storing the centroid values and potentially the indexes themselves. On the other hand, if the number of clusters is too small, the performance of the model may be so impaired that it cannot recover – it may simply not be possible for a model to function reasonably with a certain set of fixed parameters.

Weight Clustering Implementation

Like with pruning and quantization, you can either cluster the weights of an entire model or of individual layers.

Sparsity Preserving Quantization

In sparsity preserving quantization , pruning is followed by quantization (Figure 4-15).

Using code and methods discussed earlier in the pruning section, obtain a pruned_model. You can use the measurement metrics defined prior to verify that the pruning procedure was successful. Use the strip_pruning function (tfmot.sparsity.keras.strip_pruning) to remove artifacts from the pruning procedure; this is necessary in order to perform quantization.

Recall that to induce Quantization Aware Training for a model, you used the quantize_model() function and then compiled and fitted the model. Performing pruning-preserving Quantization Aware Training, however, requires an additional step. The quantize_annotate_model() function is used not to actually quantize the model, but to provide annotations indicating that the entire model should be quantized. quantize_annotate_model() is used for more specific customizations of the quantization procedure, whereas quantize_model() can be thought of as the “default” quantization method. (You may similarly recall that quantize_annotate_layer() was used for another specific customization – layer-specific quantization.)

After the entire model has been annotated, we use the quantize_apply() function to actually quantize the annotated model. In this function, we can specify the preservation of another compression method – in this case, pruning. This is specified by passing a tfmot.experimental.combine object, which indicates a compression method to be preserved when “combining” or “collaborating.” The pruning-preserving Quantization Aware Training model can then be compiled and fitted as usual.

Cluster Preserving Quantization

In cluster preserving quantization , weight clustering is followed by quantization (Figure 4-16).

Using code and methods discussed earlier in the “Weight Clustering” section, obtain a clustered_model. From here, the process is almost the same as sparsity preserving quantization: after stripping clustering artifacts from the clustered_model, annotate the model and use quantize_apply to quantize the annotated layers. When specifying which compression method to preserve in quantize_apply, use Default8BitClusterPreserveQuantizeScheme rather than Default8BitPrunePreserveQuantizeScheme.

Sparsity Preserving Clustering

In sparsity preserving clustering , pruning is followed by weight clustering (Figure 4-17).

Sparsity preserving clustering follows a slightly different process than cluster preserving quantization and sparsity preserving quantization.

Using code and methods discussed earlier in the pruning section, obtain a pruned_model. Strip away pruning artifacts with strip_pruning.

We need to import the cluster_weights function to perform weight clustering; prior, we imported it from tfmot.clustering.keras.cluster_weights. However, to use sparsity preserving clustering, we need to import the function from a different place: from tensorflow_model_optimization.python.core.clustering.keras.experimental.cluster import cluster_weights.

Then, apply the cluster_weights function to the stripped pruned model with the clustering parameters , and compile and fit (Listing 4-32).

# create sparsity preserving clustering model

spc = cluster_weights(pruned_model, **clustering_params)

# compile and fit

spc.compile(...)

spc.fit(...)

Listing 4-32

Performing sparsity preserving clustering after pruning

Extreme Collaborative Optimization

The 2016 paper “Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization, and Huffman Coding”³ by Song Han, Huizi Mao, and William J. Dally was an instrumental leap forward in collaborative optimization.

The paper proposed a three-stage compression pipeline: pruning, weight clustering and quantization (grouped together as one method), and Huffman coding (Figure 4-18). This compression pipeline progressively compresses large models like AlexNet and VGG-16 on the ImageNet dataset between 35 times and 49 times without incurring any loss in accuracy. Moreover, the latency is decreased by three to four times with three to seven times improved energy efficiency. By chaining compression methods in this order, the compression methods minimally interfered with one another, leading to surprisingly large compression:

1. 1.

   Recall that pruning is best performed as an iterative process in which connections are pruned and the network is fine-tuned on those pruned connections. In this paper, pruning reduces the model size by 9 to 13 times with no decrease in accuracy.

    
2. 2.

   Recall that weight clustering is performed by clustering weights that have similar values and setting the weights to their respective centroid values and that quantization is performed by training the model to adapt toward lower-precision weights. Weight clustering combined with quantization, after pruning, reduces the original model size by 27 to 31 times.

    
3. 3.

   Huffman coding is a compression technique that was proposed by computer scientist David A. Huffman in 1952. It allows for lossless data compression that represents more common symbols with fewer bits. Huffman coding is different from the previously discussed model compression methods because it is a post-training compression scheme; that is, there is no model fine-tuning required for the scheme to work successfully. Huffman encoding allows for even further compression – the final model is compressed with a reduction 35 to 49 times its original size.

    

Han, Mao, and Dally provide important insights into the dynamics of collaborative optimization. Pruning before quantization doesn’t hurt quantization, for instance – the performance of a model both pruned and quantized is almost identical to that of a model that has only undergone quantization (of course, the pruned and quantized model has fewer parameters) (Figure 4-19). This demonstrates a key property of ideal collaborative optimization: strength is found in a diverse array of compression attacks. By chaining a diverse set of compression methods that each attack different representation redundancies, the model is stripped of inefficient representations from all “angles” and therefore results in higher compression while still maintaining the necessary essential facilities for good performance.

Rethinking Quantization for Deeper Compression

Recall that when performing quantization, Quantization Aware Training is used to orient the model toward learning weights that are robust to quantization. This is performed by simulating a quantized environment when the model is making predictions.

However, Quantization Aware Training raises a key problem: because quantization effectively “discretizes” or “bins” a functionally “continuous” weight value, the derivative with respect to the input is zero at almost everywhere, posing problems for the gradient update calculations. In order to work around this, in practice a Straight Through Estimator is used. As implied by its name, a Straight Through Estimator estimates the output gradients of the discretized layer as its input gradients without regard for the actual derivative of the actual discretized layer. A Straight Through Estimator works with relatively less aggressive quantization (i.e., the 8-bit integer quantization implemented earlier) but fails to provide sufficient estimation for more severe compressions (e.g., 4-bit integer).

To address this problem, Angela Fan and Pierre Stock, along with Benjamin Graham, Edouard Grave, Remi Gribonval, Herve Jegou, and Armand Joulin, propose quantization noise in their paper “Training with Quantization Noise for Extreme Model Compression,”⁴ a novel approach toward orienting a compressed model to developing quantization robust weights.

Rather than simulating the quantization of the entire model, like in Quantization Aware Training, quantization noise instead simulated the quantization of part of a model – a randomly selected subset of weights are simulated quantized during each forward pass (Figure 4-20). This means that most of the weights are updated with cleaner gradients.

Responsible Compression: What Do Compressed Models Forget?

When we talk about compression, the two key figures we consider are the model performance and the compression factor. These two figures are often balanced and used to determine the success of a model compression operation. We often see an increase in compression accompanied by a decrease in performance – but do you wonder what types of data inputs are being sacrificed by the compression procedure? What is lying beneath the generalized performance metrics?

In “What Do Compressed Deep Neural Networks Forget?”,⁵ Sara Hooker, along with Aaron Courville, Gregory Clark, Yann Dauphin, and Andrea Frome, investigates just this question: how do compression methods affect what knowledge a compressed model is forced to “forget” by the compression? Hooker et al.’s findings suggest that looking merely at standard performance metrics like test set accuracy may not be enough to reveal the impact of compression on the model’s true generalization capabilities.

Pruning Identified Exemplars (PIEs) are defined inputs to a model in which there is a high level of disagreement between the predictions of pruned and unpruned models. Hooker et al. find that general metrics like test set accuracy hide important information regarding the effect of pruning on the model’s generalization capabilities; model compression methods like pruning do not uniformly affect the model’s data to process instances across the distribution of the dataset. Rather, a small subset of data is disproportionately impacted (Figure 4-21).

Data instances from the long tail of the dataset distribution – that is, less commonly represented or more complex data instances – are more often “sacrificed” in model compression. Hooker at al. asked human subjects to label components of Pruning Identified Exemplars and found that PIEs were more difficult both for humans and models to classify; PIEs are generally more complex, consisting of multiple objects, lower quality, or ambiguity (Figure 4-22). Pruning forces compressed models to sacrifice understanding of these particular instances, exposing vulnerabilities in compressed model generalization.

Compressed models, moreover, are more prone to small changes that humans would be robust to. The higher the compression, the less robust the model becomes to variations like changing brightness, contrast, blurring, zooming, and JPEG noise. This also increases the vulnerability of compressed models used in deployment to adversarial attacks, or attacks designed to subvert the model’s output by making small, cumulative changes undetectable to humans (see Chapter 2, Case Study 1, on adversarial attacks exploiting properties of transfer learning).

In addition to posing concerns for model robustness and security, these findings raise questions for the role of model compression in increasing discussion on fairness. Given that model compression disproportionately affects the model’s capacity to process less represented items in categories, model compression can amplify existing disparities in dataset representation.

Hooker et al.’s work reminds us that neural networks are complex entities that often may require more exploration and consideration than wide-reaching metrics may suggest and leaves important questions to be answered in future work on model compression.