Revisiting Knowledge Distillation for Autoregressive Language Models

TL;DR

This work investigates the limits of knowledge distillation for autoregressive language models, revealing that larger teacher models can unexpectedly degrade student performance due to token-level teaching dynamics. It introduces ATKD, a plug-and-play adaptive teaching framework that splits tokens into easy and hard groups based on a token uncertainty measure (UnC) and applies teaching modes accordingly, decoupling TKD and DKD to encourage diverse learning. The approach yields consistent improvements across model families (OPT, Pythia, LLaMA) and KD baselines, mitigating degradation from large teachers and enhancing generalization, with gains reaching up to $+3.04\%$ on average. These results highlight the importance of token-aware supervision in KD and offer a practical path to more efficient distillation of autoregressive LMs. The work also discusses ethics and reproducibility, positioning ATKD as a generalizable improvement to KD pipelines in real-world deployments.

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.

Figures (6)

• Figure 1: Comparisons of different KD methods for distilling the student (OPT-125M). The x-axis denotes the OPT-based teacher sizes, while the y-axis denotes the average performance of students on $\mathcal{S}_{\text{NLG}}$ and $\mathcal{S}_{\text{NLU}}$. The evaluation details are in §\ref{['sec:experiments']}. Notably, ATKD can be combined with various KD methods, and we only report the results of "GKD + ATKD" for ease of illustration.
• Figure 2: Comparisons of different training tokens. The y-axis denotes the average performance of students (OPT-125M) on the evaluated tasks, while the x-axis denotes the sizes of OPT-based teachers.
• Figure 3: Illustration of distributions of UnC ($p^t_{\backslash g_t}$) among different OPT-based teachers on 100 training samples (about 10K tokens). In particular, we use the kernel density estimate for visualizing, where the larger density refers to more tokens.
• Figure 4: Effect of TKD in different training tokens. Here, we report the performance of students distilled with "$\alpha\times$TKD+DKD", where $\alpha$ is varied from 0 to 1. For ease of illustration, we only illustrate the results of using OPT-1.3B and OPT-6.7B as teachers.
• Figure 5: (a) Effect of different ratios (top-$k$) for selecting hard-to-learn tokens, (b) Parameter analysis of $\alpha$ in Eq. \ref{['kd_ours']}, and (c) Comparison of different KD methods that aim to alleviate the problem of performance degrades in larger teachers. We use the Supervised KD as the baseline and report the performance of OPT-125M on $\mathcal{S}_{\text{NLG}}$.