A Comparative Analysis of Task-Agnostic Distillation Methods for Compressing Transformer Language Models (2023)

Abstract: Large language models have become a vital component in modern NLP, achieving state of the art performance in a variety of tasks. However, they are often inefficient for real-world deployment due to their expensive inference costs. Knowledge distillation is a promising technique to improve their efficiency while retaining most of their effectiveness. In this paper, we reproduce, compare and analyze several representative methods for task-agnostic (general-purpose) distillation of Transformer language models. Our target of study includes Output Distribution (OD) transfer, Hidden State (HS) transfer with various layer mapping strategies, and Multi-Head Attention (MHA) transfer based on MiniLMv2. Through our extensive experiments, we study the effectiveness of each method for various student architectures in both monolingual (English) and multilingual settings. Overall, we show that MHA transfer based on MiniLMv2 is generally the best option for distillation and explain the potential reasons behind its success. Moreover, we show that HS transfer remains as a competitive baseline, especially under a sophisticated layer mapping strategy, while OD transfer consistently lags behind other approaches. Findings from this study helped us deploy efficient yet effective student models for latency-critical applications.

## introduction large language models have become a crucial component in modern nlp. they have achieved exceptional performance on various downstream tasks @xcite @xcite and their capability shows consistent improvement with more compute, data, and model parameters @xcite @xcite . on the downside, it is becoming increasingly difficult to deploy such models in real-world environments due to their inefficiency, i.e. high computation, memory, latency and storage costs @xcite . knowledge distillation @xcite is a promising technique to overcome this challenge by transferring the knowledge of the original model (teacher) to a smaller, more efficient model @xcite . this can be conducted in either task-specific @xcite or task-agnostic manner @xcite . the latter only requires distilling a single generalpurpose student which can be directly finetuned on any downstream task. due to its high convenience, we focus on this latter approach in this study. in recent years, there have been various methods proposed for task-agnostic distillation of transformer language models. the aim of this paper is to reproduce, compare and analyze the most representative methods in this area. we generally focus on the architecture-agnostic distillation which imposes no or minimal restriction on the student architecture foot_0 : the representative methods include output distribution (od) transfer @xcite , hidden state (hs) transfer based on linear mapping @xcite and multi-head attention (mha) transfer based on minilmv2 @xcite . for hs transfer, the layer mapping strategy between teacher and student layers plays a significant role in overall performance, however, the optimal strategy remains unknown or controversial @xcite @xcite . therefore, we explore a diverse range of strategies to empirically evaluate each technique. for mha transfer, the minilmv2 approach has been shown to achieve state-of-the-art performance, however, there is relatively little understanding behind its success. therefore, we develop a novel variant named directminilm which is useful for understanding the effectiveness behind minilmv2 both theoretically and empirically. in contrast to most previous studies, all methods are reproduced on a single unified codebase for fair and consistent comparison. we also conduct distillation on 4 different student architectures, reducing the model size in various dimensions to fit various parameter and latency budgets. finally, all experiments are conducted on both monolingual and multilingual settings, distilled from open-source bert @xcite and in-house xlm-roberta @xcite , respectively. through our extensive experiments, we critically analyze the effectiveness of each distillation method and provide practical advice for both researchers and practitioners working in this area. in summary, our key findings are: ## transformer language models first, we briefly review the standard architecture of transformer language models @xcite . a transformer consists of a stack of l transformer layers, where each layer comprises two sub-layers: a multi-head attention (mha) layer followed by a fully connected feed-forward (ff) layer (figure 1 , (a) ). formally, let x denote the input sequence, d h the hidden state size, and h @xmath0. given h i , the mha layer first computes the query, key, and value mappings q i,a , k i,a , v i,a for each attention head @xmath1, a h ], which are combined to obtain the attention head output o i,a : here, d k denotes the attention head size (typically set to d h a h ) and w q,i,a , w k,i,a , w v,i,@xmath2. the output of the mha layer is the concatenation of o i,a , namely throughout this paper, we assume that both the student and teacher are transformer language models with l s and l t layers, respectively. ## distillation methods next, we introduce the representative task-agnostic distillation methods illustrated in figure 1 , (b-d ). for multi-head attention (mha) transfer, we consider two approaches: minilmv2 and its novel variant directminilm. for a survey of advanced methods and topics we could not cover in this study, please refer to appendix a. ## experimental setup we explore the task-agnostic knowledge distillation methods under two settings: foot_3 1. monolingual distillation: we train english students using the open-source bert @xcite as the teacher. these models are distilled on the same corpus used for pretraining bert, i.e., english wikipedia @xcite and bookcorpus @xcite . 2. multilingual distillation: we train multilingual students using our in-house xlm-roberta @xcite as the teacher, and distill on the cc100 dataset @xcite , which consists of data in more than 100 languages. we only use a small subset of the corpus to conduct our experiments within a reasonable computation budget while maintaining the language-wise distribution. in both settings, we use the base (12 layer) architecture for the teacher, as shown in student models to conduct a strong comparison of the representative knowledge distillation methods, we train 4 students of varying architectures and latency/parameter budgets. a summary of the student architectures, with their parameters and latency of inference, are shown in our largest student is a 6 layer model that follows the same architecture as distilbert @xcite . we also use the 6 layer model used in @xcite , which has a smaller hidden size than the teacher. our smaller 4 and 3 layer students were obtained as recommendations from a neural architecture search process @xcite to find good student architectures for distillation from the xlm-roberta teacher, conditioned to minimize the latency on cpu. please refer to appendix c for more details. layer mapping strategies the layer mapping strategy φ is a parameter that needs to be considered for both hs and mha transfer. for hs transfer, we explore the following three settings: for mha transfer, we always take the single mapping strategy and distill a single teacher layer into the last student layer, following @xcite . specifically, we experiment with the last three teacher layers as a choice for distillation for both minilmv2 and directminilm. while od transfer can be conducted from scratch, we found this converges slowly and does not perform competitively. 6 therefore, we take the style of multi-stage distillation @xcite and conduct od transfer after hs transfer, using the distilled checkpoint from hs transfer. this approach converges much faster with better final performance, hence we take this approach as the representative od transfer method. 6 our 6l monolingual student takes 49 hours on 30 v100 gpus to reach acceptable performance, while the same model achieves better scores in only 10.5 hours when initialized from the hs transferred checkpoint. ## evaluation and results for both our monolingual and multilingual models, we measure performance on the english glue benchmark @xcite and report the average score of all tasks (without cola foot_5 ). for multilingual models, we provide evaluations on the xnli dataset @xcite , a set of inference tasks which evaluates the model's performance on 15 languages after being finetuned on only english training data. we report the average score of all languages for xnli. ## conclusion this study critically analyzes the representative methods for task-agnostic distillation of language models. specifically, we compare output distribution (od), hidden state (hs), and multi-head attention (mha) transfer for different student architectures, language settings, and layer mapping strategies. through our extensive experiments, we show that mha transfer based on minilmv2 is the best option across many settings, followed by hs transfer with sophisticated 1-to-n mapping strategies. meanwhile, we did not find od transfer to be an effective alternative. finally, we propose di-rectminilm to demistify the precise advantage of the indirect (i.e. relation matrix based) optimization technique proposed in minilmv2. overall, we hope this study will be a useful guide for both researchers and practitioners working in this area. on empirical methods in natural language processing (emnlp), pages 1016-1021, online. association for computational linguistics. yimeng wu, mehdi rezagholizadeh, abbas ghaddar, md akmal haidar, and ali ghodsi. 2021. universal-kd: attention-based output-grounded intermediate layer knowledge distillation. in proceedings of the 2021 conference on empirical methods in natural language processing, pages 7649-7661, online and punta cana, dominican republic. association for computational linguistics. mengzhou xia, zexuan zhong, and danqi chen. 2022. structured pruning learns compact and accurate models. in proceedings of the 60th annual meeting of the association for computational linguistics (volume 1: long papers), pages 1513-1528, dublin, ireland. association for computational linguistics. ji xin, raphael tang, yaoliang yu, and jimmy lin. 2021. berxit: early exiting for bert with better fine-tuning and extension to regression. in proceedings of the 16th conference of the european chapter of the association for computational linguistics: main volume, pages 91-104, online. association for computational linguistics. canwen xu and julian mcauley. 2023. a survey on model compression and acceleration for pretrained language models. in proceedings of the aaai conference on artificial intelligence, volume 37, pages 10566-10575. canwen xu, wangchunshu zhou, tao ge, furu wei, and ming zhou. 2020. bert-of-theseus: compressing bert by progressive module replacing. in proceedings of the 2020 conference on empirical methods in natural language processing (emnlp), pages 7859-7869, online. association for computational linguistics. dongkuan (dk) xu, subhabrata mukherjee, xiaodong liu, debadeepta dey, wenhui wang, xiang zhang, ahmed awadallah, and jianfeng gao. 2022. fewshot task-agnostic neural architecture search for distilling large language models. in advances in neural information processing systems, volume 35, pages 28644-28656. curran associates, inc. ofir zafrir, guy boudoukh, peter izsak, and moshe wasserblat. 2019. q8bert: quantized 8bit bert. in 2019 fifth workshop on energy efficient machine learning and cognitive computing-neurips edition (emc2-nips), pages 36-39. ieee. minjia zhang, niranjan uma naresh, and yuxiong he. 2022. adversarial data augmentation for taskspecific knowledge distillation of pre-trained transformers. proceedings of the aaai conference on artificial intelligence, 36(10):11685-11693. wangchunshu zhou, canwen xu, and julian mcauley. 2022. bert learns to teach: knowledge distillation with meta learning. in proceedings of the 60th annual meeting of the association for computational linguistics (volume 1: long papers), pages 7037-7049, dublin, ireland. association for computational linguistics. yukun zhu, ryan kiros, rich zemel, ruslan salakhutdinov, raquel urtasun, antonio torralba, and sanja fidler. 2015. aligning books and movies: towards story-like visual explanations by watching movies and reading books. in 2015 ieee international conference on computer vision (iccv), pages 19-27.