R EWEIGHTED P ROXIMAL P RUNING FOR L ARGE -S CALE L ANGUAGE R EPRESENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Recently, pre-trained language representation ﬂourishes as the mainstay of the natural language understanding community, e.g., BERT.
These pre-trained language representations can create state-of-the-art results on a wide range of downstream tasks.
Along with continuous signiﬁcant performance improvement, the size and complexity of these pre-trained neural models continue to increase rapidly.
Is it possible to compress these large-scale language representation models?
How will the pruned language representation affect the downstream multi-task transfer learning objectives?
In this paper, we propose Reweighted Proximal Pruning (RPP), a new pruning method speciﬁcally designed for a large-scale language representation model.
Through experiments on SQuAD and the GLUE benchmark suite, we show that proximal pruned BERT keeps high accuracy for both the pretraining task and the downstream multiple ﬁne-tuning tasks at high prune ratio.
RPP provides a new perspective to help us analyze what large-scale language representation might learn.
Additionally, RPP makes it possible to deploy a large state-of-the-art language representation model such as BERT on a series of distinct devices (e.g., online servers, mobile phones, and edge devices).
I NTRODUCTION
Pre-trained language representations such as GPT (Radford et al., 2018), BERT (Devlin et al., 2019) and XLNet (Yang et al., 2019), have shown substantial performance improvements using self-supervised training on large-scale corpora (Dai & Le, 2015; Peters et al., 2018;
Radford et al., 2018; Liu et al., 2019a).
More interestingly, the pre-trained BERT model can be ﬁne-tuned with just one additional output layer to create state-of-the-art models for a wide range of tasks, such as question answering (Rajpurkar et al., 2016; 2018), and language inference (Bowman et al., 2015;
Williams et al., 2017), without substantial task-speciﬁc architecture modiﬁcations.
BERT is conceptually simple and empirically powerful (Devlin et al., 2019).
However, along with the signiﬁcant performance enhancement, the parameter volume and complexity of these pre-trained language representations signiﬁcantly increase.
As a result, it becomes difﬁcult to deploy these large-scale language representations into real-life computation constrained devices including mobile phones and edge devices.
Throughout this paper, we attempt to answer the following questions.
Question 1 : Is it possible to compress large-scale language representations such as BERT via weight pruning?
Question 2 : How would the weight-pruned, pre-trained model affect the performance of the downstream multi-task transfer learning objectives?
The problem of weight pruning has been studied under many types of deep neural networks (DNNs)
(Goodfellow et al., 2016), such as AlexNet (Krizhevsky et al., 2012), VGG (Simonyan & Zisserman, 2014), ResNet (He et al., 2016), and MobileNet (Howard et al., 2017).
It is shown that weight pruning can result in a notable reduction in the model size.
A suite of weight pruning techniques have been developed, such as non-structured weight pruning (Han et al., 2015), structured weight pruning (Wen et al., 2016), ﬁlter pruning (Li et al., 2016), channel pruning (He et al., 2017),
ADMM-NN (Ren et al., 2019) and PCONV (Ma et al., 2019) to name a few.
Different from pruning CNNtype models, BERT not only considers the metrics on the pre-training task, but also needs to make allowance for the downstream multi-task transfer learning objectives.
Thus, the desired weight pruning needs to preserve the capacity of transfer learning from a sparse pre-trained model to downstream ﬁne-tuning tasks.
In this work, we investigate irregular weight pruning techniques on the BERT model, including the iterative pruning method (Han et al., 2015) and one-shot pruning method (Liu et al., 2019b).
However, these methods fail to converge to a sparse pre-trained model without incurring signiﬁcant accuracy drop, or in many cases do not converge at all (see supporting results in Appendix).
Note that the aforementioned weight pruning techniques are built on different sparsity-promoting regularization schemes (Han et al., 2015; Wen et al., 2016), e.g., lasso regression ( (cid:96) 1 regularization) and ridge regression ( (cid:96) 2 regularization).
We ﬁnd that the failure of previous methods on weight pruning of
BERT is possibly due to the inaccurate sparse pattern learnt from the simple (cid:96) 1 or (cid:96) 2 based sparsitypromoting regularizer.
In fact, the difﬁculty of applying regularization to generate weight sparsity coincides with the observation in (Loshchilov & Hutter, 2018) on the imcompatibility of conventional weight decay ( (cid:96) 2 regularization) for training super-deep DNNs as BERT.
It is pointed out that the main reason is that the direct optimization of a regularization penalty term causes divergence from the original loss function and has negative effect on the effectiveness of gradient-based update.
To mitigate this limitation, (Loshchilov & Hutter, 2018) have modiﬁed the regularization in Adam by decoupling weight decay regularization from the gradient-based update , and have achieved stateof-the-art results on large-scale language pre-training and downstream multi-task transfer learning objectives (Devlin et al., 2019).
[Figure]
-Tuning
Unlabeled Sentence A and B Pair
Unlabeled Sentence A and B Pair
MRPC MNLI CoLA Mask LM NSP
Mask LM
SQuAD
Masked Sentence A
Masked Sentence B [CLS] [SEP] ... Tok
TokN
Tok1 ...
TokN [CLS] [SEP] ... Tok1 TokN
Tok1 ...
TokN [CLS] [SEP] ... Tok1 TokN
Tok1 ...
TokN
Start/End Span
Masked Sentence A
Masked Sentence B
Identical Universal
Sparsity
Question Answer Pair
Question
Paragraph
T
T N ...
T [SEP] T 1` T M` ...
E 1 E N ... E
[SEP] E 1` E M` ... E [CLS]
E 1 E N ... E
[SEP] E 1` E M` ... E [CLS]
E 1 E N ... E
[SEP] E 1` E M` ... E [CLS]
T
T N ...
T [SEP] T 1` T M` ...
T
T N ...
T [SEP] T 1` T M` ...
In this work, we aim at more accurate universal sparse pattern search (see Figure 1 for an overview of our approach) motivated by our experiments and the conclusion from Loshchilov & Hutter (2018).
We propose Reweighted Proximal Pruning (RPP), which integrates reweighted (cid:96) 1 minimization (Candes et al., 2008) with proximal algorithm (Parikh et al., 2014).
RPP consists of two parts : the reweighted (cid:96) 1 minimization and the proximal operator.
Reweighted (cid:96) 1 minimization serves as a better method of generating sparsity in DNN models matching the nature of weight pruning, compared with (cid:96) 1 regularization.
Thanks to the closed-form solution of proximal operation on a weighted (cid:96) 1 norm, in RPP the sparsity pattern search can be decoupled from computing the gradient of the training loss.
In this way the aforementioned pitfall in prior weight pruning technique on BERT can be avoided.
We show that RPP achieves effective weight pruning on BERT for the ﬁrst time to the best of our knowledge.
Experimental results demonstrate that the proximal pruned BERT model keeps high accuracy on a wide range of downstream tasks, including SQuAD (Rajpurkar et al., 2016; 2018) and GLUE (Wang et al., 2018).
We summarize our contributions as follows.
• We develop the pruning algorithm Reweighted Proximal Pruning (RPP), which acheives the ﬁrst effective weight pruning result on large pre-trained language representation model - BERT.
RPP achieves 59 .
3% weight sparsity without inducing the performance loss on both pre-training and ﬁne-tuning tasks.
• We spotlight the relationship between the pruning ratio of the pre-trained DNN model and the performance on the downstream multi-task transfer learning objectives.
We show that many downstream tasks except for SQuAD allows at least 80% pruning ratio compared with 59 .
3% under the more challenging task SQuAD.
• We observe that as the pruning ratio of the pre-trained language model increases, the performance on the downstream transfer learning tasks decreases.
The descending range varies in different downstream transfer learning tasks.
However, the proposed RPP approach is able to achieve a consistently high pruning ratio compared to iterative pruning based methods.
• We show that different from weight pruning in image classiﬁcation tasks, RPP helps to ﬁnd the structured sparsity pattern in transformer blocks used in BERT.
Moreover, we peer into the effect of network pruning on the language representation embedded in BERT.
R ELATED W ORK
BERT and prior work on model compression BERT (Devlin et al., 2019) is a self-supervised approach for pre-training a deep transformer encoder (Vaswani et al., 2017), before ﬁne-tuning it for particular downstream tasks.
Pre-training of BERT optimizes two training objectives − masked language modeling (MLM) and next sentence prediction (NSP) − which require a large collection of unlabeled text.
We use BooksCorpus (800M words) (Zhu et al., 2015) and the English instance of
Wikipedia (2,500M words) as the pre-training corpus, the same as Devlin et al. (2019).
For detailed information about the BERT model, readers can refer to the original paper (Devlin et al., 2019).
Michel et al.
(2019) mask some heads in multi-head attention modules in BERT, and then evaluate the performance on the machine translation task.
Similarly, Hao et al.
(2019) eliminates certain heads in the multi-head attention module.
First, the limited previous work do not consider the pretraining metrics and the other downstream multi-mask transfer learning objectives.
They only considered the speciﬁc machine translation task (out of over 10 transfer tasks), which is only a speciﬁc ﬁne-tuning and is limited for the universal pre-trained language representation (BERT).
Second, the multi-head attention module uses a weight sharing mechanism (Vaswani et al., 2017).
So masking some heads does not reduce the weight volume.
Finally, multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions, while single attention head inhibits this effect (Vaswani et al., 2017).
As a result, masking some heads in multi-head attention harms the weight sharing mechanism, without weight volume reduction.
In summary, the limited previous work in this area are not effective weight pruning method on BERT.
Shen et al.
(2019) reports the quantization result of BERT model, which is orthogonal to our work and can be combined for further compression/acceleration.
Reweighted (cid:96) 1 and proximal algorithm
Candes et al.
(2008) present reweighted (cid:96) 1 algorithm and demonstrate the remarkable performance and broad applicability in the areas of statistical estimation, error correction and image processing.
Proximal algorithms can be viewed as an analogous tool for non-smooth, constrained, large-scale, or distributed versions of these problems (Parikh et al., 2014).
To the best of our knowledge, ours is the ﬁrst work that applies reweighted (cid:96) 1 minimization to network compression, particularly for BERT pruning. ,
R EWEIGHTED P ROXIMAL P RUNING FOR LARGE - SCALE LANGUAGE REPRESENTATION DURING PRE - TRAINING
Pruning for pre-trained language representations should not only consider the performance of pretraining objectives, but also make allowance for the downstream ﬁne-tuning transfer learning tasks.
Let f i denote the loss function of network for downstream task T i ∼ p ( T ) , where p ( T ) denotes the distribution of tasks.
Let w denote the parameters of the pre-trained model (pre-training in BERT), and z i denote the i -th
task-speciﬁed model parameters (ﬁne-tuning in BERT).
The downstream tasks have separate ﬁne-tuned models, even though they are initialized with the same pre-trained parameters (Devlin et al., 2019).
Starting from the pre-trained parameters w , the parameters z i ( w ) are obtained through ﬁne-tuning minimize w ∈
[Equation]
3. P RUNING FORMULATION IN TRANSFER LEARNING
Following the conventional weight pruning formulation, we ﬁrst consider the problem of weight pruning during pre-training:
[Equation]
where f 0 is the loss function of pruning, p ∈ { 0 , 1 } denotes the type of regularization norm, and γ is a regularization term.
We note that the sparsity-promoting regularizer in the objective could also be replaced with a hard (cid:96) p constraint, | w (cid:107) p ≤ τ for some τ .
Let ˆ w denote the solution to problem (2), and the corresponding sparse pattern S ˆ w is given by
[Equation]
For a speciﬁc transfer task i , we allow an additional retraining/ﬁne-tuning step to train/ﬁne-tune weights starting from the pre-training results ˆ w and subject to the determined, ﬁxed sparse pattern
S ˆ w , denoted as z i ( ˆ w ; S ˆ w ) .
That is, we solve the modiﬁed problem equation
[Equation]
Here, different from (1), the task-speciﬁc ﬁne tuning weights variable
z i ( ˆ w ; S ˆ w ) is now deﬁned over S ˆ w .
Our goal is to seek a sparse (weight pruned) model during pre-training, with weight collection
ˆ w and sparsity S ˆ w , which can perform as well as the original pre-trained model over multiple new tasks (indexed by i ).
These ﬁne-tuned models z i ( ˆ w ; S ˆ w ) (for different i ) share the identical universal sparsity S ˆ w .
3. R EWEIGHTED P ROXIMAL P RUNING
In order to enhance the performance of pruning pre-trained language representation over multi-task downstream transfer learning objectives, we propose Reweighted Proximal Pruning (RPP).
RPP consists of two parts: the reweighted (cid:96) 1 minimization and the proximal operator.
Reweighted (cid:96)minimization serves as a better method of generating sparsity in DNN models matching the natural objective of weight pruning, compared with (cid:96) 1 regularization.
The proximal algorithm then separates the computation of gradient with the proximal operation over a weighted (cid:96) 1 norm, without directly optimizing the entire sparsity-penalized loss, which requires gradient backpropagation of the involved loss.
This is necessary in the weight pruning of super-deep language representation models.
3.2. R EWEIGHTED (cid:96) 1 MINIMIZATION
In the previous pruning methods (Han et al., 2015; Wen et al., 2016), (cid:96) 1 regularization is used to generate sparsity.
However, consider that two weights w i , w j ( w i < w j ) in the DNN model are penalized through (cid:96) 1 regularization.
The larger weight w j is penalized more heavily than smaller weight w i in (cid:96) 1 regularization, which violates the original intention of weight pruning, “removing the unimportant connections” (parameters close to zero) (Han et al., 2015).
To address this imbalance, we introduce reweighted (cid:96) 1 minimization (Candes et al., 2008) to the DNN pruning domain.
Our introduced reweighted (cid:96) 1 minimization operates in a systematic and iterative manner (detailed process shown in Algorithm 1), and the ﬁrst iteration of reweighted (cid:96) 1 minimization is (cid:96) 1 regularization.
This designed mechanism helps us to observe the performance difference between (cid:96) 1 and reweighted (cid:96) 1 minimization.
Meanwhile, this mechanism ensures the advancement of reweighted (cid:96)minimization over (cid:96) 1 regularization, as the latter is the single, ﬁrst step of the former.
Consider the regularized weight pruning problem (reweighted (cid:96) 1 minimization):
[Equation]
where α i ( α i > 0) factor is a positive value.
It is utilized for balancing the penalty, and is different from weight w i in DNN model.
α i factors will be updated in the iterative reweighted (cid:96) 1 minimization procedure (Step 2 in Algorithm 1) in a systematic way (Candes et al., 2008).
If we set T =for reweighted (cid:96) 1 , then it reduces to (cid:96) 1 sparse training.
3.2. P ROXIMAL METHOD
In the previous pruning methods (Han et al., 2015; Wen et al., 2016), (cid:96) 1 regularization loss is directly optimized through the back-propagation based gradient update of DNN models, and the hardthreshold is adopted to execute the pruning action at the step of pruning (all weights below the hard-threshold become zero).
In our approach, we derive an effective solution to problem (5) for given { α i } , namely, in Step 1 of Algorithm 2, in which back-propagation based gradient update is only applied on f 0 ( w ) but not γ (cid:80) i
α i
| w i | .
We adopt the proximal algorithm (Parikh et al., 2014) to satisfy this requirement through decoupling methodology.
In this way, the sparsity pattern search can be decoupled from back-propagation based gradient update of the training loss.
The proximal algorithm is shown in (Parikh et al., 2014) to be highly effective (compared with the original solution) on a wide set of non-convex optimization problems.
Additionally, our presented reweighted (cid:96) 1 minimization (5) has analytical solution through the proximal operator.
To solve problem (5) for a given α , the proximal algorithm operates in an iterative manner:
[Equation]
where the subscript k denotes the time step of the training process inside each iteration of RPP, λ k ( λ k > 0) is the learning rate, and we set the initial w to be w ( t − 1) from the last iteration of reweighted (cid:96) 1 .
The proximal operator prox λ k ,rw − (cid:96) 1 ( a ) is the solution to the problem
[Equation]
where a = w k − 1 − λ k ∇ w f ( w k − 1 ) .
The above problem has the following analytical solution (Liu et al., 2014)
[Equation]
We remark that the updating rule (6) can be interpreted as the proximal step (8) over the gradient descent step w k − 1 − λ k ∇ w f ( w k − 1 ) .
Such a descent can also be obtained through optimizers such as AdamW.
We use the AdamW (Loshchilov & Hutter, 2018) as our optimizer, the same with (Devlin et al., 2019).
The concrete process of AdamW with proximal operator is shown in Algorithm 3 of Appendix C.
Why chooses AdamW rather than Adam?
Loshchilov & Hutter (2018) proposes AdamW to improve the generalization ability of Adam (Kingma & Ba, 2014).
Loshchilov & Hutter (2018) shows that the conventional weight decay is inherently not effective in Adam and has negative effect on the effectiveness of gradient-based update, which is the reason of the difﬁculty to apply adaptive gradient algorithms to super-deep DNN training for NLU applications (like BERT).
Loshchilov & Hutter (2018) mitigates this limitation and improves regularization of Adam, by decoupling weight decay regularization from the gradient-based update (Loshchilov & Hutter, 2018).
AdamW is widely adopted in pre-training large language representations, e.g., BERT (Devlin et al., 2019), GPT (Radford et al., 2018) and XLNet (Yang et al., 2019).
Our proposed RPP also beneﬁts from the decoupling design ideology.
The difference is that RPP is for the generation of sparsity, instead of avoiding over-ﬁtting, like decoupled weight decay in AdamW.
Our new and working baseline: New Iterative Pruning (NIP).
To get the identical universal sparsity S w , we tried a series of pruning techniques, including the iterative pruning method (Han et al., 2015) and one-shot pruning method (Liu et al., 2019b).
But these methods do not converge to a viable solution.
The possible reason for non-convergence of the iterative pruning method is that directly optimizing (cid:96) p ( p ∈ { 1 , 2 } ) sparsity-promoting regularization makes the gradient computation involved and thus harms the loss convergence (We provide the loss curve and analysis in Appendix
D).
To circumvent the convergence issue of conventional iterative pruning methods, we propose a new iterative pruning (NIP) method.
Different from iterative pruning (Han et al., 2015), NIP reﬂects the naturally progressive pruning performance without any externally introduced penalty.
We hope that other pruning methods should not perform worse than NIP, otherwise, the effect of optimizing the newly introduced sparsity-promoting regularization is negative.
We will show that NIP is able to successfully prune BERT to certain pruning ratios.
We refer readers to Appendix A for the full detail about NIP, our proposed baseline algorithm.
E XPERIMENTS
In this section, we describe the experiments on pruning pre-trained BERT and demonstrate the performance on 10 downstream transfer learning tasks.
4. E XPERIMENT S ETUP
We use the ofﬁcial BERT model from Google as the startpoint.
Following the notation from Devlin et al.
(2019), we denote the number of layers (i.e., transformer blocks) as L , the hidden size as H , and the number of self-attention heads as A .
We prune two kinds of BERT model: BERT BASE ( L = 12 , H = 768 , A
= 12 , total parameters = 110M ) and BERT LARGE ( L = 24 , H = 1024 ,
A = 16 , total parameters = 340M ).
As the parameters of these transformer blocks take up more than 97% weights of the entire BERT, the weights of these transformer blocks are our pruning target.
Data: In pre-training, we use the same pre-training corpora as Devlin et al.
(2019): BookCorpus ( 800M words) (Zhu et al., 2015) and English Wikipedia ( 2 , 500M words).
Based on the same corpora, we use the same preprocessing script 1 to create the pre-training data.
In ﬁne-tuning, we report our results on the Stanford Question Answering Dataset (SQuAD) and the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018).
We use two versions of SQuAD: V1.
and V2.0 (Rajpurkar et al., 2016; 2018).
The GLUE is a collection of datasets/tasks for evaluating natural language understanding systems 2 .
Input/Output representations: We follow the input/output representation setting from Devlin et al.
(2019) for both pre-training and ﬁne-tuning.
We use the WordPiece (Wu et al., 2016) embeddings with a 30 , 000 token vocabulary.
The ﬁrst token of every sentence is always a special classiﬁcation token ([CLS]).
The sentences are differentiated with a special token ([SEP]).
Evaluation:
In pre-training, BERT considers two objectives: masked language modeling (MLM) and next sentence prediction (NSP).
For MLM, a random sample of the tokens in the input sequence is selected and replaced with the special token ([ MASK ]) .
The MLM objective is a cross-entropy loss on predicting the masked tokens.
NSP is a binary classiﬁcation loss for predicting whether two segments follow each other in the original text.
In pre-training, we use MLM and NSP as training objectives to pre-train, retrain the BERT model, and as metrics to evaluate the BERT model .
In ﬁnetuning, F1 scores are reported for SQuAD, QQP and MRPC.
Matthew’s Corr and Pearson-Spearman
Corr are reported for CoLA and SST2 respectively.
Accuracy scores are reported for the other tasks.
All the experiments execute on one Google Cloud TPU V3-512 cluster, three Google Cloud TPU
V2-512 clusters and 110 Google Cloud TPU V3-8/V2-8 instances
Baseline: As there is no public effective BERT pruning method, we use the proposed NIP pruning method on BERT as the baseline method.
Th detail of NIP is shown in Appendix A.
The progressive pruning ratio is ∇ p = 10% (prune 10% more weights in each iteration).
Starting from the ofﬁcial
BERT BASE , we use 9 iterations.
In each iteration t of NIP, we get the sparse BERT BASE with speciﬁc sparsity, as ( w t ; S w t ) .
Then we retrain the sparse BERT BASE w t over the sparsity S w t .
In the retraining process, the initial learning rate is 2 · 10 − 5 , the batch size is 1024 and the retraining lasts for 10 , 000 steps (around 16 epochs).
For the other hyperparameters, we follow the original BERT paper Devlin et al.
In each iteration, the well retrained sparse BERT BASE is the starting point for the ﬁne-tuning tasks and the next iteration.
4. R EWEIGHED P ROXIMAL P RUNING (RPP)
We apply the proposed Reweighted Proximal Pruning (RPP) method on both BERT BASE and
BERT LARGE , and demonstrate performance improvement.
Detailed process of RPP is in Appendix B.
For BERT BASE , we use the hyperparameters exactly the same with our experiments using NIP.
The initial learning rate is λ =2 · 10 − 5 and the batch size is 1024.
We iterate the RPP for six times ( T =6 ), and each iteration lasts for 100 , 000 steps (around 16 epochs).
The total number of epochs in RPP is smaller than NIP when achieving 90% sparsity ( 96 < 144 ).
There is no retraining process in RPP.
We set γ ∈ { 10 − 2 , 10 − 3 , 10 − 4 , 10 − 5 } and (cid:15) =10 − 9 in Algorithm 1.
Recall that RPP reduces to (cid:96)sparse training as t =1 .
In Figure 2, we present the accuracy versus the pruning ratio for pre-training objectives MLM and
NSP, and ﬁne-tuning task SQuAD 1.1.
Here we compare RPP with NIP.
Along with the RPP continuing to iterate, the performance of RPP becomes notably higher than NIP for both the pre-training task and the ﬁne-tuning task.
The gap further increases as the RPP iterates more times.
In Figure 2, we ﬁnd that the NSP accuracy is very robust to pruning.
Even when 90% of the weights are pruned, the NSP accuracy keeps above 95% in RPP algorithm and around 90% in NIP algorithm.
For MLM accuracy and SQuAD F1 score, the performance drops quickly as the prune ratio increases.
RPP slows down the decline trend to a great extent.
On SQuAD 1.1 dataset/task, RPP keeps the F1 score of BERT BASE at 88.5 ( 0 degradation compared with original BERT) at 41 .
2% prune ratio, while the F1 score of BERT BASE applied with NIP drops to 84.6 ( 3 . 9 degradation) at 40% prune ratio.
At 80% prune ratio, RPP keeps the F1 score of BERT BASE at 84.7 ( 3 . 8 degradation), while the F score of BERT BASE applied with NIP drops to 68.8 ( 19 . 7 degradation compared with the original
BERT).
In addition to the ﬁne-tuning task of SQuAD 1.1, the other transfer learning tasks show the same trend (RPP consistently outperforms NIP) and the detailed results are reported in Appendix C.
For BERT LARGE , we use the hyperparameters exactly the same with our experiments using NIP except for the batch size.
The initial learning rate is 2 · 10 − 5 and the batch size is 512.
We iterate the RPP for four times ( T =6 ), and each iteration lasts for 100 , 000 steps (around 8 epochs).
There is no retraining process either.
We set γ ∈{ 10 − 2 10 − 3 10 − 4 10 − 5 } and (cid:15) =
10 − 9 in Algorithm 1.
The experimental results about pruning BERT LARGE and then ﬁne-tuning are shown in Table 1.
[Figure]
[Table]
[Figure]
4. V ISUALIZING A TTENTION P ATTERN IN BERT
We visualize the sparse pattern of the kernel weights in sparse BERT model applied with RPP, and present several examples in Figure 3.
Since we directly visualize the value of identical universal sparsity S w without any auxiliary function like activation map, the attention pattern is universal and data independent.
[Figure]
BERT’s model architecture is a multi-layer, bidirectional transformer encoder based on the original implementation (Vaswani et al., 2017).
Following (Vaswani et al., 2017), the transformer architecture is based on “scaled dot-product attention.”
The input consists of queries, keys and values, denoted as matrices Q , K and V , respectively.
The output of attention model is computed as
[Equation]
where d k is the dimension.
We visualize the sparse matrices Q and K of layer 2, layer 3 and layer 11 respectively in Figure 3.
From Figure 3, we have the following observations and analyses.
Structured pattern: Figure 3 demonstrates the structured pattern of non-zero weights in a pruned transformer block.
More speciﬁcally, we found that the pruned Q and K matrices within each transformer yield interesting group-wise structures (column-wise non-sparsity for query matrix and row-wise non-sparsity for key matrix).
Interestingly, we obtained these structured sparse patterns from our proposed RPP, an irregular pruning method (namely, no group-wise sparsity is penalized).
This is different from the irregular pruning on image classiﬁers, and thus shows the specialty of pruning on language models.
We also believe that the use of reweighted (cid:96) 1 approach matters to ﬁnd these ﬁne-grained sparse patterns.
Note that the structured sparsity pattern is more friendly to hardware implementation and acceleration than the non-structured pattern.
Semantic interpretation: The structured pattern found by RPP (visualized in Figure 3) has the following semantic interpretation.
What might the large-scale language representation learn?
The answer becomes clear after the language representation is pruned by RPP.
From the perspective of attention mechanism, the query matrix Q (column-wise non-sparsity) mainly models the attention information inside each sequence, while the key matrix K (row-wise non-sparsity) mainly models the attention information between different sequences in the context.
4. t -SNE V ISUALIZATION
t -Distributed Stochastic Neighbor Embedding ( t -SNE) is a technique for dimensionality reduction that is particularly well suited for the visualization of high-dimensional datasets (Maaten & Hinton, 2008).
Pre-trained word embeddings are an integral part of modern NLP systems (Devlin et al., 2019) and one contribution of BERT is the pre-trained contextual embedding.
Hence, we visualize word embedding in the original BERT model and the BERT model applied with RPP in Figureusing t -SNE.
Since BERT is different from commonly-studied image classiﬁer in network pruning, we would like to examine if pruning on BERT will lead to signiﬁcant change on the low-dimensional manifold of the language representation.
From Figure 4, we obtain the following observations and insights.
Low-dimensional manifold
: Figure 4 illustrates that, for both original BERT and BERT pruned with RPP, the low-dimensional manifolds of the language representations are similar, showing the similar projection.
Taking the speciﬁc word “intelligent” in Figure 4 as an example, the distribution of speciﬁc words and corresponding nearest words at the low-dimensional manifold (calculated using cosine/Euclidean distance) remain the high degree of similarity.
This implies that the BERT applied with RPP keeps most of the language representation information similar to that from the original BERT.
Linguistic interpretation of proper noun: There is one salient ribbon on the upper left of the macroscopical t-SNE visualization of word embeddings in either the original BERT or the pruned
BERT through RPP.
Each point in the ribbon represents a year number in annals.
There is also one salient short line on the lower left of the macroscopical t-SNE visualization of word embeddings in either the original BERT or the BERT applied with RPP.
Each point in most of the lines represents an age number.
Other proper nouns also reveal similar characteristics.
Our proposed RPP remains the embedding information of these proper nouns from the perspective of linguistic interpretation.
[Figure]
C ONCLUSIONS AND F UTURE W ORK
This paper presents the pruning algorithm RPP , which achieves the ﬁrst effective weight pruning result on large pre-trained language representation model - BERT.
RPP achieves 59 .
3% weight sparsity without inducing the performance loss on both pre-training and ﬁne-tuning tasks.
We spotlight the relationship between the pruning ratio of the pre-trained DNN model and the performance on the downstream multi-task transfer learning objectives.
We show that many downstream tasks except
SQuAD allows at least 80% pruning ratio compared with 59 .
3% under task SQuAD.
Our proposed
Reweighted Proximal Pruning provides a new perspective to analyze what does a large language representation (BERT) learn.
A A LGORITHM OF N EW I TERATIVE P RUNING
Algorithm 2 shows the detail process of our proposed NIP algorithm.
Algorithm
1: Input: Initial model weights w , initial prune ratio p = 0% , progressive prune ratio ∇ p 2: for t = 1 , 2 , . .
, T do 3: w = w ( t − 1) 4:
Sample batch of data from the pre-training data 5:
Obtain sparsity S w through hard threshold pruning, prune ratio p t = t · ∇ p 6:
Retrain w over sparsity constraint S w 7: for all tasks in {T i } do 8:
Fine-tune z i ( w ; S w ) over sparsity S w (if the desired prune ratio p t has been reached for downstream task i ) 9: end for 10: end for
B A LGORITHM OF R EWEIGHTED P ROXIMAL P RUNING (RPP)
Algorithm 3 shows the detail process of our enhanced AdamW (Loshchilov & Hutter, 2018) with proximal operator.
Algorithm
Our enhanced AdamW (Loshchilov & Hutter, 2018) with proximal operator 1:
Given α = 0 .
001 , β 1 = 0 .
9 , β 2 = 0 .
999 , (cid:15) =
10 − 6 , λ ∈ R 2: Initialize time step k ← 0 , parameters of pre-trained model w , ﬁrst moment vector m t =0 ← 0 , second moment vector v t =0 ← 0 , schedule multiplier η k =0 ∈ R 3: repeat 4: k ← k +5: ∇ f k ( w k − 1 ) ← SelectBatch ( w k − 1 ) 6: g k ← ∇ f k ( w k − 1 ) 7: m k ← β 1 m k − 1 + (1 − β 1 ) g k 8: v k ← β 2 v k − 1 + (1 − β 2 ) g 2 k 9: ˆ m k ← m k / (cid:0) 1 − β k 1 (cid:1) 10: ˆ v k ← v k / (cid:0) 1 − β k 2 (cid:1) 11: η k ← SetScheduleMultiplier( k ) 12:
a ← w k − 1 − η k (cid:0) α ˆ m k / ( √ ˆ v k + (cid:15) ) + λ w k − 1 (cid:1) 13: w k ← prox λ k ,rw − (cid:96) 1 ( a ) 14: until stopping criterion is met 15: return optimized sparse model w in pre-training
C D OWNSTREAM T RANSFER L EARNING T ASKS
As we mentioned in our main paper, we prune the pre-trained BERT model (using NIP and RPP) and then ﬁne-tune the sparse pre-trained model to different down-stream transfer learning tasks.
In this section, we exhibit the performance of pruned BERT using NIP and RPP on a wide range of downstream transfer learning tasks to demonstrate our conclusions in the main paper.
Through ﬁnetuning the pruning BERT over different downstream tasks, we found that SQuAD the most sensitive to the pruning ratio, showing an evident performance drop after 80% pruning ratio.
By contrast, the pruning can be made more aggressively when it is evaluated under other ﬁnetuning tasks.
This is not surprising, since SQuAD is a much harder Question Answering (QA) tasks, than other simple classiﬁcation tasks with limited solution space.
On the other hand, as the prune ratio of the pre-trained BERT increases, the performances on different transfer learning tasks descend generally.
The descending ranges differ in different downstream transfer learning tasks.
The descending range on SQuAD is the largest.
Our proposed RPP mitigates the descending trend on all downstream transfer learning tasks to a great extent, compared with NIP.
The intrinsic reason of this descending trend is left to future work.
C. QQP
Quora Question Pairs is a binary classiﬁcation task where the goal is to determine if two questions asked on Quora are semantically equivalent.
[Figure]
Finetuning setting: for ﬁne-tuning on QQP, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne tuned for 3 epochs.
C. MRPC
Microsoft Research Paraphrase Corpus consists of sentence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent.
Dolan & Brockett (2005)
Finetuning setting: for ﬁne-tuning on MRPC, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne-tune for 3 epochs.
[Figure]
C. MNLI
Multi-Genre Natural Language Inference is a large-scale, crowdsourced entailment classiﬁcation task Williams et al.
Given a pair of sentences, the goal is to predict whether the second sentence is an entailment, contradiction, or neutral with respect to the ﬁrst one.
Finetuning setting: for ﬁne-tuning on MNLI, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne-tune for 3 epochs.
[Figure]
(RPP)
MNLI(RPP) t=1 t= t= t= t=
[Figure]
t=
C. MNLIM
Multi-Genre Natural Language Inference has a separated evaluation MNLIM.
Following (Devlin et al., 2019), the ﬁne-tuning process on MNLIM is separated from MNLI.
So we present our results on MNLIM in this subsection.
Finetuning setting: for ﬁne-tuning on MNLIM, we set learning rate λ = 2 · 10 − 5 , batch sizeand ﬁne-tune for 3 epochs.
[Figure]
C. QNLI
Question Natural Language Inference is a version of the Stanford Question Answering Dataset (Rajpurkar et al., 2016) which has been converted to a binary classiﬁcation task Wanget al., 2018a).
The positive examples are (question, sentence) pairs which do contain the correct answer, and the negative examples are (question, sentence) from the same paragraph which do not contain the answer.
Finetuning setting: for ﬁne-tuning on QNLI, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne-tune for 3 epochs.
[Figure]
C. SST-
The Stanford Sentiment Treebank is a binary single-sentence classiﬁcation task consisting of sentences extracted from movie reviews with human annotations of their sentiment(Socher et al., 2013).
Finetuning setting: for ﬁne-tuning on SST-2, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne-tune for 3 epochs.
To consistent with the GLUE benchmark (Wang et al., 2018), the PearsonSpearman Corr score is reported.
[Figure]
C. C O LA
The Corpus of Linguistic Acceptability is a binary single-sentence classiﬁcation task, where the goal is to predict whether an English sentence is linguistically acceptable or not (Warstadt et al., 2018).
Finetuning setting: for ﬁne-tuning on CoLA, we set learning rate λ = 2 · 10 − 5 , batch size 32 and ﬁne-tune for 3 epochs.
To consistent with the GLUE benchmark (Wang et al., 2018), the Matthew’s
Corr score is reported.
[Figure]
' s C o rr f o r C o
L A
[Figure]
D N ON CONVERGENCE OF P RUNING BERT USING PREVIOUS METHODS
[Figure]
As we mentioned in our main paper, we investigate a series of pruning techniques to prune BERT, include the iterative pruning method (Han et al., 2015) and the one-shot pruning (Liu et al., 2019b).
However, most of the previous pruning techniques requires to directly optimize the (cid:96) 1 /(cid:96) 2 regularization using the back-propagation based gradient update in the original training of DNN models.
We execute a school of experiments and ﬁnd that, this kind of method to optimize regularization might not be compatible with BERT.
We show the experiment results about this incompatibility in this section.
For the sake of fair comparison, we not only adopt the same hyperparameters (in our experiments about NIP and RPP) on iterative pruning and one-shot pruning, we execute a wide set of hyperparamters to make the iterative pruning and one-shot pruning work.
We set the learning λ ∈ { 2 · 10 − 4 , 10 − 4 , 5 · 10 − 5 , 3 · 10 − 5 , 2 · 10 − 5 , 1 · 10 − 5 , 1 · 10 − 6 , 1 · 10 − 7 , 1 · 10 − 8 } , batch size B ∈ { 256 , 512 , 1024 , 2048 } .
We execute the same hyperparameters (with NIP and RPP) and attempt more hyperparameters on the iterative pruning and one-shot pruning, but iterative and one-shot pruning could not converge to a valid solution.
Figure A8 illustrates training loss curve of applying iterative pruning, one-shot pruning and RPP on BERT.
It is clear that iterative pruning and one-shot pruning leads to a non-convergence result, while different settings of RPP ( T = 0 , T = 1 , T = 2) converge well.
From the perspective of optimization and convergence, we make the following analysis:
The previous method, such as Iterative Pruning (IP) and one-shot pruning, relies on directly optimizing the (cid:96) 1 / (cid:96) 2 penalized training loss to conduct DNN pruning (this is discussed by Han et al. (2015) on iterative pruning, Section 3.1).
As a result, a simultaneous back-propagation (for updating model weights) is conducted over both the original training loss as well as the non-smooth sparsity regularizer.
When the penalty term is back-propagated together with the loss function, this affects the convergence of the original loss function.
The convergence performance is signiﬁcantly degraded for extremely large DNN model like BERT.
This phenomenon is also observed in the training of BERT (ADAM weight decay) that decouples the regularization term with the original loss function, instead of using an overall back-propagation.
For the super-deep DNN model (like BERT), it becomes harder to count on the one-time backpropagation ﬂow to solve both the original training objective and the sparse regulariztion at the same time.
Loshchilov & Hutter (2018) notices this limitation and improves regularization of Adam, by decoupling weight decay regularization from the gradient-based update (Loshchilov & Hutter, 2018).
AdamW is widely adopted in pretraining large language representations, e.g., BERT (Devlin et al., 2019), GPT (Radford et al., 2018) and XLNet (Yang et al., 2019).
The difference is that, the decoupled weight decay in AdamW is to avoid overﬁtting, while our purpose is to generate sparsity.
Moreover, previous algorithms with directly optimizing (cid:96) 1 through the back-propagation based gradient update penalty on TPU will easily lead to the gradient NaN.
Hence we proposed New Iterative Pruning (NIP) as our working baseline.
We believe that NIP works since NIP reﬂects the naturally progressive pruning performance without any externally introduced penalty.
As a ﬁx of IP, NIP simpliﬁes the training objective by removing the non-smooth sparsity regularizer.
This simple ﬁx improves the convergence of the training process, and makes new iterative pruning doable for BERT.
We hope that other pruning methods should not perform worse than
NIP, otherwise, the effect of optimizing the newly introduced sparsity-promoting regularization is negative.
To further improve the pruning performance, we need to ﬁnd a better pruning method that exploits our composite objective structure (original training loss + sparsity regularization), so that the backpropagation is not affected for the original training objective of BERT.
Motivated by that, proximal gradient provides an elegant solution, which splits the updating rule into a) gradient descent over the original training loss, and b) proximal operation over non-smooth sparsity regularizers.
Moreover, reweighted (cid:96) 1 minimization serves as a better sparsity generalization method, which self-adjusting the importance of sparsity penalization weights.
Furthermore, the incorporation of reweighted (cid:96)will not affect the advantage of the proximal gradient algorithm.
Thanks to the closed-form solution (equation 8) of proximal operation on a weighted (cid:96) 1 norm, Reweighted Proximal Pruning (RPP) is a desired pruning method on BERT model.
We hope RPP proves to be effective in more kinds of
DNN models in the future.