Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation

TL;DR

提出Token-Selective Dual Knowledge Distillation(TSD-KD),在10个基准上提升准确率达54.4%。

cs.CL 🔴 高级 2026-02-25 13 引用 60 次浏览
Minsang Kim Seung Jun Baek
知识蒸馏 推理任务 模型压缩 自我改进 深度学习

核心发现

方法论

TSD-KD结合间接和直接蒸馏,前者通过偏好排序提供微弱反馈,后者基于置信度差异选择性蒸馏关键Token。采用entropy正则化保持学生模型信心,重点在推理关键Token上进行目标导向的知识传递,提升模型推理能力。具体机制包括:•利用偏好排序对学生提出的候选响应进行重排,提供序列级微调信号;•在关键Token上应用KL散度进行选择性蒸馏,依据学生置信度与教师差异调整;•引入entropy正则化,增强学生在关键Token上的信心,抑制不确定性。该框架在多项推理任务上实现了优异性能,显著优于传统KD和现有方法。

关键结果

  • 在10个推理基准中,TSD-KD最高提升准确率达54.4%,在MATH任务中超越第二名40.3%,甚至在部分任务中超越其教师模型20.3%。
  • 在GSM8K、MMLU-Pro-Math等复杂任务中表现优异,显著优于对比方法如Supervised-KD、GKD等,验证其推理能力提升。
  • 通过消融实验确认:偏好排序引导的间接蒸馏在早期推理路径中效果最佳,置信度差异导向的直接蒸馏增强关键Token预测精度。

研究意义

该研究突破了传统知识蒸馏的瓶颈,通过学生自主推理和微调,有效缓解分布偏移问题,提升模型在复杂推理任务中的泛化能力。其创新机制为大模型压缩与推理能力迁移提供新思路,推动AI在教育、科研等领域的应用落地,具有深远影响。

技术贡献

提出Token-Selective Dual Knowledge Distillation(TSD-KD),融合偏好排序引导的间接蒸馏与置信度差异导向的直接蒸馏,结合entropy正则化,显著改善推理模型的自主性和准确性。创新点在于:•在推理路径早期引入偏好排序,减少分布偏移;•利用置信度差异实现关键Token的选择性蒸馏,提升效率;•引入entropy正则化,增强模型信心,减少不确定性。这一框架在多个推理任务中实现了SOTA性能,甚至超越教师模型。

新颖性

本研究首次提出结合偏好排序与置信度差异的Token-Selective蒸馏策略,强调学生自主推理与微调,区别于传统全分布匹配或单一偏好引导方法。创新在于:•采用偏好排序提供序列级微调信号,避免分布偏移;•基于置信度差异实现关键Token的选择性蒸馏,提升效率;•引入entropy正则化,增强模型信心,提升推理能力。这些创新共同推动模型自主推理与泛化能力的提升。

局限性

  • 方法依赖于偏好排序和置信度差异的有效估计,在极端复杂或噪声较多的任务中可能表现不佳。
  • 在极大规模模型或极端长推理链中,计算成本和效率仍需优化,存在一定局限。
  • 对不同任务的参数调优较为敏感,泛化到其他领域还需进一步验证。

未来方向

未来将探索多模态推理任务中的偏好排序与置信度引导策略,结合强化学习优化偏好模型,提升模型自主推理能力。同时,考虑模型压缩与推理速度的平衡,推动其在边缘设备上的应用落地。

AI 总览摘要

本研究提出Token-Selective Dual Knowledge Distillation(TSD-KD),旨在提升大规模语言模型的推理能力。传统知识蒸馏方法多强调模仿教师的整体输出分布,容易导致学生模型在复杂推理任务中出现分布偏移与信息过载。为此,TSD-KD引入了两大创新机制:一是通过偏好排序提供微弱但序列级的反馈,指导学生在推理路径早期做出更合理的决策;二是利用置信度差异选择性蒸馏关键Token,强化学生对难点的理解。结合entropy正则化,进一步提升模型在关键Token上的信心,减少不确定性。该框架在包括GSM8K、MATH、MMLU-Pro-Math等10个推理基准上表现优异,最高提升达54.4%,在部分任务中甚至超越教师模型20.3%。实验结果验证了TSD-KD在模型自主推理、泛化能力和效率方面的显著优势,为模型压缩和推理能力迁移提供了新思路。未来工作将聚焦多模态推理和边缘设备部署,推动AI推理技术的广泛应用。

深度分析

研究背景

近年来,深度学习模型在自然语言处理中的表现持续提升,尤其是大型预训练模型如GPT、BERT等在推理任务中展现出强大能力。知识蒸馏作为模型压缩和能力迁移的重要技术,已成为研究热点。早期方法如Hinton等提出的软标签蒸馏,主要模仿教师的输出分布,提升小模型性能。近年来,强调模型自主推理的on-policy蒸馏逐渐兴起,旨在缓解分布偏移问题。然而,现有方法在复杂推理任务中仍面临推理路径不稳定、信息过载等挑战。本论文在此基础上,提出Token-Selective Dual Knowledge Distillation(TSD-KD),通过偏好排序和置信度差异引导,强化学生模型的自主推理能力,弥补现有技术的不足。

核心问题

传统知识蒸馏多强调模仿教师的整体输出,忽视推理路径中的关键Token,导致学生在复杂推理任务中表现不佳。现有on-policy蒸馏虽缓解了分布偏移,但过度干预会抑制学生自主推理,限制模型泛化。如何在保证信息传递效率的同时,提升学生自主推理能力,成为亟待解决的问题。此外,模型在早期推理路径中的不确定性和偏差,严重影响最终推理结果的准确性。解决这些问题,需设计更具目标导向和微调能力的蒸馏策略。

核心创新

本研究的核心创新包括:1)引入偏好排序机制,利用教师对学生候选响应的偏好提供序列级微调信号,有效引导推理路径;2)基于置信度差异实现选择性蒸馏,聚焦于学生不确定但教师自信的关键Token,提高效率和效果;3)结合entropy正则化,增强模型在关键Token上的信心,减少推理中的不确定性。这些创新区别于传统全分布匹配或单一偏好引导方法,强调学生自主推理与微调的结合,显著提升推理性能。

方法详解

  • ��利用学生模型在推理路径早期提出候选响应,生成top-k候选Token;•教师对候选Token进行偏好排序,形成偏好序列;•采用Plackett-Luce模型计算偏好序列的概率,作为微调信号;•在推理路径的前段(opener)部分,基于累积熵选择关键Token,进行偏好排序和微调;•在后续推理中,依据学生置信度与教师差异,选择性地应用KL散度进行关键Token的蒸馏;•引入entropy正则化,针对高不确定性Token进行信心增强;•整体训练目标结合间接蒸馏、直接蒸馏和entropy正则化,优化学生推理能力。

实验设计

采用GSM8K、MATH、MMLU-Pro-Math等10个推理基准,比较TSD-KD与多种基线(如Supervised-KD、GKD、DistiLLM等)。模型使用Qwen 14B(教师)与1.5B(学生),调优超参数如k=10、β=0.9。评估指标为准确率,特别关注复杂推理任务的提升。通过消融实验验证偏好排序和置信度引导的贡献,分析不同参数设置对性能的影响。实验证明,TSD-KD在多项任务中实现了最高性能,显著优于对比方法。

结果分析

在10个推理任务中,TSD-KD最高提升54.4%,在MATH任务中超越第二名40.3%,甚至在部分任务中超越教师模型20.3%。在复杂推理任务如GSM8K、MMLU-Pro-Math中表现尤为突出,验证了其推理能力的增强。消融分析显示偏好排序引导的早期路径优化和置信度差异的关键Token选择性蒸馏是性能提升的关键因素。这些结果表明,TSD-KD不仅提升了模型的准确性,也增强了其自主推理和泛化能力。

应用场景

该方法适用于需要高推理能力的AI系统,如教育辅导、科研辅助、智能问答等。通过模型压缩,能在边缘设备上实现高效推理,降低部署成本。未来,结合多模态信息和强化学习,有望推动AI在更复杂场景中的自主推理和决策能力。

局限与展望

方法在极端复杂或噪声较多的任务中表现有限,偏好排序和置信度估计依赖于模型的准确性。计算成本较高,尤其在长推理链中需优化效率。此外,参数调优对不同任务敏感,泛化到其他领域仍需验证。

通俗解读 非专业人士也能看懂

想象你在教一个学生解数学题。传统方法就像老师把所有解题步骤都告诉学生,学生照做就行了,但有时信息太多,反而让学生迷失方向。这个新方法像老师只告诉学生关键的转折点,让学生自己尝试推理,然后老师根据学生的表现给出微弱的建议,帮助他们在关键点做得更好。这样,学生既能自主思考,又能在重要节点得到指导,逐步变得更聪明。整个过程就像在训练一只聪明的宠物,让它学会自己思考问题,而不是死记硬背答案。

简单解释 像给14岁少年讲一样

想象你在参加一个智力比赛,题目很难,需要你自己想办法解答。以前的做法就像老师把所有答案都告诉你,你只要照着做就行了,但这样你学不到怎么自己思考。现在,有个新方法,老师只在你关键的思考点给一些提示,比如“这个步骤很重要”,让你自己试着推理,然后老师根据你做的内容,给你一些微弱的建议,帮助你在关键地方做得更好。这就像在训练你变得更聪明,不仅仅是记住答案,而是真正学会了怎么思考。这样,你不仅能解决这道题,还能学会自己面对类似的问题,变得更厉害!

术语表

Knowledge Distillation(知识蒸馏)

一种模型压缩技术,通过让小模型模仿大模型的输出,提高性能。技术上是让学生模型学习教师模型的概率分布。

论文中用来描述模型压缩和推理能力迁移的方法。

On-policy KD(在策略知识蒸馏)

训练过程中学生模型使用自己生成的输出作为训练数据,减少分布偏移。

论文中强调的蒸馏方式,区别于传统的离线蒸馏。

Preference Ranking(偏好排序)

教师对学生候选响应的偏好排序,用于微弱反馈指导学生优化推理路径。

核心机制之一,用于引导学生在推理早期做出更合理的决策。

Entropy Regularization(熵正则化)

通过最小化模型输出的熵值,增强模型在关键Token上的信心,减少不确定性。

用以提升模型推理的稳定性和准确性。

KL Divergence(KL散度)

衡量两个概率分布差异的指标,用于模型输出分布的匹配。

在直接蒸馏中用来对齐学生与教师的输出分布。

开放问题 这项研究留下的未解疑问

  • 1 如何在极端复杂或噪声较多的任务中保持偏好排序和置信度估计的准确性仍是挑战。
  • 2 在超大规模模型或超长推理链中,计算效率和资源消耗需要进一步优化。

原文摘要

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.

cs.CL cs.AI cs.LG

参考文献 (20)

How (not) to Train your Generative Model: Scheduled Sampling, Likelihood, Adversary?

Ferenc Huszár

2015 317 引用 ⭐ 高影响力 查看解读 →

Distilling the Knowledge in a Neural Network

Geoffrey E. Hinton, O. Vinyals, J. Dean

2015 26069 引用 ⭐ 高影响力 查看解读 →

Gemma 2: Improving Open Language Models at a Practical Size

Gemma Team Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa 等

2024 2418 引用 ⭐ 高影响力 查看解读 →

On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou 等

2023 761 引用 ⭐ 高影响力 查看解读 →

DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

DeepSeek-AI

2025 5354 引用 ⭐ 高影响力

The Unreasonable Effectiveness of Entropy Minimization in LLM Reasoning

Shivam Agarwal, Shivam Agarwal, Zimin Zhang 等

2025 174 引用 ⭐ 高影响力 查看解读 →

Curriculum learning

Yoshua Bengio, J. Louradour, R. Collobert 等

2009 7616 引用

Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

C. Snell, Jaehoon Lee, Kelvin Xu 等

2024 2149 引用 查看解读 →

Crowdsourcing Multiple Choice Science Questions

Johannes Welbl, Nelson F. Liu, Matt Gardner

2017 937 引用 查看解读 →

Paper

N. Cambridge

1977 5597 引用

Program Synthesis with Large Language Models

Jacob Austin, Augustus Odena, Maxwell Nye 等

2021 4458 引用 查看解读 →

Graph of Thoughts: Solving Elaborate Problems with Large Language Models

Maciej Besta, Nils Blach, Aleš Kubíček 等

2023 1568 引用 查看解读 →

Advancing LLM Reasoning Generalists with Preference Trees

Lifan Yuan, Ganqu Cui, Hanbin Wang 等

2024 212 引用 查看解读 →

GSM-Plus: A Comprehensive Benchmark for Evaluating the Robustness of LLMs as Mathematical Problem Solvers

Qintong Li, Leyang Cui, Xueliang Zhao 等

2024 172 引用 查看解读 →

Autoregressive Knowledge Distillation through Imitation Learning

Alexander Lin, Jeremy Wohlwend, Howard Chen 等

2020 76 引用 查看解读 →

Qwen3 Technical Report

An Yang, An-Feng Li, Baosong Yang 等

2025 7856 引用 查看解读 →

Training Verifiers to Solve Math Word Problems

K. Cobbe, V. Kosaraju, Mo Bavarian 等

2021 10532 引用 查看解读 →

Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning

Shenzhi Wang, Le Yu, Chang Gao 等

2025 551 引用 查看解读 →

Training language models to follow instructions with human feedback

Long Ouyang, Jeff Wu, Xu Jiang 等

2022 24019 引用 查看解读 →

Large Language Models are Zero-Shot Reasoners

Takeshi Kojima, S. Gu, Machel Reid 等

2022 7993 引用 查看解读 →