Exploring the Limits of Simple Learners in Knowledge Distillation for Document Classification with D oc BERT (2020)

Abstract: Fine-tuned variants of BERT are able to achieve state-of-the-art accuracy on many natural language processing tasks, although at significant computational costs. In this paper, we verify BERT’s effectiveness for document classification and investigate the extent to which BERT-level effectiveness can be obtained by different baselines, combined with knowledge distillation—a popular model compression method. The results show that BERT-level effectiveness can be achieved by a single-layer LSTM with at least 40\times fewer FLOPS and only {\sim}3\% parameters. More importantly, this study analyzes the limits of knowledge distillation as we distill BERT’s knowledge all the way down to linear models—a relevant baseline for the task. We report substantial improvement in effectiveness for even the simplest models, as they capture the knowledge learnt by BERT.

## introduction transformer-based @xcite pretrained contextual word embedding models such as bert @xcite and xlnet @xcite currently power many of the state-of-the-art models across various natural language processing (nlp) tasks. however, these models consume immense computational resources @xcite . with the surge of such pre-trained models being developed in quick succession, there is a need for effective compression techniques for their inexpensive deployment. knowledge distillation (kd; @xcite has been shown to be a fairly straightforward and effective model-agnostic compression method, which transfers knowledge learnt by huge models into more efficient models. in this work, we investigate if bert-level effectiveness can be achieved by more efficient models using kd. and more importantly, if so, how simple can these models be? we investigate these questions through the lens of document classification-a setting where these computational concerns are particularly relevant due to potentially long document lengths. further, in previous work, neural networks as an architectural choice have been questioned owing to the effectiveness of simple bag-of-words baselines @xcite . we first confirm that a fine-tuned bert model leads to state-of-the-art model quality by a substantial margin on standard document classification benchmarks. following this, we investigate the extent to which bert-level effectiveness can be obtained by various different baselines, combined with kd. we demonstrate, quite surprisingly, that it is possible to apply kd successfully on impoverished student models, such as a single-layer convolutional neural network (cnn) @xcite and even linear models. the key contributions of this work are as follows: ## background and methods typically, the task of document classification deals with classifying long texts (documents). more often than not, a document may be associated with more than one label, thus exposing the classifiers to multi-label classification and class imbalance. here, we review a subset of approaches developed to solve the task and highlight the methods that we compare and build upon in this work. ## datasets we use the following four datasets to evaluate bert: reuters-21578 (reuters; @xcite , arxiv academic paper dataset (aapd; @xcite , imdb reviews, and yelp 2014 reviews. reuters and aapd are multi-label datasets while documents in imdb and yelp '14 contain only a single label per document. for reuters, we use the standard modapté splits @xcite ; for aapd, we use the splits provided by @xcite ; for imdb and yelp, following @xcite , we randomly sample 80% of the data for training and 10% each for validation and test. ## training and hyperparameters as a simple and straightforward adaptation of bert models @xcite for document classification, we introduce a fully-connected layer over the final hidden state corresponding to the [cls] input token. during fine-tuning, we optimize the entire model end-to-end, with the additional softmax classifier parameters @xmath0, where h is the dimension of the hidden state vectors and k is the number of classes. we minimize the cross-entropy and binary cross-entropy loss for single-label and multi-label tasks, respectively. while fine-tuning bert, we optimize the number of epochs, batch size, learning rate, and maximum sequence length (msl; i.e., the number of tokens that documents are truncated to). for knowledge distillation, we train the lstm, kimcnn, and lr to capture the learnt representations from bert large using the objective of the type shown in equation ( 1 ). depending upon the dataset, we use cross-entropy or binary cross-entropy loss as l classif ication , equation (1). for l distill , following @xcite , we minimize the kullback-leibler (kl) divergence kl(p||q) where p and q are the class probabilities produced by the student and the teacher models, respectively. to build an effective transfer set for distillation as suggested by @xcite , we augment the training splits of the datasets by applying pos-guided word swapping and random masking, as in @xcite , along with randomizing the order of the sentences of documents in the training set. the transfer set sizes for reuters, imdb and aapd are 3×, 4×, and 4× their training splits, respectively; for yelp2014, no data augmentation was performed due to computational restrictions. refer to the appendix for further details regarding the training hyperparameters. ## results and discussion in rows 1-7 report the model quality of pre-bert models (that do not take advantage of pre-training). as observed by @xcite , lr and svm trained with tf-idf vectors form effective baselines as they challenge many neural networkbased baselines (e.g., han) on multiple datasets. this raises the question whether neural networks are a suitable architectural choice for document classification. however, at a much higher computational cost, the regularized lstm @xcite (row 7) achieves the best numbers for the class of models that do not exploit pre-training. consistent with @xcite , the bertbased models achieve state-of-the-art results on all four datasets (see surprisingly, distilled lstm (kd-lstm, row 10) achieves parity with bert base on average for reuters, aapd, and imdb. in fact, it outperforms bert base (on both dev and test) in at least one of the five runs. for yelp, we see that kd-lstm reduces the difference between bert base and lstm, but not to the same extent as in the other datasets. next, we explore the limits of kd by further distilling bert base all the way down to kimcnn @xcite ) (a single-layer cnn) and lr. it is not surprising that these models don't come close to bert base owing to their limited expressivity. however, interestingly, we see massive leaps in model quality of these models after distillation (rows 1-3; 11-12). specifically for multi-label datasets, both these models beat or come close to han and sgm, which are far more complex models. to put things in perspective, lr is a simple fully-connected layer and kimcnn contains merely ∼ 0.4% parameters of bert base . these results demonstrate that kd can yield a broad spectrum of baselines for varying computational costs, all of which can be useful depending on the requirements. ## conclusion and future work in this paper we improve baselines for document classification by fine-tuning bert (docbert). using docbert, we show the effectiveness of kd over a range of efficient models-a single-layer lstm model, a single layer cnn, and a logistic regression trained on tf-idf. this provides us with a spectrum of baselines for varying tradeoffs in classification accuracy and complexity. in fact, we show that the distilled lstm model achieves bert base parity on a majority of datasets, using only ∼ 3% parameters of the latter. while distillation is an effective way to reduce computational cost during inference, it doesn't aid in reducing resources needed for training. thus, methods for reducing the computational resources required while training deserve attention in future.