Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation
Proposed TSD-KD improves reasoning by combining preference-based indirect and confidence-based direct distillation, outperforming baselines by up to 54.4%.
Key Findings
Methodology
TSD-KD integrates dual mechanisms: preference ranking provides subtle, sequence-level feedback guiding early reasoning paths; confidence-based selective KL divergence focuses on uncertain tokens with low student confidence; entropy regularization enhances model confidence on critical tokens. The approach emphasizes minimal teacher intervention, promoting student autonomy. It employs top-k candidate responses, re-ranked by teacher preferences via Plackett-Luce models, and applies entropy minimization on high-uncertainty tokens, all within an end-to-end training framework. This targeted, student-centric design effectively improves reasoning performance across multiple benchmarks.
Key Results
- TSD-KD achieves up to 54.4% accuracy improvement over baseline students, with the largest gains on GSM8K, MATH, and MMLU-Pro-Math datasets. Notably, in some tasks, the student surpasses its teacher by 20.3%, demonstrating enhanced reasoning capabilities.
- On 10 reasoning benchmarks, TSD-KD outperforms recent state-of-the-art methods like GKD and DistiLLM, with significant margins, especially in complex tasks requiring multi-step reasoning.
- Ablation studies confirm that preference-guided early path correction and confidence-based token distillation are key to performance gains, validating the effectiveness of the combined approach.
Significance
This work addresses the core challenge of enabling smaller models to perform complex reasoning without extensive supervision. By focusing on key tokens and early reasoning steps, TSD-KD reduces distribution mismatch and improves generalization. Its student-centric design promotes autonomous reasoning, making it suitable for deployment in resource-constrained environments like edge devices. The approach paves the way for more efficient, accurate AI systems capable of complex problem-solving, impacting education, research, and industry.
Technical Contribution
The paper introduces a novel dual distillation framework combining preference-based indirect feedback with confidence-driven direct distillation, integrated with entropy regularization. It innovatively leverages token importance via entropy, employs Plackett-Luce models for preference alignment, and selectively distills critical tokens based on uncertainty gaps. These mechanisms collectively enable the student model to learn more effectively, with improved reasoning and robustness, surpassing existing methods in accuracy and efficiency.
Novelty
This is the first work to unify preference ranking and confidence-based token selection within a single distillation framework tailored for reasoning tasks. Unlike prior methods that either mimic teacher distributions or rely solely on explicit supervision, TSD-KD emphasizes student autonomy by guiding early reasoning paths and focusing on uncertain tokens. Its dual mechanism and entropy regularization represent a significant departure from traditional, monolithic distillation approaches, setting new standards for model self-improvement.
Limitations
- The reliance on accurate preference ranking and entropy estimation may limit performance in noisy or highly complex scenarios. Computational overhead remains a concern for very large models or long reasoning chains.
- Parameter sensitivity, especially in the choice of entropy thresholds and gating functions, could affect robustness across tasks. Further research is needed to adaptively tune these hyperparameters.
- While effective on benchmark datasets, real-world deployment may require additional optimization for efficiency and scalability.
Future Work
Future directions include extending the framework to multi-modal reasoning tasks, integrating reinforcement learning to dynamically refine preferences, and optimizing computational efficiency for real-time applications. Exploring adaptive thresholds for entropy and preference signals could further enhance robustness. Additionally, applying TSD-KD to larger models and diverse domains will validate its generalizability and practical impact.
AI Executive Summary
This paper introduces Token-Selective Dual Knowledge Distillation (TSD-KD), a novel framework designed to enhance the reasoning capabilities of smaller language models. Traditional knowledge distillation methods often force models to mimic entire output distributions, which can overwhelm limited-capacity students and cause distribution mismatch, especially in complex reasoning tasks. TSD-KD addresses this by focusing on key tokens that are critical for reasoning, employing a dual mechanism: indirect distillation guided by teacher preferences and direct distillation based on confidence gaps. The indirect component leverages preference ranking over candidate responses, providing subtle, sequence-level feedback that guides the early stages of reasoning. The direct component selectively distills tokens where the student exhibits high uncertainty but the teacher is confident, using KL divergence to fine-tune these critical points. Entropy regularization further stabilizes learning by boosting confidence on uncertain tokens, reducing ambiguity. Extensive experiments on 10 benchmarks, including GSM8K, MATH, and MMLU-Pro-Math, demonstrate that TSD-KD achieves state-of-the-art results, with accuracy improvements up to 54.4% over baseline students and surpassing some teacher models by over 20%. These results highlight the method’s ability to foster autonomous, robust reasoning in compact models, with broad implications for AI deployment in resource-constrained environments. The framework’s modular design and targeted feedback mechanisms open new avenues for model compression, self-improving AI, and scalable reasoning systems.
Deep Analysis
Background
Recent advances in large language models (LLMs) like GPT-4, PaLM, and Claude have revolutionized NLP, especially in reasoning tasks. Knowledge distillation (KD) emerged as a key technique to compress these models, enabling deployment on edge devices. Early KD methods focused on mimicking output logits or probability distributions, but these often led to issues like distribution mismatch and poor generalization. On-policy distillation, where models learn from their own generated outputs, partially alleviated these problems but still struggled with complex reasoning paths. Recent works such as GKD and DistiLLM introduced preference-based and reinforcement learning approaches, yet they lacked targeted focus on critical reasoning tokens. This paper builds upon these foundations, proposing a more nuanced, student-centric approach that emphasizes early reasoning steps and key tokens, aiming to improve reasoning accuracy and robustness.
Core Problem
Despite progress, current KD methods face significant challenges in complex reasoning tasks. They tend to treat all tokens equally, ignoring the importance of early decision points, which are often high-uncertainty tokens. Overly aggressive imitation of teacher distributions can lead to model overfitting and reduced generalization, especially when the student model capacity is limited. Moreover, existing approaches lack mechanisms to focus learning on critical reasoning junctures, resulting in suboptimal reasoning paths and reduced interpretability. Addressing these issues requires a targeted, flexible distillation strategy that emphasizes key tokens and early reasoning stages, allowing the student to develop autonomous, reliable reasoning capabilities while maintaining efficiency.
Innovation
The core innovations include: 1) Preference-guided indirect distillation, where the teacher ranks candidate responses proposed by the student, providing sequence-level subtle feedback that guides early reasoning; 2) Confidence-based direct distillation, focusing on tokens where the student is uncertain but the teacher is confident, using a gating function to selectively apply KL divergence; 3) Entropy regularization on critical tokens, reducing uncertainty and improving confidence. These mechanisms collectively enable the student to learn more effectively by focusing on pivotal reasoning points, reducing unnecessary distribution matching, and fostering autonomous reasoning. Unlike prior work that uniformly distills entire distributions, TSD-KD emphasizes targeted, adaptive guidance, leading to superior reasoning performance.
Methodology
- ��学生模型在推理路径早期提出候选响应，生成top-k候选Token；•教师对候选Token进行偏好排序，形成偏好序列；•采用Plackett-Luce模型计算偏好序列的概率，作为微调信号；•在推理路径的前段（opener）部分，基于累积熵选择关键Token，进行偏好排序和微调；•在后续推理中，依据学生置信度与教师差异，选择性地应用KL散度进行关键Token的蒸馏；•引入entropy正则化，针对高不确定性Token进行信心增强；•整体训练目标结合间接蒸馏、直接蒸馏和entropy正则化，优化学生推理能力。
Experiments
采用GSM8K、MATH、MMLU-Pro-Math等10个推理基准，比较TSD-KD与多种基线（如Supervised-KD、GKD、DistiLLM等）。模型使用Qwen 14B（教师）与1.5B（学生），调优超参数如k=10、β=0.9。评估指标为准确率，特别关注复杂推理任务的提升。通过消融实验验证偏好排序和置信度引导的贡献，分析不同参数设置对性能的影响。实验证明，TSD-KD在多项任务中实现了最高性能，显著优于对比方法。
Results
在10个推理任务中，TSD-KD最高提升54.4%，在MATH任务中超越第二名40.3%，甚至在部分任务中超越教师模型20.3%。在复杂推理任务如GSM8K、MMLU-Pro-Math中表现尤为突出，验证了其推理能力的增强。消融分析显示偏好排序引导的早期路径优化和置信度差异的关键Token选择性蒸馏是性能提升的关键因素。这些结果表明，TSD-KD不仅提升了模型的准确性，也增强了其自主推理和泛化能力。
Applications
该方法适用于需要高推理能力的AI系统，如教育辅导、科研辅助、智能问答等。通过模型压缩，能在边缘设备上实现高效推理，降低部署成本。未来，结合多模态信息和强化学习，有望推动AI在更复杂场景中的自主推理和决策能力。
Limitations & Outlook
方法在极端复杂或噪声较多的任务中表现有限，偏好排序和置信度估计依赖于模型的准确性。计算成本较高，尤其在长推理链中需优化效率。此外，参数调优对不同任务敏感，泛化到其他领域仍需验证。
Plain Language Accessible to non-experts
想象你在教一个学生解数学题。传统方法就像老师把所有解题步骤都告诉学生，学生照做就行了，但有时信息太多，反而让学生迷失方向。这个新方法像老师只告诉学生关键的转折点，让学生自己尝试推理，然后老师根据学生的表现给出微弱的建议，帮助他们在关键点做得更好。这样，学生既能自主思考，又能在重要节点得到指导，逐步变得更聪明。整个过程就像在训练一只聪明的宠物，让它学会自己思考问题，而不是死记硬背答案。
ELI14 Explained like you're 14
想象你在参加一个智力比赛，题目很难，需要你自己想办法解答。以前的做法就像老师把所有答案都告诉你，你只要照着做就行了，但这样你学不到怎么自己思考。现在，有个新方法，老师只在你关键的思考点给一些提示，比如“这个步骤很重要”，让你自己试着推理，然后老师根据你做的内容，给你一些微弱的建议，帮助你在关键地方做得更好。这就像在训练你变得更聪明，不仅仅是记住答案，而是真正学会了怎么思考。这样，你不仅能解决这道题，还能学会自己面对类似的问题，变得更厉害！
Glossary
Knowledge Distillation
A model compression technique where a smaller model learns to mimic a larger, more complex model's output distribution.
Describes the core method for transferring reasoning abilities in the paper.
On-policy KD
A distillation approach where the student learns from its own generated outputs, reducing distribution mismatch.
Emphasized as a key aspect of the proposed framework.
Preference Ranking
A process where the teacher ranks candidate responses proposed by the student, guiding the student towards better reasoning paths.
Used to provide subtle, sequence-level feedback.
Entropy Regularization
A technique to minimize the output entropy of the model on critical tokens, increasing confidence and reducing uncertainty.
Applied to improve reasoning stability.
KL Divergence
A measure of difference between two probability distributions, used here to align student and teacher output distributions.
Applied selectively on uncertain tokens.
Open Questions Unanswered questions from this research
- 1 如何在极端复杂或噪声较多的任务中保持偏好排序和置信度估计的准确性仍是挑战。
- 2 在超大规模模型或超长推理链中，计算效率和资源消耗需要进一步优化。
Abstract
Knowledge Distillation (KD) can transfer the reasoning abilities of large models to smaller ones, which can reduce the costs to generate Chain-of-Thoughts for reasoning tasks. KD methods typically ask the student to mimic the teacher's distribution over the entire output. However, a student with limited capacity can be overwhelmed by such extensive supervision causing a distribution mismatch, especially in complex reasoning tasks. We propose Token-Selective Dual Knowledge Distillation (TSD-KD), a framework for student-centric distillation. TSD-KD focuses on distilling important tokens for reasoning and encourages the student to explain reasoning in its own words. TSD-KD combines indirect and direct distillation. Indirect distillation uses a weak form of feedback based on preference ranking. The student proposes candidate responses generated on its own; the teacher re-ranks those candidates as indirect feedback without enforcing its entire distribution. Direct distillation uses distribution matching; however, it selectively distills tokens based on the relative confidence between teacher and student. Finally, we add entropy regularization to maintain the student's confidence during distillation. Overall, our method provides the student with targeted and indirect feedback to support its own reasoning process and to facilitate self-improvement. The experiments show the state-of-the-art performance of TSD-KD on 10 challenging reasoning benchmarks, outperforming the baseline and runner-up in accuracy by up to 54.4\% and 40.3\%, respectively. Notably, a student trained by TSD-KD even outperformed its own teacher model in four cases by up to 20.3\%. The source code is available at https://github.com/kmswin1/TSD-KD.
References (20)
How (not) to Train your Generative Model: Scheduled Sampling, Likelihood, Adversary?
Ferenc Huszár
Distilling the Knowledge in a Neural Network
Geoffrey E. Hinton, O. Vinyals, J. Dean
Gemma 2: Improving Open Language Models at a Practical Size
Gemma Team Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa et al.
On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou et al.
DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
DeepSeek-AI
The Unreasonable Effectiveness of Entropy Minimization in LLM Reasoning
Shivam Agarwal, Shivam Agarwal, Zimin Zhang et al.
Curriculum learning
Yoshua Bengio, J. Louradour, R. Collobert et al.
Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters
C. Snell, Jaehoon Lee, Kelvin Xu et al.
Crowdsourcing Multiple Choice Science Questions
Johannes Welbl, Nelson F. Liu, Matt Gardner
Paper
N. Cambridge
Program Synthesis with Large Language Models
Jacob Austin, Augustus Odena, Maxwell Nye et al.
Graph of Thoughts: Solving Elaborate Problems with Large Language Models
Maciej Besta, Nils Blach, Aleš Kubíček et al.
Advancing LLM Reasoning Generalists with Preference Trees
Lifan Yuan, Ganqu Cui, Hanbin Wang et al.
GSM-Plus: A Comprehensive Benchmark for Evaluating the Robustness of LLMs as Mathematical Problem Solvers
Qintong Li, Leyang Cui, Xueliang Zhao et al.
Autoregressive Knowledge Distillation through Imitation Learning
Alexander Lin, Jeremy Wohlwend, Howard Chen et al.
Training Verifiers to Solve Math Word Problems
K. Cobbe, V. Kosaraju, Mo Bavarian et al.
Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning
Shenzhi Wang, Le Yu, Chang Gao et al.
Training language models to follow instructions with human feedback
Long Ouyang, Jeff Wu, Xu Jiang et al.
Large Language Models are Zero-Shot Reasoners
Takeshi Kojima, S. Gu, Machel Reid et al.