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

1. Problem Statement

Knowledge Distillation (KD) aims to transfer reasoning capabilities from large "teacher" models to smaller "student" models to reduce inference costs, particularly for Chain-of-Thought (CoT) tasks. However, existing methods face significant limitations:

• Distribution Mismatch: Traditional off-policy KD trains students on static teacher data, causing a mismatch between the training distribution and the student's own generated distribution during inference.
• Overwhelming Supervision: On-policy KD methods (where the student trains on its own outputs) often force the student to mimic the teacher's distribution at every token. For complex reasoning tasks, a student with limited capacity may be overwhelmed by this extensive supervision, leading to a "distribution mismatch" where the student fails to generalize because it cannot replicate the teacher's exact reasoning path.
• Lack of Student Agency: Current methods often treat the student as a passive imitator rather than an active reasoner, preventing the student from developing its own reasoning strategies ("explaining in its own words").

2. Methodology: Token-Selective Dual Knowledge Distillation (TSD-KD)

The authors propose TSD-KD, a student-centric, on-policy framework that combines indirect and direct distillation, both applied selectively to specific tokens rather than the entire sequence.

A. Core Philosophy

The method is built on the observation that reasoning difficulty is not uniform. High-entropy tokens (where the model is uncertain) often appear early in reasoning traces and represent critical "branching points." TSD-KD focuses supervision on these critical moments while allowing the student freedom elsewhere.

B. Indirect Distillation (Teacher-Guided Preference Ranking)

Instead of forcing the student to match the teacher's full probability distribution, TSD-KD uses a preference-based approach:

1. Student Generation: The student generates its own response and proposes top-$k$ candidate tokens for the next step.
2. Teacher Re-ranking: The teacher acts as an oracle, re-ranking these student-proposed candidates based on its own logits.
3. Loss Function: The student is trained to align its ranking of candidates with the teacher's preference using a Plackett-Luce (PL) model.
   • Key Innovation: The student actively proposes reasoning paths; the teacher only guides the selection. This reduces distribution shift and encourages the student to explore its own reasoning space.
4. Token Selection (Openers): This indirect distillation is applied only to the "opener"—the initial segment of the response where cumulative entropy reaches a threshold (e.g., the first 10% of cumulative entropy). This targets the critical initial direction of reasoning.

C. Direct Distillation (Uncertainty Gap Gating)

For tokens where the student is highly uncertain but the teacher is confident, TSD-KD applies direct distribution matching:

1. Uncertainty Gap: It calculates the difference in entropy between the student ($H_t(p_S)$) and the teacher ($H_t(p_T)$).
2. Soft Gating: A sigmoid gating function $\sigma(H_t(p_S) - H_t(p_T))$ modulates the distillation loss.
   • If the student is uncertain and the teacher is confident (large gap), the loss is reinforced.
   • If the student is already confident, the loss is suppressed, allowing the student to generate freely.
3. Objective: Uses Generalized Jensen-Shannon Divergence (JSD) to balance mode-seeking and mode-covering behaviors.

D. Entropy Regularization

To boost confidence in critical tokens:

• The method identifies the top 10% of tokens with the highest entropy in the student's output.
• It applies an entropy minimization loss specifically to these tokens.
• This forces the student to become more deterministic (confident) on the most difficult reasoning steps without stifling exploration on easier steps.

E. Final Objective

The total loss is a weighted sum:
$$\mathcal{L}_{total} = \alpha \mathcal{L}_{Indirect} + \mathcal{L}_{Direct} + \mathcal{L}_{EM}$$
Where $\mathcal{L}_{Indirect}$ is the preference ranking loss, $\mathcal{L}_{Direct}$ is the gated JSD loss, and $\mathcal{L}_{EM}$ is the selective entropy minimization.

3. Key Contributions

1. Student-Centric Framework: Shifts from "mimicry" to "guided self-improvement," allowing students to reason in their own words while receiving targeted feedback.
2. Dual Distillation Mechanism: Combines indirect preference ranking (for global reasoning structure) and direct uncertainty-gap distillation (for local token accuracy).
3. Token-Selectivity: Introduces the concept of "Openers" and uncertainty gaps to apply supervision only where it is most needed, preventing the student from being overwhelmed.
4. Empirical Superiority: Demonstrates that a small student model can outperform its much larger teacher model in specific reasoning tasks.

4. Experimental Results

The authors evaluated TSD-KD on 10 challenging reasoning benchmarks (Math, Science, Code, Logic) using Qwen2.5 (14B Teacher $\to$ 1.5B Student) and Gemma2 (9B Teacher $\to$ 2B Student).

• State-of-the-Art Performance: TSD-KD outperformed all baselines (including MiniLLM, GKD, DistilLLM) across almost all tasks.
  • Improvement over Baseline: Up to 54.4% accuracy improvement over the baseline student.
  • Improvement over Runner-up: Up to 40.3% improvement over the second-best method.
• Student Surpassing Teacher: In several cases, the distilled student outperformed the teacher model itself:
  • MATH Benchmark: Student (26.1%) > Teacher (21.7%) by 20.3%.
  • SciQ Benchmark: Student (93.0%) > Teacher (86.4%) by 7.6%.
  • IFEval & MBPP (Gemma2): Student surpassed the 9B teacher.
• Ablation Studies:
  • Removing the token-selective gating or indirect distillation significantly degraded performance.
  • The "Opener" selection (focusing on the first 10% of entropy) was critical; applying distillation to the whole sequence (100%) performed worse than selective distillation.

5. Significance

• Efficiency: TSD-KD proves that high-level reasoning can be compressed into small models without requiring the student to blindly copy the teacher's entire thought process.
• Generalization: By encouraging the student to "explain in its own words" via indirect feedback, the method improves generalization to unseen tasks compared to standard imitation learning.
• Paradigm Shift: The results challenge the assumption that a student must always be weaker than its teacher. With the right distillation strategy (targeted, student-centric), smaller models can learn to reason more effectively than their larger counterparts on specific complex tasks.
• Practicality: The method is compatible with Parameter-Efficient Fine-Tuning (PEFT/LoRA), making it viable for resource-constrained environments.

In summary, TSD-KD represents a significant advancement in model compression for reasoning, moving away from rigid distribution matching toward a flexible, targeted, and student-driven learning paradigm.