Sparse Distillation: Speeding Up Text Classification by Using Bigger Student Models (2022)

Abstract: Distilling state-of-the-art transformer models into lightweight student models is an effective way to reduce computation cost at inference time. The student models are typically compact transformers with fewer parameters, while expensive operations such as self-attention persist. Therefore, the improved inference speed may still be unsatisfactory for real-time or high-volume use cases. In this paper, we aim to further push the limit of inference speed by distilling teacher models into bigger, sparser student models – bigger in that they scale up to billions of parameters; sparser in that most of the model parameters are n-gram embeddings. Our experiments on six single-sentence text classification tasks show that these student models retain 97% of the RoBERTa-Large teacher performance on average, and meanwhile achieve up to 600x speed-up on both GPUs and CPUs at inference time. Further investigation reveals that our pipeline is also helpful for sentence-pair classification tasks, and in domain generalization settings.

## introduction large pre-trained transformers @xcite are highly successful, but their large inference costs mean that people who host low-latency applications, or who are simply concerned with their cloud computing costs have looked for ways to reduce the costs. prior work mainly achieves this by leveraging knowledge distillation @xcite , which allows for the capabilities of a large well-performing model known as the teacher to be transferred to a smaller student model. for example, distillbert @xcite is a smaller transformer model distilled from bert @xcite , which reduces bert's size by 40% and becomes 60% faster during inference. however, such speed-up may be still @xcite and knowledge distillation, we obtain a student model with competitive performance on imdb dataset, while being 607x faster than roberta-large, and 20x faster than bi-directional lstms at inference time. counter to the convention of using "smaller, faster, @xcite cheaper" @xcite student models, our work explores a new area of the design space, where our fast and cheap student model is actually several times larger than the teacher. the student model we use is modified from deep averaging network (dan) in @xcite . dans take a simple architecture by mapping the n-grams in the input sentence into embeddings, aggregating the embeddings with average pooling, and then using multiple linear layers to perform classification (see fig. 2 ). this architecture is reminiscent of the high expressive power of billion-parameter ngram models @xcite from before the existence of pre-trained language models. by selecting the n-gram vocabulary and the embedding dimension, dans also scale up to billions of parameters. meanwhile, the inference costs are kept low as dans are sparsely-activated. one weakness of dans is that they are restricted in modeling high-level meanings in long-range contexts, as compared to the self-attention operator in transformers. however, recent studies have shown that large pre-trained transformers are rather insensitive to word order @xcite and that they still work well when the learned selfattention is replaced with hard-coded localized attention @xcite or convolution blocks @xcite . taken together, these studies suggest that on some tasks it may be possible to get competitive results without computationally expensive operations such as self-attention. to verify our hypothesis, we use six singlesentence text classification tasks 2 and apply knowledge distillation to dans. we observe that the resulting student models retain 97% of the roberta-large teacher performance on average. we also show that our method falls outside of the pareto frontier of existing methods; compared to a baseline of distilling to a lstm student, our method gives comparable accuracy at less than 1/20 the inference cost (see fig. 1 ). based on our empirical results, we conclude that faster and larger student models provide a valuable benefit over existing methods. we further examine our method (1) with qqp, a sentence-pair task, (2) in privacy-preserving settings (i.e., no access to task-specific data during distillation), and (3) in domain generalization and adaptation settings (i.e., student models are applied and adapted to new data domains), where we find our method continues to bring improvements over non-distillation baselines. ## conclusions & future work we investigated a new way of using knowledge distillation to produce a faster student model by reversing the standard practice of having the student be smaller than the teacher and instead allowed the student to have a large table of sparsely-activated embeddings. this enabled the student model to essentially memorize task-related information that if an alternate architecture were used would have had to be computed. we tested this method on six single-sentence classification tasks with models that were up to 1 billion parameters in size, approximately 3x as big as the roberta-large teacher model, and found that the student model was blazing fast and performed favorably. we hope that our work can lead to further exploration of sparse architectures in knowledge distillation. there are multiple directions for future work, including extending the dan architecture to better support tasks with long range dependencies like natural language inference or multiple inputs like text similarity. additionally, more work is needed to test the idea on non-english languages where n-gram statistics can be different from english. stage 3, we set the batch size to be 32 and select the learning rate from {3e-4, 1e-4, 3e-5}. we train the model for 10 epochs on d train . for all training procedures, we validate the model at the end of each epoch in the case of fine-tuning, or every 100,000 steps in the case of knowledge distillation. we save the best checkpoint based on dev accuracy. due to the analysis nature of this work and the scale of experiments, performance are computed using dev set and based on one single run. a.4 hardware model training. except for the parallel training attempt in inference speed tests. all inference speed tests are done with the batch size of 32. gpu inference is performed with one quadro rtx 8000 gpu, and cpu inference is performed with 56 intel xeon cpu e5-2690 v4 cpus.