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: The paper proposes a new approach to prune weights that is designed keeping large scale pre-trained language representations like BERT. Such a method is desirable for deploying such models on devices with limited memory like phones etc. Experiments on Squad and Glue datasets show that a pruned version of the model maintains high accuracy for these tasks.