AD - KD : Attribution-Driven Knowledge Distillation for Language Model Compression (2023)

Abstract: Knowledge distillation has attracted a great deal of interest recently to compress large language models. However, existing knowledge distillation methods suffer from two limitations. First, the student model simply imitates the teacher’s behavior while ignoring the reasoning behind it. Second, these methods usually focus on the transfer of sophisticated model-specific knowledge but overlook data-specific knowledge. In this paper, we present a novel attribution-driven knowledge distillation approach, which explores the token-level rationale behind the teacher model based on Integrated Gradients (IG) and transfers attribution knowledge to the student model. To enhance the knowledge transfer of model reasoning and generalization, we further explore multi-view attribution distillation on all potential decisions of the teacher. Comprehensive experiments are conducted with BERT on the GLUE benchmark. The experimental results demonstrate the superior performance of our approach to several state-of-the-art methods.

## introduction transformer-based pre-trained language models (plms), such as bert @xcite and roberta @xcite , have aroused widespread interest among natural language processing (nlp) researchers in recent years. these language models are first pre-trained on large-scale unlabeled corpora to learn the general representation of language, and then fine-tuned on specific downstream tasks to effectively transfer the general knowledge to target domains. this pre-training and fine-tuning paradigm leads to state-of-the-art performances in various nlp tasks such as natural language understanding. however, with the rapid growth of the model scale, the deployment of large-scale plms becomes challenging, especially in low-resource scenarios. to this end, a variety existing knowledge distillation methods can be divided into three categories, namely responsebased, feature-based, and relation-based @xcite . while response-based methods @xcite directly distill the final output, e.g. probability distribution, from the top of the teacher, feature-based @xcite and relation-based methods @xcite try to align the features from intermediate layers of teacher and student models and minimize the difference. to transfer comprehensive knowledge from the teacher, a common practice is to combine response-based methods with the other two @xcite . however, due to the capacity gap between the teacher and the student, feature-based and relation-based methods may not necessarily bring improvement to response-based methods @xcite . to sum up, existing knowledge distillation methods have two limitations. first, they mainly focus on understanding what the teacher's behavior is, instead of why the teacher behaves like this, hinder-ing the reasoning and generalization ability of the student model. second, they pay more attention to distilling sophisticated model-specific knowledge from intermediate layers but neglect data-specific knowledge, which may contain valuable rationale information to understand how the teacher model arrives at a prediction. to address the above limitations, in this paper we propose a novel attribution-driven knowledge distillation (ad-kd) approach that transfers attribution-based knowledge from the teacher to the student. as shown in figure 1 , the attribution information reflects the importance of different tokens towards the prediction, which contains reasoning knowledge of the model and can be complementary to the soft-label knowledge. by transferring such attribution knowledge, the student is allowed to learn the token-level rationale behind the teacher's behavior and thus generalizes better. specifically, we utilize integrated gradients (ig) @xcite , a well-established gradient-based attribution method, to calculate the importance score of each input token. to reduce the influence of trivial dimensions in the teacher's input embeddings, we further adopt the top-k strategy to filter out dimensions with low attribution scores. the remaining attribution scores are aggregated and normalized to denote the importance of individual tokens. moreover, we extract the attribution knowledge for all possible predictions rather than just the prediction with the highest probability. by transferring the multi-view attribution knowledge, the student learns a more comprehensive understanding of the teacher's soft-label distribution. extensive experiments are conducted with bert @xcite on the glue benchmark @xcite . the experimental results demonstrate the effectiveness and superiority of our approach over several state-of-the-art baselines. furthermore, we show that attribution knowledge from different layers contains different information, while the input layer contains the most prominent attribution knowledge for distillation. to summarize, the main contributions are threefold. first, we propose a novel attribution-driven knowledge distillation framework for language model compression that effectively transfers attribution knowledge from the teacher to the student. second, we extract multi-view attribution knowledge based on model predictions to learn comprehensive reasoning knowledge. third, we systematically validate ad-kd on the glue benchmark and show its superior performance over state-of-the-art baselines. ## conclusion in this paper, we propose ad-kd, a novel knowledge distillation framework for language model compression. unlike other distillation methods, ad-kd investigates the model knowledge from the perspective of input attribution, which is vital yet easy to transfer between the teacher and the student. moreover, top-k method is adopted to obtain noiseless attribution maps among input tokens, and multi-view attribution is conducted for a more comprehensive distillation. to our knowledge, this is the first work that incorporates attribution into knowledge distillation. extensive experiments including ablation studies are carried out to show the effectiveness of ad-kd and its components. with the recent emergence of large language models (llms), gradient-based attribution methods are infeasible due to the unavailable parameters. however, the idea of ad-kd can still be potentially extended to these black-box models by using occlusion-based attribution or using chainof-thoughts @xcite as the rationale for distillation. we will leave it to future work.