SelecTKD: Selective Token-Weighted Knowledge Distillation for LLMs
Haiduo Huang, Jiangcheng Song, Yadong Zhang, Pengju Ren
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 Speck. 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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on27
- Fast Inference from Transformers via Speculative DecodingYaniv Leviathan, Matan Kalman, Yossi MatiasICML 2023 · 1,472 citations
- WizardLM: Empowering Large Pre-Trained Language Models to Follow Complex InstructionsCan Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng et al.ICLR 2024 · 1,206 citations
- Decoupled Knowledge DistillationBorui Zhao, Quan Cui, Renjie Song, Yiyu Qiu et al.CVPR 2022 · 835 citations
- MetaMath: Bootstrap Your Own Mathematical Questions for Large Language ModelsLonghui Yu, Weisen Jiang, Han Shi, Jincheng Yu et al.ICLR 2024 · 637 citations
- On-Policy Distillation of Language Models: Learning from Self-Generated MistakesRishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk et al.ICLR 2024 · 311 citations
Related papers
- Revisiting Knowledge Distillation for Autoregressive Language ModelsQihuang Zhong, Liang Ding, Li Shen, Juhua Liu et al.ACL 2024
- AdaSPEC: Selective Knowledge Distillation for Efficient Speculative DecodersYuezhou Hu, Jiaxin Guo, Xinyu Feng, Tuo ZhaoNeurIPS 2025 · 9 citations
- Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge DistillationMinsang Kim, Seung Jun BaekICLR 2026 · 15 citations
- Hybrid Policy Distillation for LLMsWenhong Zhu, Ruobing Xie, Rui Wang, Pengfei LiuICML 2026 · 2 citations
- ToDi: Token-wise Distillation via Fine-Grained Divergence ControlSeongryong Jung, Suwan Yoon, DongGeon Kim, Hwanhee LeeEMNLP 2025
