Block Pruning For Faster Transformers (2021)

Abstract: Pre-training has improved model accuracy for both classification and generation tasks at the cost of introducing much larger and slower models. Pruning methods have proven to be an effective way of reducing model size, whereas distillation methods are proven for speeding up inference. We introduce a block pruning approach targeting both small and fast models. Our approach extends structured methods by considering blocks of any size and integrates this structure into the movement pruning paradigm for fine-tuning. We find that this approach learns to prune out full components of the underlying model, such as attention heads. Experiments consider classification and generation tasks, yielding among other results a pruned model that is a 2.4x faster, 74% smaller BERT on SQuAD v1, with a 1% drop on F1, competitive both with distilled models in speed and pruned models in size.

## introduction pre-trained transformer models are the standard for nlp tasks in both classification and generation tasks @xcite . the recent trend is for models to continue to grow in size while yielding improved performance on standard benchmarks @xcite . this development highlights the need to reduce the storage size and increase the efficiency of pre-trained models. pruning methods have shown to be extremely effective at reducing the storage size of models fine-tuned for a specific task. approaches such as magnitude pruning @xcite , l0 regularization @xcite , lottery ticket hypothesis @xcite , diff pruning @xcite , and movement pruning @xcite have demonstrated remarkable reductions in model size. movement pruning produces 77% savings in parameter storage for a 1% drop in accuracy on squad v1.1. however, these models yield very little actual efficiency benefits, as to run them in standard hardware often requires reconstructing the original dense shape. on the other hand distillation methods have been more effective at producing faster models as has been shown by distilbert @xcite , tinybert @xcite or mobilebert @xcite . these approaches utilize targeted distillation to produce smaller models with a dense structure that is fast on standard hardware. however without careful engineering and size selection these models are much larger than pruned ones. in this work, we target closing this gap through block pruning. unlike pruning individual parameters, this approach encourages pruning that can be optimized on dense hardware. it is a less rigid approach than row or column-based pruning typically used in structured approaches @xcite , which have been difficult to apply effectively to transformers. we integrate this approach with movement pruning @xcite , a simple method for pruning pre-trained models during fine-tuning. the final method foot_0 has few additional hyperparameters or training requirements. experiments consider a large variety of different benchmark datasets comparing accuracy and efficiency. we find a surprising result that despite utilizing sub-row square blocks during training, the approach learns to eliminate full components of the model, effectively dropping a large number of attention heads. this effect allows the model to achieve speedups even beyond standard structured pruning of feed-forward layers. results show a 2.4x speedup on squad v1.1 with a 1% drop of f1, and a 2.3x speedup on qqp with a 1% loss of f1. experiments on summarization also show a 1.39x speedup for an average of 2 points drop on all rouge metrics on cnn/dailymail, and for a reduction of decoder weights of 3.5x. ## related work there has been a growing interest in the compression of pre-trained language models. we consider three varieties of methods: distillation, pruning, and structured pruning. knowledge distillation, introduced by @xcite , is a popular compression technique. researchers have applied this method to a variety of nlp models @xcite @xcite . distillation has been used to obtain significantly smaller bert models achieving competitive performances. @xcite distills bert into shallower students during the pre-training stage and optionally during the finetuning stage. mobilebert @xcite and tinybert @xcite are obtained thanks to a layer-wise distillation strategy. while the distillation of former is task-agnostic, the one used to obtain the latter is task-specific. other previous work has focused on unstructured pruning @xcite @xcite . when targeting transformer models, it is typical to select the weights to prune based on their magnitude @xcite , or by computing an importance score using a firstorder method @xcite . while these methods allow for a significant reduction in model size, specialized hardware is required to make use of the resulting unstructured sparse matrices in order to speed up inference. in contrast, structured pruning removes coherent groups of weights @xcite @xcite @xcite . recent works @xcite show that some heads can be removed without significant degradation in performance, leading to the conclusion that most heads provide redundant information. other authors have worked on combining matrix factorization and weight pruning. while @xcite combine svd-based matrix factorization with unstructured pruning, @xcite use structured pruning in order to reduce the rank. related to our approach, kim and awadalla (2020) and @xcite both apply structured pruning on the heads of the multi-head attention (mha) and on the inner-layer nodes of the feed-forward network (ffn). the former uses predefined pruning ratios, shared across all layers, in order to select the modules to prune after sorting them given an importance score. @xcite compares dif-ferent methods to compute the prunable module masks and find l0 regularization to perform the best. ## background starting with a transformer model with parameters θ, our goal is to produce a set of parameters θ that are both fine-tuned for a specific end-task and smaller in such a way that inference can be efficiently computed on parallel hardware. the two largest lines in the transformer parameter budget are the feed-forward network sublayer (ffn) and the multi-head attention sub-layer (mha). the ffn parameters consist of two matrices (w 1 and w 2 ) of transposed shape r d model ×d ff and r d ff ×d model where d model is the hidden size and d ff d model is the inner size. these are used in the standard fashion by the network. the mha parameters consist of 4 projection matrices (w q , w k , w v and w o ) of size r d model ×d model (query, key, value, out). these are used to project the hidden vector to and from the component attention parts. in implementations, this projection is made with the matrices in their folded tensor form r n heads × d model n heads ×d model where n heads is the number of attention heads. in standard fine-tuning, starting from θ, we optimize the loss l (for instance, cross-entropy for classification): score-based pruning methods @xcite modify the model by introducing score parameters s for each parameter i and replace the original parameter matrices with a masked version @xmath0. for instance, in the simplest version of magnitude pruning, the mask would just zero-out parameters with low absolute values. movement pruning @xcite ) is a scorebased pruning approach that encourages the model to optimize these score parameters. specifically, we focus on the soft-movement variant of movement pruning that sets m (s) = 1(@xmath1, and optimizes a regularized objective, this pruning objective encourages the model to fine-tune the parameters while lowering the scores of unimportant parameters and thus encouraging more sparsity. in order to train through the threshold, a straight-through estimator @xcite is used. movement pruning, combined with distillation, has shown to be a very effective method to reduce the number of parameters in an existing model yielding 94% pruning in our tests for a f1 of 87.5 on squad v1.1 (bert-base is 88.5). this results in significantly smaller models than distillation alone. however, even with this sparsity level, the model is not substantially faster when run on most standard hardware that cannot significantly take advantage of this style of sparse matrix-vector product. ## model: block movement pruning in this work, we extend movement pruning to work on blocks of local parameters. specifically, each matrix in the transformer is partitioned into fixedsized blocks. this setting goes beyond the arbitrary pruning of unstructured methods, with the goal of encouraging the data locality closer to what would be needed for efficiency. 2 our approach is extremely simple. for each parameter matrix @xmath2, we assume a fixedsized block structure (m , n ). each of these blocks acts as an individual group in the regularization with a shared score parameter derived from the corresponding score matrix @xmath3. computing the masked weight is done by expanding the thresholded values, i.e. as in past work, this model is trained with distillation to match the performance of a teacher model. unlike other distillation approaches that require fully specifying the new model structure, our method only requires the size and shapes of the blocks, i.e. the set of (m , n ) for each parameter matrix in the model. if blocks are too large, then they are difficult to prune, but if they are too small they do not support efficient inference. to reduce the search space, we will limit ourselves to test (m , n ) att and (m , n ) ff : the same block size will be used for all layers for attention weights w q , w k , w v and w o on one hand, and for the feed-forward weights w 1 and w 2 on the other hand. we split the movement pruning regularization term into: this allows us to take into account the difference in terms of gradient received by the score parameters. to reduce further the search space, we will test on two kinds of blocks: these block sizes allow for efficient models: blocks of size at least (16, 16) are efficient to compute with appropriate gpu kernels, whereas full rows, columns or heads can be entirely removed from the matrix: the remaining matrix is then dense. we also include two additional baseline block types used to verify the approach: the first considers small blocks, and the second considers very large functional blocks. cnn/dailymail ("cnn") is formulated as a conditional generation task. we report the performance on the development set as measured by the accuracy for mnli and sst-2, f1 for qqp, the exact match (em) and f1 for squad and rouge for cnn/dailymail. we experiment with task-specific pruning of transformer language models. we use bert @xcite (an encoder-only transformer language model with 110m parameters, among which 85m are part of the linear layers present in the transformer layers) for sentence classification and question answering (340m and 227m respectively for bert-large), and bart @xcite (an encoder-decoder language model with 139m parameters, among which 99m are part of the linear layers present in the transformer layers) for summarization (406m and 353m for bart-large). we compare against several baselines. movement pruning is a fully unstructured approach and gives an upper bound on the sparsity trade-offs we hope to achieve, even if it provides little speed benefit. we also compare our results against state-ofthe-art approaches developed for fast inference of transformer-based language models. distilbert @xcite is obtained by distilling through pre-training a pre-trained bert into a smaller model. tinybert @xcite distills a finetuned model while using data augmentation. mo-bilebert @xcite is the result of a large architecture search. dbart (shleifer and rush, 2020) is obtained by arbitrarily copying equally spaced layers of a large model to a smaller one. to measure inference speed on gpu, we use a 24gb 3090 rtx and an intel i7 cpu, using a large batch size (128) for evaluation and using pytorch cuda timing primitives. we measure the speed of other models in this same setup. results may be different from original papers, as latency and throughput characteristics are different for each platform. we also provide the number of parameters in the linear layers of the transformer layers for each of our models and for the reference ones: as the linear layers represent most of the flops, this is a good proxy for the computation required and to some extent for the compute time, when the model characteristics are equivalent. ## analysis large model pruning to test that this approach scales to large models, we apply hybrid pruning on bert-large on squad v1.1. we observe similar results: a 18% dense bert-large has a f1 of 90.2, with a speedup of 3.2x compared to bert-large with a f1 of 93.2. this pruned model is actually faster than a bert-base model ( ## conclusion we have shown that we can extract small pruned models that are at an equivalent or better than distilled networks. this approach can be done during fine-tuning and not pre-training. the method does not resort to techniques such as data augmentation or architecture search, and it works on a diverse set of tasks and base models. as better and larger models are published at an increasing pace, we can rely on a simple and robust method to accelerate them on specific tasks without sacrificing accuracy and distribute these models easily while keeping most of the original model accuracy.