SelecTKD: Selective Token-Weighted Knowledge Distillation for LLMs
This addresses the issue of inefficient compression for LLMs, offering a plug-and-play solution that improves performance in domains like instruction following and code generation, though it is incremental as it builds on existing distillation methods.
The paper tackled the problem of noisy supervision in knowledge distillation for compressing Large Language Models by introducing SelecTKD, a selective token-weighted framework that focuses on where to apply learning, achieving state-of-the-art results for small models across multiple tasks without architectural changes.
Knowledge distillation (KD) is a standard route to compress Large Language Models (LLMs) into compact students, yet most pipelines uniformly apply token-wise loss regardless of teacher confidence. This indiscriminate supervision amplifies noisy, high-entropy signals and is especially harmful under large teacher-student capacity gaps. We introduce SelecTKD, a plug-and-play Selective Token-Weighted distillation framework that shifts the focus from "how to measure divergence" to "where to apply learning". At each step, the student proposes tokens that are verified by the teacher through a robust propose-and-verify procedure with two variants: greedy Top-k and non-greedy Spec-k. Accepted tokens receive full loss, while rejected tokens are masked or down-weighted. This objective-agnostic design works with on- and off-policy data, induces an implicit curriculum quantified by Token Acceptance Rate (TAR), and stabilizes optimization. Across instruction following, mathematical reasoning, code generation, and a VLM setting, SelecTKD consistently improves strong baselines and achieves state-of-the-art results for small models without architectural changes or extra reference models.