Gradient-based Intra-attention Pruning on Pre-trained Language Models (2023)

Abstract: Pre-trained language models achieve superior performance but are computationally expensive. Techniques such as pruning and knowledge distillation have been developed to reduce their sizes and latencies. In this work, we propose a structured pruning method GRAIN (gradient-based intra-attention pruning), which performs task-specific pruning with knowledge distillation and yields highly effective models. Different from common approaches that prune each attention head as a whole, GRAIN inspects and prunes intra-attention structures, which greatly expands the structure search space and enables more flexible models. We also propose a gradient separation strategy that reduces the interference of distillation on pruning for a better combination of the two approaches. Experiments on GLUE, SQuAD, and CoNLL 2003 show that GRAIN notably outperforms other methods, especially in the high sparsity regime, and achieves 6 7x speedups while maintaining 93% 99% performance. Under extreme compression where only 3% transformer weights remain, the pruned model is still competitive compared to larger models.

## introduction transformer-based @xcite pretrained language models (plms) have achieved great success and become the backbones of various natural language processing tasks. however, plms are computationally expensive and slow in inference due to their large sizes, which limits their applications in real-world scenarios. hence, a growing interest has been in developing compression and acceleration methodologies for plms. a common approach to model compression is structured pruning, which compresses the model by removing groups of consecutive parameters, namely the pruning units. in applying structured 1 code is available at https://github.com/airaria/ grain . pruning on plms, recent works have investigated removing units such as hidden dimensions in feedforward layers, attention heads in the multi-head attention @xcite , and coarse-grained units such as multi-head attention layers and feed-forward layers @xcite . however, these pruning units only span a small space of model structures and limit the exploration for better structures. for example, in the pruning of bert base @xcite , which contains 144 attention heads, the possible choices of attention heads for the pruned model are limited. block pruning @xcite extends pruning units by considering blocks in the weight matrices, but block pruning is not a fully structured pruning method and can not achieve large speedups. in this work, we propose grain (gradientbased intra-attention pruning), a structured pruning method that prunes plms with finer pruning units. in the following, we present the method from three aspects: pruning units, pruning algorithm, and training objectives. pruning units unlike attention heads pruning where the pruning unit is a single head, we propose intra-attention pruning, which inspects and prunes the structures inside attention heads. intra-attention pruning greatly expands the search space of model structures, making the resulting models more likely to find better structures. however, directly applying intra-attention pruning yields fragmented models, i.e., models with many small heads. the fragmented models have relatively large latencies on devices like gpus. to overcome the shortcoming, we introduce structure regularization, which encourages prioritizing specific units for pruning. structure regularization helps generate more regular structures and achieve lower latencies. pruning algorithm pruning algorithms decide which units to be removed. we adapt the gradientbased pruning algorithm @xcite for intra-attention pruning. gradient-based pruning is a light-weighted method that estimates the importance of the pruning units with gradient-based scores and then prunes the least important ones. in addition, we conduct the pruning in an iterative manner @xcite , i.e., the model is gradually pruned during fine-tuning. the iterative approach has been employed in combination with pruning algorithms such as movement pruning @xcite and magnitude pruning @xcite , but few works have combined it with gradient-based pruning. we find that iterative gradient-based pruning is especially effective despite its simplicity. training objectives as another common approach to model compression, knowledge distillation offers highly effective training objectives @xcite . pruning with distillation objective shows improved performance @xcite . however, in gradient-based pruning, the distillation objectives may disturb the estimation of importance scores. we propose a gradient separation strategy that uses different gradients for model optimization and importance score estimation. we show that this method leads to better performance. grain performs task-specific pruning without additional pre-training or data augmentation. in the experiments, we compare grain with strong pruning and distillation baselines on glue, squad, and conll 2003. grain notably outperforms the comparable methods in the high-sparsity regime. a demonstration of the results on mnli is shown in figure 1 . while keeping 5% parameters in transformers, grain maintains 93% ∼ 99% performance of bert base and 6 ∼ 7× speedups across different tasks. furthermore, grain still achieves competitive results even under extreme compression where only 3% transformer weights remain. ## related work a growing number of works have been devoted to the compression and acceleration of plms. most of the works have combined multiple techniques. knowledge distillation @xcite is a training technique that trains a student model to mimic the outputs and intermediate representations of the teacher model @xcite . distilbert @xcite and tinybert @xcite are both small bert-like models distilled with general and task-specific distillation. mobilebert @xcite and kroneckerbert @xcite have designed novel structures for student models. @xcite proposes to extract a subnetwork from the teacher and then perform distillation. autotinybert @xcite combine distillation with neural architecture search to find optimal hyperparameters. dynabert @xcite apply task-specific distillation and can flexibly adjust the model size. in this work, we only apply task-specific distillation, which consumes fewer resources. structured pruning on plms remove different types of units from the models, like attention heads @xcite , ffn hidden dimensions @xcite , blocks of weights @xcite , mha layers or ffn layers @xcite . many works combine pruning with other methods. @xcite presents a structured pruning approach with low-rank factorization of weight matrices. @xcite and @xcite apply pruning with knowledge distillation. in this work, we apply matrix factorization on the embeddings and use distillation and pruning to reduce the size of transformers. unstructured pruning removes each weight individually based on its magnitude @xcite @xcite , or the score computed by first-order @xcite or second-order @xcite ) method. unstructured pruning yields higher sparsity models but is hard to speed up without specialized devices for sparse matrix operations. in this work, we only consider structured pruning. besides model compression, another group of acceleration methods is dynamic inference, where the computation cost is determined at test time @xcite @xcite . @xcite and @xcite have proposed to integrate model compression with dynamic inference. we do not consider dynamic inference in this work and leave it for future work. ## preliminaries 3.1 transformers a transformer block @xcite is mainly composed of a multi-head attention (mha) layer and a feed-forward network (ffn) layer. let @xmath0, where n is the length, and d is the hidden size. an attention head is parameterized by the matrices following the mha layer is the feed-forward network layer. it consists of two linear layers and a gelu activation @xcite ) a transformer block contains other components, such as layernorm and residual connection, but they only take up a few parameters. ## methodology grain performs task-specific intra-attention pruning together with knowledge distillation. the overview of grain is depicted in figure 2 . following previous works, we only include the encoder in counting the model size unless otherwise specified. we refer to the size of the pruned model relative to the unpruned model as model density: sparsity is equal to one minus model density. ## conclusion this paper proposes grain, a gradient-based structured pruning method that expands the structure search space by pruning with intra-attention structures. we provide a structure regularization strategy that encourages finding regular structures and helps achieve lower latencies. we also combine pruning with distillation. we propose to separate the gradients from different losses to reduce the interference. grain is computationally efficient since it does not require pre-training or data augmentation. experiments show that grain achieves impressive high performance and outperforms other methods at different model densities on various natural language understanding tasks and meanwhile maintains competitive speedups.