Reweighted Proximal Pruning for Large-Scale Language Representation

1 INTRODUCTION . 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 fine-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-specific architecture modifications . BERT is conceptually simple and empirically powerful ( Devlin et al. , 2019 ) . However , along with the significant performance enhancement , the parameter volume and complexity of these pre-trained language representations significantly increase . As a result , it becomes difficult 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 ) , filter 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 fine-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 significant 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 ( ` 1 regularization ) and ridge regression ( ` 2 regularization ) . We find that the failure of previous methods on weight pruning of BERT is possibly due to the inaccurate sparse pattern learnt from the simple ` 1 or ` 2 based sparsitypromoting regularizer . In fact , the difficulty of applying regularization to generate weight sparsity coincides with the observation in ( Loshchilov & Hutter , 2018 ) on the imcompatibility of conventional weight decay ( ` 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 modified 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 ) . 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 ` 1 minimization ( Candes et al. , 2008 ) with proximal algorithm ( Parikh et al. , 2014 ) . RPP consists of two parts : the reweighted ` 1 minimization and the proximal operator . Reweighted ` 1 minimization serves as a better method of generating sparsity in DNN models matching the nature of weight pruning , compared with ` 1 regularization . Thanks to the closed-form solution of proximal operation on a weighted ` 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 first 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 first 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 fine-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 classification tasks , RPP helps to find 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 . 2 RELATED WORK . 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 fine-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 specific machine translation task ( out of over 10 transfer tasks ) , which is only a specific fine-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 ` 1 and proximal algorithm Candes et al . ( 2008 ) present reweighted ` 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 first work that applies reweighted ` 1 minimization to network compression , particularly for BERT pruning . , 3 REWEIGHTED PROXIMAL PRUNING 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 fine-tuning transfer learning tasks . Let fi denote the loss function of network for downstream task Ti ∼ 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 zi denote the i-th task-specified model parameters ( fine-tuning in BERT ) . The downstream tasks have separate fine-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 zi ( w ) are obtained through fine-tuning minimize w∈Rd fi ( w ) ( 1 ) 3.1 PRUNING FORMULATION IN TRANSFER LEARNING . Following the conventional weight pruning formulation , we first consider the problem of weight pruning during pre-training : minimize w∈Rd f0 ( w ) + γ‖w‖p ( 2 ) where f0 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 ` p constraint , |w‖p ≤ τ for some τ . Let ŵ denote the solution to problem ( 2 ) , and the corresponding sparse pattern Sŵ is given by Sŵ = { i|ŵi = 0 , ∀i ∈ [ d ] } ( 3 ) For a specific transfer task i , we allow an additional retraining/fine-tuning step to train/fine-tune weights starting from the pre-training results ŵ and subject to the determined , fixed sparse pattern Sŵ , denoted as zi ( ŵ ; Sŵ ) . That is , we solve the modified problem equation 1 minimize zi fi ( zi ( ŵ ; Sŵ ) ) ( 4 ) Here , different from ( 1 ) , the task-specific fine tuning weights variable zi ( ŵ ; Sŵ ) is now defined over Sŵ . Our goal is to seek a sparse ( weight pruned ) model during pre-training , with weight collection ŵ and sparsity Sŵ , which can perform as well as the original pre-trained model over multiple new tasks ( indexed by i ) . These fine-tuned models zi ( ŵ ; Sŵ ) ( for different i ) share the identical universal sparsity Sŵ .

Summary: Models such as BERT are pretrained language models which provide significant improvement for different tasks, however they suffer from high huge size and complexity. This paper has proposed using proximal gradient descent to find sparse weights for BERT to reduce the number of parameters and make the model smaller. They concentrate on the drawbacks of the previous sparse-based approaches and claimed that they have convergence issues (they have provided some evidence in the appendix). therefore, they propose to use reweighed sparse method and optimise it using proximal gradient descent which provides a closed form solution for sparse constraint.