LLM-Oriented Token-Adaptive Knowledge Distillation

TL;DR

AdaKD introduces token-level adaptability into knowledge distillation for LLMs by coupling a Loss-Driven Adaptive Token Focusing module with an Inverse Difficulty Temperature Scaling strategy, both guided by a unified token-difficulty metric based on Hellinger distance. LATF dynamically selects the most informative tokens to distill, while IDTS assigns per-token temperatures that sharpen hard-token learning and smooth easy-token distributions for better generalization. The approach is plug-and-play, yielding consistent improvements across multiple baselines and model families, with ablations confirming the necessity and synergy of its components. Empirical results on instruction-following benchmarks demonstrate notable performance gains and robustness, while analyses show favorable gradient dynamics and negligible efficiency costs. This work offers a practical framework for token-aware KD that adapts to the student’s learning state in real time, enhancing distillation effectiveness in diverse LLM settings.

Abstract

Knowledge distillation (KD) is a key technique for compressing large-scale language models (LLMs), yet prevailing logit-based methods typically employ static strategies that are misaligned with the dynamic learning process of student models. These methods typically treat all tokens indiscriminately and apply a single, fixed temperature, resulting in suboptimal knowledge transfer. To address these limitations, we propose LLM-Oriented Token-Adaptive Knowledge Distillation (AdaKD), a novel framework that adapts the distillation process to the real-time learning state of each token. AdaKD consists of two synergistic modules driven by a unified token difficulty metric. First, our Loss-Driven Adaptive Token Focusing (LATF) module dynamically adjusts the distillation focus by monitoring the student's learning stability, concentrating computational resources on the most valuable tokens at each training phase. Second, we introduce Inverse Difficulty Temperature Scaling (IDTS), a counterintuitive yet effective token-level temperature strategy. It employs low temperatures for difficult tokens for targeted error correction, and high temperatures for easy tokens to encourage students to learn from the teacher's complete and smooth output distribution, thereby enhancing generalization. As a plug-and-play framework, AdaKD can consistently improve the performance of various distillation methods on multiple model architectures and benchmarks.

Figures (4)

• Figure 1: Analysis of token difficulty and gradient dynamics. Tokens are grouped into Hard, Mid, and Easy based on difficulty (Hellinger distance). (a) Evolution of token difficulty across training stages. (b) Cosine similarity of each token group's gradient with the SFT gradient. (c) Each group's gradient norm percentage and its cosine similarity with the total batch gradient.
• Figure 2: Illustration of the AdaKD framework. The bar charts visualize simplified teacher (blue) and student (purple) probability distributions. The top charts depict the initial learning gaps for "hard" and "mid-difficulty" tokens. After the LATF module filters tokens based on difficulty calculated via indicator, the IDTS module (bottom) applies low temperature to hard tokens for a sharp, corrective signal, and high temperature to easier tokens for a smoother distribution that enhances generalization.
• Figure 3: From left to right: (a) The loss and sample ratio of our adaptive LATF during training. (b) A comparison of loss curves for fixed scheduling strategies. (c) The effect of IDTS modulation intensity $c$ on normalized scores for individual datasets. (d) The average ROUGE-L score across datasets, showing the optimal choice for $c$.
• Figure 4: These histograms display the distribution of token counts (y-axis) across different metrics. The rows compare the model's state 'Before Training' (top) with 'Late in Training' (bottom). The columns, from left to right, show the distributions for assigned temperature, student’s output entropy before IDTS, and entropy after IDTS. Tokens are categorized into 'hard' (blue) and 'easy' (orange) groups, with dashed vertical lines indicating their respective means.