Structured Pruning of Large Language Models (2020)

Abstract: Large language models have recently achieved state of the art performance across a wide variety of natural language tasks. Meanwhile, the size of these models and their latency have significantly increased, which makes their usage costly, and raises an interesting question: do language models need to be large? We study this question through the lens of model compression. We present a generic, structured pruning approach by parameterizing each weight matrix using its low-rank factorization, and adaptively removing rank-1 components during training. On language modeling tasks, our structured approach outperforms other unstructured and block-structured pruning baselines at various compression levels, while achieving significant speedups during both training and inference. We also demonstrate that our method can be applied to pruning adaptive word embeddings in large language models, and to pruning the BERT model on several downstream fine-tuning classification benchmarks.

## introduction recent advances in language modeling have led to remarkable improvements on a variety of natural language tasks @xcite @xcite @xcite @xcite . these models, however, have grown increasingly large, rendering them slow and expensive for realworld applications. through the use of model compression, we aim to reduce this overhead, and to better understand the role of model capacity in large language models. a common approach to model compression is known as weight pruning @xcite @xcite . model weights are progressively removed, resulting in sparse matrices across the network. earlier work focuses mostly on unstructured pruning, where weights are pruned individually @xcite . while this method is effective, it results in unstructured sparse matrices that are difficult to support on common hardware @xcite , making it challenging to obtain training and inference speedups despite a significant reduction in model size. on the other hand, structured pruning imposes structured sparse patterns by removing groups of consecutive parameters, such as rows, columns or k×k sub-blocks of the weight matrix @xcite @xcite . these methods lead to significant speedup, but tend to give lower performance than unstructured pruning given the same parameter budget @xcite . another caveat is that some of these methods require special linear algebra implementations @xcite or hardware @xcite in order to accelerate matrix multiplication, therefore limiting their application to a broad set of existing models. we propose a generic, improved structured pruning approach based on adaptive low-rank factorization. as an alternative to unstructured sparse and block sparse representations, low-rank factorization retains the full dense structure of weight matrices, eliminating the need for special linear algebra primitives and hardware for computation speedup. compared to row (and column) based pruning, lowrank factorization better preserves the linear transformation of the un-compressed matrices. during training, our method adaptively learns which lowrank components to remove in order to achieve a strong performance-compression trade-off. we show that a simple magnitude based pruning strategy is sufficient to accomplish strong results. in addition, we further increase performance via an improved l 0 regularization @xcite technique which uses an augmented lagrangian method to directly control the final compression level of the model. our method, which we refer to as flop (factorized low-rank pruning) applies to any matrix multiplication. pruning large language models introduces unique challenges with the handling of large input and output layers. although our method is generic, it is particularly well suited to this task. in particular, we show that flop can dynamically learn the embedding dimensions of different word clusters, effectively extending the idea of adaptive embeddings and softmax @xcite . since these embedding layers take a significant amount of parameters in the language models, learning flexible dimensions instead of specifying them manually results in a more optimal trade-off between parameter reduction and performance. we evaluate our method on common language modeling and language understanding tasks including the wiki-103, enwiki8 and glue benchmarks, and by testing our method on both recurrent networks and transformer @xcite . our results demonstrate that factorization based pruning significantly outperforms block-structured pruning and even surpasses unstructured pruning, while using our improved l 0 regularization further improves the performance in most cases. when pruning a large word-level language model with adaptive embeddings for example, our method achieves 50% compression while losing only 0.8 perplexity. moreover, our method is able to achieve over 2x speed-up during both training and inference with no additional hardware or software requirements. our method will be released as a pytorch @xcite library. ## related work the development of model compression techniques can be categorized into three areas of research: weight pruning @xcite , knowledge distillation @xcite @xcite , and quantization @xcite @xcite . recent efforts have successfully applied compression on various architectures and nlp applications, such as pruning multi-head attentions for machine translation @xcite , learning adaptive embeddings and softmax layers for lan-guage models @xcite @xcite , and compressing bert models via distillation @xcite @xcite @xcite . only one of the compression techniques such as distillation has been used in these works for simplicity. however, these techniques can be combined to achieve greater compression @xcite . our pruning method is compatible with quantization and distillation, as it can be applied to compress any matrix multiplication in a network. previous work has considered different weight pruning approaches such as unstructured pruning based on magnitude @xcite , dropout @xcite @xcite , and structured pruning @xcite . model weights are often removed via thresholding and l 1 regularization during the pruning process @xcite . our method differs from previous work by using lowrank parameterization for compression. furthermore, we extend l 0 regularization using an augmented lagrangian optimization method to control the final model size. ## background we formalize the task of model pruning as an endto-end learning problem with l 0 regularization, following the prior work of @xcite . consider a given neural network model f (•; θ) parameterized by θ = {θ j } n @xmath0, where each θ j represents an individual parameter weight or a block of weights (e.g. a column of a weight matrix) and n denotes the number of blocks. a pruning strategy of the model can be parameterized by introducing additional binary variables @xmath1, 1} and θ = θ z ∀j θj = θ j z j . here θ = { θj } denotes the set of model parameters after pruning and its l 0 norm, θ 0 = n @xmath2, measures the effective size of the pruned model. the choice of binary variables z can be regulated by some prior distribution and optimized given the training data. that is, let q j (z) be the density function of the learnable prior of z j . the optimization objective during training can be formulated as min-imizing the expected training loss the key to the method of @xcite , called the re-parameterization trick, enables z to be differentiable and jointly trained with the model parameter θ. specifically, the random variables z are relaxed as continuous variables distributed within the interval [0, 1]. in addition, instead of learning the probability density function q j (z), the re-parameterization trick proposes to learn the inverse of the cumulative density function (cdf). note that if g() is the inverse of cdf for a variable z, then z can be easily sampled by first sampling u ∼ u (0, 1) and computing @xmath3. assuming the inverse cdf function is parameterized by some learnable parameters α = {α j } n @xmath4; α) is differentiable, we obtain an overall end-to-end learning objective, following previous work, we choose the hard concrete distribution for the random variables @xmath5. the inverse of cdf g(•; α) of this distribution is defined as follows ## method in this section, we introduce flop , an improved structured pruning method. flop proposes a different parameterization of the weight matrices using low-rank factorization. in addition, we introduce a revised optimization objective that allows for an explicit control of the compression size. ## experimental setup tasks we evaluate the performance of our method on language modeling and bert fine-tuning. specifically, we consider the following task setup. 1. recurrent word-level language models on the wiki-103 dataset. we adopt sru @xcite as the recurrent architecture and tied adaptive embedding and softmax layers @xcite . our base model consists of 12 recurrent layers, 100m parameters in total. about 50% of the parameters are used for the adaptive layers. ## results word-level language model the sru base model (unpruned) achieves a test perplexity of 24.5, being a strong starting point and competitive with top-performing models such as transformer @xcite . the pruning results conform to our expectations that pruning a large model is consistently better than training a small model from scratch, and using low-rank based pruning yields better performance than removing matrix columns and input features. flop exceeds the performance of fac, np-l 0 and agp baselines at all compression levels tested. the performance of flop-agp, especially in comparison with its unstructured counterpart agp, highlights the effectiveness of factorization based pruning. moreover, we achieve a test perplexity (ppl) of 25.3 with flop-l 0 method, a loss of 0.8 perplexity score, while removing 50% of the model parameters. this result is impressive since our base model adopts the adaptive word embedding and softmax layers, which already reduce the model size significantly. figure 1 illustrates how our method adaptively controls the size of different model components. ## analysis in this section, we perform an analysis of several aspects of our method. factorization one of the key hypotheses outlined in this paper is that pruning input dimensions (equivalently rows or columns of weight matrices) is a more restrictive form of pruning compared to our factorization based strategy. however, one could also argue that the factorization method works better simply because the hidden size can be initially set much larger than an unfactorized model, not because of pruning itself. for instance, variants size 0% 70% 80% 85% 90% np-l 0 37m 1.30 1.31 (-0.8%) 1.34 (-3.2%) 1.37 (-5.4%) 1.43 (-10.0%) 66m 1.25 1.28 (-2.4%) 1.31 (-4.8%) 1.32 (-5.6%) 1.37 (-9.6%) flop -l 0 35m speed analysis thanks to its structured nature, flop can achieve significant computation speedup. as shown in learning dynamics figure 2 demonstrates the training dynamics of the hardconcrete distribution. we plot the histogram of hardconcrete parameters α after every few thousands of training iterations. a negative value of α indicates that the associated parameter is likely to be pruned while a positive value indicates the opposite. the magnitude of the value reflects the certainty of the pruning decision. as illustrated by the figure, the distribution of α becomes bi-modal after initial exploration. certain parameters within each layer are completely pruned while others are kept with (almost) absolute certainty. in addition, the dynamics vary across different layers. for instance, for sru the first recurrent layer gets pruned more aggressively than the last layer. ## conclusion in this work, we present a generic structured pruning method based on adaptive low-rank factorization. we systematically evaluate the performance of this method on large language models. we show that our method can provide significant speedups and compression rates on large models while losing minimal performance compared to other methods, including unstructured magnitude pruning. this work contributes to reducing the growing overhead of large language models, and shines a light on the role of model capacity in language modeling.