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

1 INTRODUCTION . Pre-trained feature extractors , such as BERT ( Devlin et al. , 2018 ) for natural language processing and VGG ( Simonyan & 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 CO2 ( 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 task-generic 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 overpruned 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 . 2 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.1 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.2 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 & 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.3 2.1 MAGNITUDE WEIGHT PRUNING . In this work , we focus on weight magnitude pruning because it is one of the most fine-grained and effective pruning methods . It also has a compelling saliency criterion ( Han et al. , 2015 ) : if a weight is close to zero , then its input is effectively ignored , which means the weight can be pruned . 1If weights are pruned , however , the weight matrices become sparse . Sparse matrix multiplication is difficult to optimize on current GPU architectures ( Han et al. , 2016 ) , although progress is being made . 2Interestingly , recent work used compression not to induce simplicity but to measure it ( Arora et al. , 2018 ) . 3Sparse networks are difficult to train from scratch ( Evci et al. , 2019 ) . However , Dettmers & Zettlemoyer ( 2019 ) and Mostafa & Wang ( 2019 ) present methods to do this by allowing SGD to search over the space of possible subnetworks . Our findings suggest that these methods might be used to train sparse BERT from scratch . Magnitude weight pruning itself is a simple procedure : 1 . Pick a target percentage of weights to be pruned , say 50 % . 2 . Calculate a threshold such that 50 % of weight magnitudes are under that threshold . 3 . Remove those weights . 4 . Continue training the network to recover any lost accuracy . 5 . Optionally , return to step 1 and increase the percentage of weights pruned . This procedure is conveniently implemented in a Tensorflow ( Abadi et al. , 2016 ) package4 , which we use ( Zhu & Gupta , 2017 ) . Calculating a threshold and pruning can be done for all network parameters holistically ( global pruning ) or for each weight matrix individually ( matrix-local pruning ) . Both methods will prune to the same sparsity , but in global pruning the sparsity might be unevenly distributed across weight matrices . We use matrix-local pruning because it is more popular in the community.5 For information on other pruning techniques , we recommend Gale et al . ( 2019 ) and Liu et al . ( 2019b ) . 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 ) . 3.1 IMPLEMENTING BERT PRUNING . BERT-Base consists of 12 encoder layers , each of which contains 6 prunable matrices : 4 for the multi-headed self-attention and 2 for the layer ’ s output feed-forward network . Recall that self-attention first projects layer inputs into key , query , and value embeddings via linear projections . While there is a separate key , query , and value projection matrix for each attention head , implementations typically stack matrices from each attention head , resulting in only 3 parameter matrices : one for key projections , one for value projections , and one for query projections . We prune each of these matrices separately , calculating a threshold for each . We also prune the linear output projection , which combines outputs from each attention head into a single embedding.6 We prune word embeddings in the same way we prune feed-foward networks and self-attention parameters.7 The justification is similar : if a word embedding value is close to zero , we can assume it ’ s zero and store the rest in a sparse matrix . This is useful because token / subword embeddings tend to account for a large portion of a natural language model ’ s memory . In BERT-Base specifically , the embeddings account for ∼21 % of the model ’ s memory . Our experimental code for pruning BERT , based on the public BERT repository , is available here.8 3.2 PRUNING DURING PRE-TRAINING . We perform weight magnitude pruning on a pre-trained BERT-Base model.9 We select sparsities from 0 % to 90 % in increments of 10 % and gradually prune BERT to this sparsity over the first 10k steps of training . We continue pre-training on English Wikipedia and BookCorpus for another 90k steps to regain any lost accuracy.10 The resulting pre-training losses are shown in Table 1 . We then fine-tune these pruned models on tasks from the General Language Understanding Evaluation ( GLUE ) benchmark , which is a standard set of 9 tasks that include sentiment analysis , natural 4https : //www.tensorflow.org/api docs/python/tf/contrib/model pruning 5The weights in almost every matrix in BERT-Base are approximately normally distributed with mean 0 and variance between 0.03 and 0.05 ( Table A ) . This similarity may imply that global pruning would perform similarly to matrix-local pruning . 6We could have calculated a single threshold for the entire self-attention layer or for each attention head separately . Similar to global pruning vs. matrix-local pruning , it ’ s not clear which one should be preferred . 7Interestingly , pruning word embeddings is slightly more interpretable that pruning other matrices . See Figure 9 for a heatmap of embedding magnitudes , which shows that shorter subwords tend to be pruned more than longer subwords and that certain dimensions are almost never pruned in any subword . 8URL omitted for anonymity 9https : //github.com/google-research/bert 10 Evaluation curves leveled out at 20k steps . language inference , etc . We avoid WNLI , which is known to be problematic.11 We also avoid tasks with less than 5k training examples because the results tend to be noisy ( RTE , MRPC , STS-B ) . We fine-tune a separate model on each of the remaining 5 GLUE tasks for 3 epochs and try 4 learning rates : [ 2 , 3 , 4 , 5 ] × 10−5 . The best evaluation accuracies are averaged and plotted in Figure 1 . Individual task results are in Table 1 . BERT can be used as a static feature-extractor or as a pre-trained model which is fine-tuned end-toend . In all experiments , we fine-tune weights in all layers of BERT on downstream tasks .

Summary: This work is an empirical study of testing how pruning at the pre-training stage affects subsequent transfer learning (through fine-tuning) stage. The main idea is to carefully control the amount of sparsity injected into BERT through weight magnitude pruning and study the impact on accuracy. The experimental setup is mostly well done, especially the part that disentangles the complexity restriction and information deletion. During the exploration, the authors made several interesting observations, such as 30-40% model weights do not encode any useful inductive bias, which could help shed some light for future work on both training and compressing BERT-like models.