T ext B rewer: A n O pen- S ource K nowledge D istillation T oolkit for N atural L anguage P rocessing (2020)

Abstract: In this paper, we introduce TextBrewer, an open-source knowledge distillation toolkit designed for natural language processing. It works with different neural network models and supports various kinds of supervised learning tasks, such as text classification, reading comprehension, sequence labeling. TextBrewer provides a simple and uniform workflow that enables quick setting up of distillation experiments with highly flexible configurations. It offers a set of predefined distillation methods and can be extended with custom code. As a case study, we use TextBrewer to distill BERT on several typical NLP tasks. With simple configurations, we achieve results that are comparable with or even higher than the public distilled BERT models with similar numbers of parameters.

## introduction large pre-trained language models, such as gpt @xcite , bert @xcite , roberta @xcite and xlnet @xcite have achieved great success in many nlp tasks and greatly contributed to the progress of nlp research. however, one big issue of these models is the high demand for computing resources -they usually have hundreds of millions of parameters, and take several gigabytes of memory to train and inference -which makes it impractical to deploy them on mobile devices or online systems. from a research point of view, we are tempted to ask: is it necessary to have such a big model that contains hundreds of millions of parameters to achieve a high performance? motivated by the above considerations, recently, some researchers in the nlp community have tried to design lite models @xcite , or resort to knowledge distillation (kd) technique to compress large pretrained models to small models. kd is a technique of transferring knowledge from a teacher model to a student model, which is usually smaller than the teacher. the student model is trained to mimic the outputs of the teacher model. before the birth of bert, kd had been applied to several specific tasks like machine translation @xcite in nlp. while the recent studies of distilling large pre-trained models focus on finding general distillation methods that work on various tasks and are receiving more and more attention @xcite @xcite @xcite @xcite . though various distillation methods have been proposed, they usually share a common workflow: firstly, train a teacher model, then optimize the student model by minimizing some losses that are calculated between the outputs of the teacher and the student. therefore it is desirable to have a reusable distillation workflow framework and treat different distillation strategies and tricks as plugins so that they could be easily and arbitrarily added to the framework. in this way, we could also achieve great flexibility in experimenting with different combinations of distillation strategies and comparing their effects. in this paper, we introduce textbrewer, a pytorch-based distillation toolkit for nlp that aims to provide a unified distillation workflow, save the effort of setting up experiments and help users to distill more effective models. textbrewer provides simple-to-use apis, a collection of distillation methods, and highly customizable configurations. it has also been proved able to distill bert models efficiently and reproduce the state-of-theart results on typical nlp tasks. the main features of textbrewer are: textbrewer also provides some useful utilities such as model size analysis and data augmentation to help model design and distillation. ## related work recently some distilled bert models have been released, such as distilbert @xcite , tinybert @xcite , and ernie slim foot_1 . distilbert performs distillation on the pretraining task, i.e., masked language modeling. tinybert performs transformer distillation at both the pre-training and task-specific learning stages. ernie slim distills ernie @xcite )on a sentiment classification task. their distillation code is publicly available, and users can replicate their experiments easily. however, it is laborious and error-prone to change the distillation method or adapt the distillation code for some other models and tasks, since the code is not written for general distillation purposes. there also exist some libraries for general model compression. distiller @xcite and paddleslim foot_2 are two versatile libraries supporting pruning, quantization and knowledge distillation. they focus on models and tasks in computer vision. in comparison, textbrewer is more focused on knowledge distillation on nlp tasks, more flexible, and offers more functionalities. based on pytorch, it provides simple apis and rich customization for fast and clean implementations of experiments. ## architecture and design figure 1 shows an overview of the main functionalities and architecture of textbrewer. to support different models and different tasks and meanwhile stay flexible and extensible, textbrewer provides distillers to conduct the actual experiments and configuration classes to configure the behaviors of the distillers. ## experiments in this section, we conduct several experiments to show textbrewer's ability to distill large pretrained models on different nlp tasks and achieve results are comparable with or even higher than the public distilled bert models with similar numbers of parameters. foot_3 4.1 settings datasets and tasks. we conduct experiments on both english and chinese datasets. for english datasets, we use mnli @xcite for text classification task, squad1.1 @xcite for span-extraction machine reading comprehension (mrc) task and @xcite for named entity recognition (ner) task. for chinese datasets, we use the chinese part of xnli models. all the teachers are bert base -based models. for english tasks, teachers are initialized with the weights released by google foot_4 and converted into pytorch format via transformers foot_5 . for chinese tasks, teacher is initialized with the pre-trained roberta-wwm-ext foot_6 @xcite . we test the performance of the following student models: t3-small and t4-tiny are initialized randomly. the model structures of the teacher and students are summarized in ## results on chinese datasets the results on chinese datasets are presented in ## conclusion and future work in this paper, we present textbrewer, a flexible pytorch-based distillation toolkit for nlp research and applications. textbrewer provides rich customization options for users to compare different distillation methods and build their strategies. we have conducted a series of experiments. the results show that the distilled models can achieve state-of-the-art results with simple settings. textbrewer also has its limitations. for example, its usability in generation tasks such as machine translation has not been tested. we will keep adding more examples and tests to expand textbrewer's scope of application. apart from the distillation strategies, the model structure also affects the performance. in the future, we aim to integrate neural architecture search into the toolkit to automate the searching for model structures.