Compressing BERT : Studying the Effects of Weight Pruning on Transfer Learning (2020)

Abstract: Pre-trained universal feature extractors, such as BERT for natural language processing and VGG for computer vision, have become effective methods for improving deep learning models without requiring more labeled data. While effective, feature extractors like BERT may be prohibitively large for some deployment scenarios. We explore weight pruning for BERT and ask: how does compression during pre-training affect transfer learning? We find that pruning affects transfer learning in three broad regimes. Low levels of pruning (30-40%) do not affect pre-training loss or transfer to downstream tasks at all. Medium levels of pruning increase the pre-training loss and prevent useful pre-training information from being transferred to downstream tasks. High levels of pruning additionally prevent models from fitting downstream datasets, leading to further degradation. Finally, we observe that fine-tuning BERT on a specific task does not improve its prunability. We conclude that BERT can be pruned once during pre-training rather than separately for each task without affecting performance.

## introduction pre-trained feature extractors, such as bert (devlin et al., 2018) for natural language processing and vgg (simonyan and zisserman, 2014) for computer vision, have become effective methods for improving the performance of deep learning models. in the last year, models similar to bert have become state-of-the-art in many nlp tasks, including natural language inference (nli), named entity recognition (ner), sentiment analysis, etc. these models follow a pre-training paradigm: they are trained on a large amount of unlabeled text via a task that resembles language modeling (yang et al., 2019; chan et al., 2019) and are then fine-tuned on a smaller amount of "downstream" data, which is labeled for a specific task. pre-trained models usually achieve higher accuracy than any model trained on downstream data alone. the pre-training paradigm, while effective, still has some problems. while some claim that language model pre-training is a "universal language learning task" (radford et al., 2019), there is no theoretical justification for this, only empirical evidence. second, due to the size of the pre-training dataset, bert models tend to be slow and require impractically large amounts of gpu memory. bert-large can only be used with access to a google tpu, and bert-base requires some optimization tricks such as gradient checkpointing or gradient accumulation to be trained effectively on consumer hardware (sohoni et al., 2019). training bert-base from scratch costs ∼$7k and emits ∼1438 pounds of co 2 (strubell et al., 2019). model compression (bucila et al., 2006), which attempts to shrink a model without losing accuracy, is a viable approach to decreasing gpu usage. it might also be used to trade accuracy for memory in some low-resource cases, such as deploying to smartphones for real-time prediction. the main questions this paper attempts to answer are: does compressing bert impede it's ability to transfer to new tasks? and does fine-tuning make bert more or less compressible? to explore these questions, we compressed english bert using magnitude weight pruning (han et al., 2015) and observed the results on transfer learning to the general language understanding evaluation (glue) benchmark (wang et al., 2019), a diverse set of natural language understanding tasks including sentiment analysis, nli, and textual similarity evaluation. we chose magnitude weight pruning, which compresses models by removing weights close to 0, because it is one of the most fine-grained and effective compression methods and because there are many interesting ways to view pruning, which we explore in the next section. our findings are as follows: low levels of pruning (30-40%) do not increase pre-training loss or affect transfer to downstream tasks at all. medium levels of pruning increase the pre-training loss and prevent useful pre-training information from being transferred to downstream tasks. this information is not equally useful to each task; tasks degrade linearly with pre-train loss, but at different rates. high levels of pruning, depending on the size of the downstream dataset, may additionally degrade performance by preventing models from fitting downstream datasets. finally, we observe that fine-tuning bert on a specific task does not improve its prunability or change the order of pruning by a meaningful amount. to our knowledge, prior work had not shown whether bert could be compressed in a taskgeneric way, keeping the benefits of pre-training while avoiding costly experimentation associated with compressing and re-training bert multiple times. nor had it shown whether bert could be over-pruned for a memory / accuracy trade-off for deployment to low-resource devices. in this work, we conclude that bert can be pruned prior to distribution without affecting it's universality, and that bert may be over-pruned during pre-training for a reasonable accuracy trade-off for certain tasks. ## pruning: compression, regularization, architecture search neural network pruning involves examining a trained network and removing parts deemed to be unnecessary by some heuristic saliency criterion. one might remove weights, neurons, layers, channels, attention heads, etc. depending on which heuristic is used. below, we describe three different lenses through which we might interpret pruning. compression pruning a neural network decreases the number of parameters required to specify the model, which decreases the disk space required to store it. this allows large models to be deployed on edge computing devices like smartphones. pruning can also increase inference speed if whole neurons or convolutional channels are pruned, which reduces gpu usage. @xcite regularization pruning a neural network also regularizes it. we might consider pruning to be a form of permanent dropout (molchanov et al., 2017) or a heuristic-based l0 regularizer (louizos et al., 2018). through this lens, pruning decreases the complexity of the network and therefore narrows the range of possible functions it can express. @xcite the main difference between l0 or l1 regularization and weight pruning is that the former induce sparsity via a penalty on the loss function, which is learned during gradient descent via stochastic relaxation. it's not clear which approach is more principled or preferred. (gale et al., 2019) sparse architecture search finally, we can view neural network pruning as a type of sparse architecture search. liu et al. (2019b) and frankle and carbin (2019) show that they can train carefully re-initialized pruned architectures to similar performance levels as dense networks. under this lens, stochastic gradient descent (sgd) induces network sparsity, and pruning simply makes that sparsity explicit. these sparse architectures, along with the appropriate initializations, are sometimes referred to as "lottery tickets." foot_3 ## experimental setup bert is a large transformer encoder; for background, we refer readers to vaswani et al. (2017) or one of these excellent tutorials (alammar, 2018; klein et al., 2017). ## conclusion and future work we've shown that encoding bert's inductive bias requires many more weights than are required to fit downstream data. future work on compressing pre-trained models should focus on maintaining that inductive bias and quantifying its relevance to various tasks during accuracy/memory trade-offs. for magnitude weight pruning, we've shown that 30-40% of the weights do not encode any useful inductive bias and can be discarded without affecting bert's universality. the relevance of the rest of the weights vary from task to task, and fine-tuning on downstream tasks does not change the nature of this trade-off by changing which weights are pruned. in future work, we will investigate the factors that influence language modeling's relevance to downstream tasks and how to improve compression in a task-general way. figure 6 : we show how weight sort order movements are distributed during fine-tuning, given a weight's starting magnitude. we see that higher magnitude weights are more stable than lower magnitude weights and do not move as much in the sort order. this plot is nearly identical for every model and learning rate, so we only show it once. figure 7 : a heatmap of the weight magnitudes of the 12 horizontally stacked self-attention key projection matrices for layer 1. a banding pattern can be seen: the highest values of the matrix tend to cluster in certain attention heads. this pattern appears in most of the self-attention parameter matrices, but it does not cause pruning to prune one head more than another. however, it may prove to be a useful heuristic for attention head pruning, which would not require making many passes over the training data.