T ext P runer: A Model Pruning Toolkit for Pre-Trained Language Models (2022)

Abstract: Pre-trained language models have been prevailed in natural language processing and become the backbones of many NLP tasks, but the demands for computational resources have limited their applications. In this paper, we introduce TextPruner, an open-source model pruning toolkit designed for pre-trained language models, targeting fast and easy model compression. TextPruner offers structured post-training pruning methods, including vocabulary pruning and transformer pruning, and can be applied to various models and tasks. We also propose a self-supervised pruning method that can be applied without the labeled data. Our experiments with several NLP tasks demonstrate the ability of TextPruner to reduce the model size without re-training the model.

## introduction large pre-trained language models (plms) @xcite have achieved great success in a variety of nlp tasks. however, it is difficult to deploy them for real-world applications where computation and memory resources are limited. reducing the pre-trained model size and speeding up the inference have become a critical issue. pruning is a common technique for model compression. it identifies and removes redundant or less important neurons from the networks. from the view of the model structure, pruning methods can be categorized into unstructured pruning and structured pruning. in the unstructured pruning, each model parameter is individually removed if it reaches some criteria based on the magnitude or importance score @xcite @xcite . the unstructured pruning results in sparse matrices and allows for significant model compression, but the inference speed can hardly be improved without specialized devices. while in the structured pruning, rows or columns of the parameters are removed from the weight matrices @xcite @xcite @xcite . thus, the resulting model speeds up on the common cpu and gpu devices. pruning methods can also be classified into optimization-free methods @xcite and the ones that involve optimization @xcite . the latter usually achieves higher performance, but the former runs faster and is more convenient to use. pruning plms has been of growing interest. most of the works focus on reducing transformer size while ignoring the vocabulary @xcite . pruning vocabulary can greatly reduce the model size for multilingual plms. in this paper, we present textpruner, a model pruning toolkit for plms. it combines both transformer pruning and vocabulary pruning. the purpose of textpruner is to offer a universal, fast, and easy-to-use tool for model compression. we expect it can be accessible to users with little model training experience. therefore, we implement the structured optimization-free pruning methods for its convenient use and fast computation. pruning a base-sized model only requires several minutes with textpruner. textpruner can also be a useful analysis tool for inspecting the importance of the neurons in the model. textpruner has the following highlights: ## pruning methodology we briefly recall the multi-head attention (mha) and the feed-forward network (ffn) in the transformers @xcite . then we describe how we prune the attention heads and the ffn based on the importance scores. ## experiments in this section, we conduct several experiments to show textpruner's ability to prune different pretrained models on different nlp tasks. we mainly focus on the text classification task. we list the results on the mrc task and ner task with different pre-trained models in the appendix. ## conclusion and future work this paper presents textpruner, a model pruning toolkit for pre-trained models. it leverages optimization-free pruning methods, including vocabulary pruning and transformer pruning to reduce the model size. it provides rich configuration options for users to explore and experiment with. textpruner is suitable for users who want to prune their model quickly and easily, and it can also be used for analyzing pre-trained models by pruning, as we did in the experiments. for future work, we will update textpruner to support more pre-trained models, such as the generation model t5 @xcite . we also plan to combine textpruner with our previously released knowledge distillation toolkit textbrewer @xcite ) into a single framework to provide more effective model compression methods and a uniform interface for knowledge distillation and model pruning.