SelecTKD: Selective Token-Weighted Knowledge Distillation for LLMs

TL;DR

SelecTKD introduces a selective, token-weighted knowledge distillation method that improves the training of compact LLMs by focusing on high-confidence tokens, leading to state-of-the-art results without architectural changes.

Contribution

The paper proposes SelecTKD, a novel selective distillation framework that dynamically chooses tokens for supervision, enhancing LLM compression efficiency and stability.

Findings

Consistently improves baseline models across multiple tasks.

Achieves state-of-the-art results for small models.

Works with on- and off-policy data without architectural changes.

Abstract

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.