Revisiting Knowledge Distillation for Autoregressive Language Models (2024)

Abstract: Knowledge distillation (KD) is a common approach to compress a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, in the context of autoregressive language models (LMs), we empirically find that larger teacher LMs might dramatically result in a poorer student. In response to this problem, we conduct a series of analyses and reveal that different tokens have different teaching modes, neglecting which will lead to performance degradation. Motivated by this, we propose a simple yet effective adaptive teaching approach (ATKD) to improve the KD. The core of ATKD is to reduce rote learning and make teaching more diverse and flexible. Extensive experiments on 8 LM tasks show that, with the help of ATKD, various baseline KD methods can achieve consistent and significant performance gains (up to +3.04% average score) across all model types and sizes. More encouragingly, ATKD can improve the student model generalization effectively.

## introduction autoregressive language models (lms), such as @xcite , palm @xcite and llama2 @xcite , have achieved great success in a numerous tasks @xcite @xcite . however, with the scaling of model size, the inference and deployment of these lms become more computationally expensive and memory intensive, hindering the development of industrial applications. hence, it is crucial and green to compress these lms and accelerate the inference, while not losing much performance @xcite . to achieve this goal, a common approach is knowledge distillation (kd), which aims to compress a large teacher model by distilling its knowledge into a small student model @xcite , the evaluation details are in §4. notably, atkd can be combined with various kd methods, and we only report the results of "gkd + atkd" for ease of illustration. although a few works aim to investigate this problem and propose to fill the gap, they are mostly studied for vision models @xcite or discriminative language understanding models @xcite , while the autoregressive kd for generative lms is yet to be explored. in this work, we investigate this problem from the perspective of the distillation objective, which is at the core of autoregressive kd. specifically, taking the classical token-level kd objective, i.e., forward kl-divergence, as an example, we first reformulate it as two parts: 1) targetoriented knowledge distillation (tkd), which enforces the student model to learn the target-related information ; 2) diversity-oriented knowledge distillation (dkd), which encourages the student to learn more diverse knowledge from the teacher in the non-target classes. these two parts are tied by a token-wise factor, which reflects the teacher's uncertainty and we denote it as uncertainty coefficient (unc). after reformulating the distillation objective, we conduct a series of preliminary analyses on the popular opt-family @xcite models, and find that: ❶ unc measures the learning difficulties of tokens, where the hard-to-learn ones are more important for kd. based on these observations, we can conclude that different tokens have different teaching modes, and (one of) the limitations of kd comes from the neglect of this principle. to address this limitation, we propose a simple yet effective adaptive teaching method (referred to as atkd) to improve the kd. the core of atkd is to reduce rote learning and make teaching more diverse and flexible. specifically, atkd skips the target-oriented teaching for the (less-informative) easy-to-learn tokens and pays more attention to the diverse learning of hard-to-learn tokens. we evaluate atkd on a variety of lm benchmarks, including 5 language generation tasks and 3 language understanding tasks, upon 3 types of autoregressive lms: opt @xcite , pythia @xcite and llama @xcite . results show that atkd can not only alleviate the problem of performance degradation in larger teachers, but also bring consistent and significant improvements (up to +3.04% average score) into various baseline kd methods among all model types and sizes. moreover, compared to the standard kd, atkd can effectively improve the generalization of distilled students. contributions. to summarize, our contributions are three-fold: (1) our study reveals that different tokens have different teaching modes, neglecting which will cause the sub-optimal distillation performance, especially in larger teachers. (2) we propose a simple yet effective, plug-and-play approach (atkd) to alleviate this problem and improve the quality of teaching. (3) extensive experiments show that atkd outperforms the standard kd with up to +3.04% average gains and improves the student's model generalization effectively. ## rethinking knowledge distillation for autoregressive lms in this section, we first delve into the mechanism of classic knowledge distillation and then present the empirical analyses of this strategy in detail. ## improving knowledge distillation with adaptive teaching modes based on the observations in §2, we recognize that different tokens have different teaching modes, and the side effect (i.e., problem degrades in larger teachers) of kd mainly comes from the neglect of this principle. to this end, we propose to improve the autoregressive kd with adaptive teaching modes (atkd). in this section, we introduce the atkd approach in detail. motivation and overview of atkd. in addition to the empirical findings in §2, our atkd is also inspired by a famous education initiative @xcite , "teach less, learn more", which highlights that reducing rote learning and making education more diverse and flexible can improve the quality of teaching and enhance student learning. intuitively, due to the large capability gap between teacher and student models, target-oriented learning of easy-to-learn tokens may encourage the student to simply mimic the teacher's shallow style but not to learn its dark knowledge @xcite . that is, the student might fall short in generalizing to more tasks, leading to sub-optimal performance. motivated by this, our atkd aims to encourage the students to learn from different perspectives for different tokens. in short, atkd skips the target-oriented teaching for the easy-to-learn tokens, and pays more attention to the learning of diverse knowledge in the hard-to-learn tokens. by doing so, our atkd forces the student to learn more flexible and diverse knowledge, and thus improve overall performance. to achieve this goal, we should first obtain the easy-/hard-to-learn tokens. as mentioned in ❶ of §2.2, unc can effectively measure the learning difficulties of tokens, and we thus use it as a metric to select the easy-/hard-to-learn tokens. specifically, for each mini-batch, we rank the training tokens according to unc and select the top-k foot_4 tokens as hard-to-learn tokens, while the others are easy-tolearn. then, atkd performs the kd processes with adaptive teaching modes as follows: adaptive teaching modes of atkd. as aforementioned, tkd and dkd contribute differently in easy-/hard-to-learn tokens. thus, instead of using a unified teaching mode for all tokens, we use adaptive teaching modes for easy-to-learn and hardto-learn tokens, respectively. specifically, we decouple the tkd and dkd (i.e., dkd will not be suppressed by the unc) to enhance the diverse learning of students. moreover, for the easy-tolearn tokens, considering that the student can easily learn the target-class information, we skip the target-oriented teaching, i.e., removing tkd. on the other hand, both tkd and dkd are used for hard-to-learn tokens, as we empirically found that target-oriented teaching is essential to the learning of hard-to-learn tokens. the learning objectives of different tokens can be formulated as: ## related works recently, autoregressive lms @xcite @xcite have shown their superior performance by solving various nlp tasks in a generative manner. despite their success, they usually suffer from unbearable inference latency @xcite . to this end, several model compression approaches are proposed to reduce the model size and accelerate the inference @xcite @xcite . among these efforts, kd strategy @xcite , which aims at training a smaller student model with the guidance of a teacher model, has attracted great attention recently @xcite @xcite . although these kd methods realize promising performance when distilling (relatively) smaller lms, they might fall short in distilling larger lms (e.g., opt-6.7b) especially when the student is of a small scale. in fact, this phenomenon has been observed in the vision community @xcite and language understanding models @xcite . to alleviate this problem, a few studies including teacher assistant-based @xcite and studentfriendly @xcite @xcite distillation have been recently explored. the above efforts are generally used for vision models or discriminative lms, while the autoregressive kd for generative lms is yet to be explored. to the best of our knowledge, we are the (nearly) first to alleviate the problem of performance degradation in larger autoregressive teacher lms. different from the previous methods that aim to directly bridge the performance gap between teacher and student, we attempt to improve the quality of teaching by exploring and addressing the limitations of existing kd objectives. ## conclusion in this paper, we reveal and address the limitations of kd in compressing the larger autoregressive teachers. based on a series of preliminary analyses, we find that equally adopting the same teaching modes for all tokens is sub-optimal, as learning more target-oriented knowledge of the easyto-learn tokens might lead to overfitting and result in poor performance. to address these limitations, we improve kd with a novel adaptive teaching algorithm. it skips the target-oriented teaching for easy-to-learn tokens and pays more attention to the diverse learning of hard-to-learn tokens. experiments show that our approach consistently and significantly improves distillation performance across all model architectures. in-depth analyses prove that our approach indeed alleviates the problem, and further improves the model generalization.