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.
Figures & tables
Figure 1: Probe-KD vs Typical Distillation. In standard logit distillation (left), soft labels come from the teacher’s output layer, which projects hidden states onto answer tokens via a fixed, task-agnostic readout. This bottleneck produces noisy supervision even when the correct answer is encoded internally. Probe-KD (right) bypasses this bottleneck by training a probe to decode hidden states directly, learning a task-aligned projection. On AQuA-RAT, the probe achieves 52% accuracy versus the teacher’s 45%, demonstrating that cleaner readouts yield cleaner labels which in turn yields more performant students.
Figure 2: Data efficiency comparison on AQuA-RAT. Test accuracy (%) as a function of training data percentage. We distill from Qwen2.5-7B-Instruct (teacher) to DeBERTa-v3-base (student, 86M parameters). Probe-KD variants consistently outperform standard distillation baselines across all data regimes, with gains most pronounced in low-data settings.
Domain adapted
Compute efficient
Uses LLM representation
LLM Distillation
No
Yes
Yes
Train on labels only
Yes
Yes
No
Fine-tune and distil
Yes
No
Yes
Probe-KD
Yes
Yes
Yes
Table 1: Comparison of knowledge distillation approaches for building smaller task-specific models from LLMs. Probe-KD is the only approach to be domain adaptable, compute efficient, and to make use of LLMs internal representations.
Dataset
Teacher
Logistic
MLP
Δ
AQuA-RAT
44.7
50.6
50.3
+5.6
ARC-Easy
96.6
96.6
97.2
+0.6
ARC-Challenge
89.7
90.3
91.2
+1.5
MMLU
71.4
70.1
73.5
+2.1
Table 2: Probe accuracy vs. teacher accuracy (%). MLP probes on hidden states outperform teacher 5-shot outputs, with the largest gap on AQuA-RAT.
Figure 3: Probe-KD Calibration analysis on AQuA-RAT. We plot test accuracy against mean prediction confidence, with the x-axis reversed so that lower confidence (better calibration) appears rightward. Arrows indicate the distillation path from source (Teacher or MLP Probe) to student.
Method
Teacher Acc.
Student Acc.
Base Teacher (Qwen2.5-7B, no fine-tuning)
Logit-KD
44.7
26.6
Probe-KD (MLP)
44.7
29.4
Probe-KD (CCS) †
44.7
28.5
Fine-tuned Teacher (+LoRA, 3 epochs)
Logit-KD
52.3
27.8
Table 4: Fine-tuning vs. probe training on AQuA-RAT. The teacher is fine-tuned with LoRA for 3 epochs. Probe-KD improves distillation from both base and fine-tuned teachers.
Training Samples
Method
n =333
n =667
n =1000
Baseline Distillation
Logit-KD
25.4
24.9
25.3
Feature-KD
24.7
24.9
24.2
Patient-KD
25.9
25.9
25.9
Ours
Table 5: Data scaling on AQuA-RAT. Probe-KD consistently outperforms baseline distillation methods across data regimes.
Distilling robust reasoning capabilities from large language models (LLMs) into smaller, computationally efficient student models remains an unresolved challenge. Despite recent advances, distilled models frequently suffer from superficial pattern memorization and subpar generalization. To overcome these limitations, we introduce a novel distillation framework that moves beyond simple mimicry to instill a deeper conceptual understanding. Our framework features two key innovations. First, to address pattern memorization, Explanatory Inversion (EI) generates targeted ``explanatory probes'' that compel the student to articulate the underlying logic behind an answer, rather than just memorizing it. Second, to improve generalization, Explanatory GRPO (EXGRPO) uses a reinforcement learning algorithm with a novel Dialogue Structure Utility Bonus, which explicitly rewards the student for maintaining a coherent reasoning process across these probes. Extensive evaluations on 12 datasets demonstrate significant improvements. Using Gemma-7b as the student model, our method yields an average 20.39% increase over zero-shot performance and a 6.02% improvement over the state-of-the-art distillation baselines. Moreover, models distilled with our method show remarkable training efficiency (e.g., surpassing vanilla fine-tuning with 10-25% training data) and strong generalization to out-of-distribution tasks. Implementation is released at https://github.com/Zhen-Tan-dmml/ExGRPO.git.
Zhen Tan, Chengshuai Zhao, Song Wang +3
Arizona State University · University of Virginia · University of North Carolina at Chapel Hill
Distilling reasoning traces from strong large language models into smaller ones is a promising route to improve intelligence in resource-constrained settings. Existing approaches face a fundamental trade-off: offline distillation from teacher-generated traces provides high-quality, sample-efficient supervision but suffers from distributional drift: during training, the student model conditions on teacher-generated prefixes, whereas during inference the student autoregresses on self-generated prefixes, leading to compounding errors over long reasoning trajectories. Meanwhile, on-policy or self-distillation methods better match the student's inference-time distribution, but require costly online sampling and often produce low-quality traces in early training. We propose a principled offline reasoning distillation framework that preserves the efficiency and supervision quality of offline teacher-generated data while correcting teacher-student distribution drift. It adaptively emphasizes teacher supervision that is better aligned with the student's on-policy distribution. Evaluations on mathematical reasoning benchmarks of GSM8K, MATH, MATH500, and harder held-out competition-style tasks, including AMC, AIME, and OlympiadBench, show that our method improves reasoning accuracy over prior offline distillation algorithms and yields more stable reasoning traces while preserving instruction-following capabilities. Our work shows that lightweight, distribution-correction-aware training can substantially strengthen offline reasoning distillation without online rollouts.
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.