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

Minsang Kim; Seung Jun Baek

cs.CLFeb 25, 2026

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

Authors: Minsang Kim, Seung Jun Baek

Organizations: Korea University · SK Telecom

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.

Figures & tables

Figure 1: Average per token entropy of reasoning traces for ( Yuan et al., 2025 ) .

Figure 2: Overview of TSD-KD. 1) Indirect distillation. Unlike traditional KD, the student actively suggests reasoning candidates, and the teacher chooses from them. 1.1) The student generates a reasoning response. 1.2) Early part of the response containing important branching of reasoning (called opener ) is selected. 1.3) The student sequentially generates candidates of partial responses from the opener. Top candidates are proposed to the teacher; for example, \raisebox{-0.9pt}{1}⃝ “ To solve x+3=5, we subtract 5 ” and \raisebox{-0.9pt}{2}⃝“ To solve x+3=5, we subtract 3 ”. The teacher provides preference ranking (prefers \raisebox{-0.9pt}{2}⃝) for better reasoning. The preference signal is used as an indirect form of distillation. 2) Direct distillation performs selective distillation of critical tokens about which the student is uncertain but the teacher is confident. 3) Entropy regularization minimizes the entropy of critical tokens, reducing uncertainty and maintaining the student’s confidence during distillation.

Figure 3: Performance with varying $c$ .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Performances using different top- $k$ of $\mathcal{L}_{\text{Indirect}}$ .

Figure 5: Performances with token selection ratio for entropy regularization.

Figure 7: Comparison of mode-covering (forward KL, $\mathcal{D}(P\|Q)$ , red ) and mode-seeking (reverse KL, $\mathcal{D}(Q\|P)$ , blue ) behavior. The student ( $Q$ ) is trained to match the bimodal teacher ( $P$ ). The mode-seeking distribution fits one mode precisely, while the mode-covering distribution spreads its mass to cover both. Adapted from ( Le, 2017 ) .

Q: Write a python function to convert the given string to lower case.

Qwen2.5-1.5B: import re def is_lower(s): # Check if all characters in the string are lowercase return s.islower()
GKD: import re def is_lower(s): # Check if all characters in the string are lowercase return s.islower()
Ours:
import re def is_lower(s): s = str.lower(s) return s

Appendix

Figure 8: Qualitative Analysis of Code generation (MBPP).

Q: Trey receives a $5\%$ commission on every sale he makes. On the sale of a $\$ 60 $coat (before any discounts), how many more cents will he receive if his commission is based on the original price of the coat rather than the price of the coat after a$ 20%$ discount?
GT: 60
Qwen2.5-1.5B: To determine how much more Trey will receive as a commission if the sale is based on the original price of the coat rather than the discounted price, we need to calculate the commissions for both scenarios and then find the difference. First, let’s calculate the commission based on the original price of the coat: The original price of the coat is $60. With a 5% commission rate, the commission Trey would receive is:$ 60\times 0.05=3\text{ dollars} $Next, let’s calculate the commission based on the discounted price of the coat: The coat is sold at a 20% discount, so the selling price before tax is:$ 60\times(1-0.20)=60\times 0.80=48\text{ dollars} $With a 5% commission rate, the commission Trey would receive is:$ 48\times 0.05=2.4\text{ dollars} $Now, we find the difference between the two commissions:$ 3-2.4=0.6\text{ dollars}$ Since the question asks for the difference
GKD: To determine how much more Trey will receive as a commission when the sale is based on the original price of the coat rather than the discounted price, let’s follow these steps: 1. Calculate the original price of the coat before any discount: The original price of the coat is $60. 2. Determine the commission based on the original price: Trey receives a 5% commission on the original price of the coat. So, the commission from the original price is calculated as follows: Commission from original price = 0.05 x 60 = 3 dollars 3. Calculate the price of the coat after a 20% discount: A 20% discount means the customer pays 80% of the original price. Thus, the discounted price is: Discounted price = 60 x (1 - 0.20) = 60 x 0.80 = 48 dollars 4. Determine the commission based on the discounted price: With the discounted price, the commission is Not complete
Ours:
The original price of the coat is $60$ dollars. Trey gets a $5\%$ commission on sales. So he gets $60\times.05=3$ dollars from the sale. Now if the price was discounted by $20\%$ , the new price would be $60(1-.2)=48$ dollars. He would get $48\times.05=2.4$ dollars from the sale. So he gets $3-2.4=0.6$ dollars more when the price is not discounted. In cents, this is $0.6\times 100=60$ cents. Therefore, he will receive $60$ more cents. Final Answer: The final answer is 60 cents.

Appendix

Figure 9: Qualitative Analysis of Math solution (MATH).

Explore similar work

Task-Specific Knowledge Distillation via Intermediate Probes
Feb 18, 2026Ryan Brown, Chris RussellLanguage Model Distillation Knowledge Distillation
Oxford Internet Institute, University of Oxford, Oxford, United Kingdom.
Finding the Evidence: Discovering Decision-Supporting Tokens for On-Policy Reasoning Distillation
Jun 22, 2026Jinwei Xiao, Zhuowen Han, Yueqing Sun +6CoT Distillation On-Policy Distillation
Meituan Longcat Team · TJUNLP Lab, College of Intelligence and Computing, Tianjin University · Nanjing University
TeacherGRPO: Closing the Capacity Gap in Reasoning Distillation via Teacher Alignment
Sep 27, 2026Zhenyu Lei, Zihan Chen, Yaochen Zhu +5Reasoning Distillation RL for Language Model Reasoning
University of Virginia · Netflix · University of Washington +2

Feb 18, 2026J/K move · Enter open · S save

Task-Specific Knowledge Distillation via Intermediate Probes

Ryan Brown, Chris Russell

Oxford Internet Institute, University of Oxford, Oxford, United Kingdom.

Language Model Distillation Knowledge Distillation

Open paper AskPDF arXiv

Knowledge distillation from large language models (LLMs) assumes that the teacher's output distribution is a high-quality training signal. On reasoning tasks, this assumption is frequently violated. A model's intermediate representations may encode the correct answer, yet this information is lost or distorted through the vocabulary projection, where prompt formatting and answer-token choices creates brittle, noisy outputs. We introduce \method{}, a distillation framework that bypasses this bottleneck by training lightweight probes on frozen teacher hidden states and using the probe's predictions, rather than output logits, as supervision for student training. This simple change yields consistent improvements across four reasoning benchmarks (AQuA-RAT, ARC Easy/Challenge, and MMLU), with gains most pronounced under limited data. Probes trained on intermediate representations provide cleaner labels than the teacher's own outputs, effectively denoising the distillation signal. \method{} requires no architectural changes to student or teacher, is architecture-agnostic, and adds minimal compute since probe training is cheap and teacher representations can be cached. By exploiting internal representations, \method{} enables practitioners to extract more value from large teacher models without additional training data or architectural complexity.

Abstract

Explore similar work

Task-Specific Knowledge Distillation via Intermediate Probes

Finding the Evidence: Discovering Decision-Supporting Tokens for On-Policy Reasoning Distillation

TeacherGRPO: Closing the Capacity Gap in Reasoning Distillation via Teacher Alignment