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.