Source-linked AI summary
Explain in Your Own Words: Improving Reasoning via Token-Selective Dual Knowledge Distillation
Minsang Kim, Seung Jun Baek
TL;DR
Full-distribution distillation can mismatch the reasoning of limited-capacity students on complex tasks. TSD-KD instead combines token-selective direct and preference-based indirect feedback with entropy regularization, achieving leading benchmark results and surpassing teachers in some cases.
Problem
Full-distribution supervision can overwhelm limited-capacity students and create distribution mismatch when teacher and student reasoning processes differ.
Method
TSD-KD combines student-proposed candidate re-ranking, confidence-based selective direct distillation, and entropy regularization over important tokens.
Results
TSD-KD achieves the best performance across six tasks, improves over the baseline by up to 54.4%, and surpasses its teacher on MATH by 20.3%.
Takeaways & Limitations
Student-centered, token-selective distillation can improve reasoning performance while allowing students to explain reasoning in their own words.
Takeaways & Limitations
The entropy-regularization design fixes the selected-token ratio at the top 10%, with alternative ratios studied only in the appendix.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper argues that full-distribution distillation can overwhelm limited-capacity students on complex reasoning tasks, motivating targeted and indirect supervision. TSD-KD selects important tokens, uses student-generated reasoning, and combines preference feedback with selective direct distillation.
- Knowledge distillation compresses large reasoning models into smaller students, reducing inference costs for generating long Chains-of-Thoughts.
- Off-policy KD can harm generalization because teacher-generated training data differs from the student’s own output distribution.
- On-policy methods reduce distribution shift but may overwhelm limited-capacity students by forcing them to match the teacher at every generated token.
- The paper motivates targeted and indirect distillation to match the student’s reasoning level while allowing room for self-improvement.
- High-entropy tokens mark uncertain outputs, often occur early in reasoning traces, and motivate focusing distillation on important branching points.
- TSD-KD combines student-proposed candidate feedback, selective direct distillation, and entropy regularization while allowing later reasoning to remain student-generated.
2 PRELIMINARIES
The preliminaries distinguish on-policy from off-policy distillation and describe distribution-matching objectives such as KL and generalized Jensen-Shannon divergence. The section also introduces TSD-KD’s preference-guided, student-generated reasoning setup.
- On-policy KD trains on student-generated outputs, whereas off-policy KD relies on fixed teacher-generated data.
- On-policy training can better fit the student’s distribution and mitigate the distribution shift associated with off-policy distillation.
- Autoregressive language models define token probabilities with a softmax over logits conditioned on the input and preceding tokens.
- KL divergence measures mismatch between teacher and student distributions, with forward and reverse forms favoring mode-covering and mode-seeking behavior, respectively.
- Generalized Jensen-Shannon divergence interpolates between forward- and reverse-KL behavior through a weighting parameter.
- TSD-KD’s indirect distillation has the student propose top-k candidates while the teacher only re-ranks those proposals.
3 PROPOSED METHOD
TSD-KD combines student-generated preference feedback with selective distribution matching, targeting difficult reasoning tokens while leaving other tokens freer for self-expression.
- TSD-KD combines indirect preference distillation, direct uncertainty-gap distillation, and token-selective entropy regularization.The framework is designed to provide targeted supervision rather than feedback on every generated token.
- In indirect distillation, the student proposes top-k candidate tokens and the teacher re-ranks only those candidates using its own logits.The student supplies the candidates, while the teacher provides preference feedback instead of enforcing its full output distribution.
- The preference-ranking objective aligns the student’s candidate ordering with the teacher’s ordering, using the Plackett-Luce model.The Bradley-Terry model is the k = 2 special case, and the loss is optimized over token positions.
- Direct distillation selects tokens where the student has high entropy and the teacher has low entropy, focusing supervision on large uncertainty gaps.This selective mechanism is intended to address cases where the student’s candidate reasoning path may be substantially wrong.
- Indirect distillation provides sequence-level guidance for early reasoning, whereas direct distillation targets important tokens with token-level guidance.Their combination is presented as complementary: indirect distillation shapes the early reasoning path, while direct distillation refines selected tokens.
- Entropy regularization minimizes entropy for the student’s most uncertain tokens while withholding feedback from remaining tokens.The method targets the top-10% highest-entropy tokens to boost confidence on critical reasoning tokens.
4 EXPERIMENT
Across Qwen2.5, Gemma2, and Qwen3 experiments, TSD-KD generally achieves the strongest reasoning performance while using token-selective distillation. Results also show that concentrating feedback early and adding the proposed components improves outcomes.
- Main Results: TSD-KD achieves the best performance across six Qwen2.5 tasks, improving accuracy by up to 54.4% over the baseline student.
- Main Results: On four more challenging Qwen2.5 tasks, TSD-KD leads the runner-up by 9.8% and surpasses its teacher by 7.6% on SciQ.
- Ablation Study: The ablation study shows that indirect distillation improves performance by up to 15.5%, while selective entropy minimization raises MATH from 18.1 to 22.3.
- Ablation Study: The fully integrated TSD-KD model improves performance by up to 44.2% over the baseline, indicating gains from combining its components.
- Generalization to Gemma2 Models: On Gemma2, TSD-KD achieves the highest performance across five tasks and surpasses the teacher on IFEval and MBPP.
- Generalization to Qwen3 Models: On Qwen3, TSD-KD achieves state-of-the-art performance on most benchmarks, with the student exceeding the teacher on MATH by 25.0%.
- Analysis of Openers in Indirect Distillation: Peak average performance occurs at c = 10%, whereas applying distillation to all tokens at c = 100% worsens performance.
5 RELATED WORK
Prior work developed knowledge distillation and reasoning strategies to make complex reasoning more efficient, but both distribution shift and computational cost remain challenges.
- Sequence-level and distribution-matching distillation methods train students from teacher-generated outputs but can suffer from distribution shift.
- Chain-of-Thought and Tree-of-Thoughts improve complex reasoning through intermediate or multiple reasoning paths, but generating these processes is computationally expensive.
- Knowledge distillation transfers large teacher models into smaller student models to reduce the computational cost of generating reasoning processes.
6 CONCLUSION
The paper concludes that TSD-KD improves student reasoning by selectively distilling important tokens through complementary indirect and direct supervision.
- TSD-KD focuses distillation on important tokens to enhance student-model reasoning performance.
- The framework combines teacher preference rankings, selective distribution matching, and entropy regularization to support student-centered distillation.
- Experiments show substantial gains over existing on-policy methods, with students surpassing their teachers in several cases.
8 REPRODUCIBILITY STATEMENT
The study reports its implementation resources, evaluation sources, and public code to support reproducibility.
- Training uses a TRL-based framework with newly developed loss functions.
- Inference datasets are generated with vLLM.
- Math and code tasks are evaluated using open-access repositories in lm-evaluation-harness.
- The source code is available in a public repository.
- Experiments use eight A100 GPUs with 80GB of VRAM each.
A.1 PROOF OF PROPOSITION 1
The proof relates teacher and student logits to preference distributions, establishing that the indirect-distillation objective is a preference-model likelihood.
- The preference distribution is expressed using a Plackett–Luce model over sub-responses.
- Teacher preference between sub-responses is determined by comparing the teacher model’s logits.
- The proof concludes that indirect distillation is equivalent to the negative log-likelihood of the preference distribution summed across reasoning positions.
A.2 ANALYSIS OF SOFT TOKEN SELECTION VIA στ IN DIRECT DISTILLATION.
Soft token selection rescales direct-KD gradients according to the relative uncertainty of teacher and student tokens. It strengthens supervision for less-confident students early in training and weakens it later to support reasoning diversification.
- Gradient rescaling: When the teacher is more confident than the student on a token, KD provides positive feedback that raises the token’s logit.The analysis treats the high-probability answer token as the ground-truth token for reasoning supervision.
- Gradient rescaling: A weaker student receives an additional multiplicative gradient gain through the token-selection rescaling factor.The factor exceeds 1 when the weaker student is less confident than a later, stronger version of itself.
- Soft token selection: Token-selective KD uses entropy-based sigmoid gating to softly select tokens for direct distillation.The gating compares teacher and student token entropy while using stop-gradient for the selection weight.
- Training-stage effects: When student confidence is low, the rescaling factor exceeds 1; when confidence is higher, it falls below 1.This makes the selection effect stronger early in training and weaker at later stages.
- Training-stage effects: Early training receives reinforced gradients toward the best token, whereas later training receives weaker gradients that diversify reasoning similarly to label smoothing.The paper links these two regimes to establishing logical abilities first and expanding the reasoning process afterward.
A.3 ANALYSIS OF TOP-k IN INDIRECT DISTILLATION.
Indirect distillation performs best with a focused candidate set: too few candidates limit learning, while broad selection introduces noise into the supervision signal.
- Top-k selection: Performance peaks at k = 10 in top-k sub-response selection for indirect distillation.Using k = 5 is insufficient, while increasing k to 20 reduces the score to 44.5.
- Top-k selection: The authors select k = 10 because forcing the student to match a wide token distribution appears to introduce noise.The result supports focused rather than broad distillation.
A.4 ANALYSIS OF THE TOKEN SELECTION RATIO FOR ENTROPY REGULARIZATION
The analyses favor selective supervision: entropy regularization peaks at 10% selection, JSD peaks near β = 0.9, and on-policy training consistently improves over off-policy training. TSD-KD also retains strong performance under parameter-efficient fine-tuning.
- Entropy-regularization ratio: 10% entropy-regularization token selection raises performance from 42.9 without regularization to a peak of 45.7.Performance declines above 20% and reaches 41.8 when all tokens are selected.
- Entropy-regularization ratio: The authors choose s = 10% because full-token entropy minimization can cause overconfidence and the useful-token estimate is 10–20%.The selected ratio is also consistent with the cited estimate for crucial reasoning tokens in reinforcement learning.
- JSD coefficient: JSD performance peaks at β = 0.9 with a score of 45.7, compared with 43.4 for β = 0.Increasing β to 1.0 slightly lowers performance to 45.4, motivating the setting β = 0.9.
- On-policy versus off-policy: On-policy TSD-KD averages 45.7 versus 43.4 for off-policy training across six benchmarks, with gains of +4.1 on GSM8K and +3.4 on GSM-Plus.The comparison supports guiding students along their own generated reasoning paths.
- Adaptive selection: Adaptive-c performs similarly to fixed-c overall and performs better on MATH and SciQ.The authors interpret this as evidence that adaptive-c can respond to distillation-sample difficulty.
- Parameter-efficient fine-tuning: Under LoRA-based PEFT, TSD-KD achieves the highest average score of 43.5 and ranks first on 5 of 6 tasks.Its top-rank results include the challenging GSM8K and MMLU-STEM benchmarks.