Task-Specific Knowledge Distillation via Intermediate Probes
Not on the map. Why?
Abstract (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.
Key facts. Cited by 2 papers in the map (data through 7 Oct 2026). Builds on 29 of its 41 references that are in the map. First posted on arXiv 18 Feb 2026.
Builds on
29 of 41 references.
References · 29
- 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
- 2015Distilling the Knowledge in a Neural Networkcited by 10,076
- 2015Rethinking the Inception Architecture for Computer Visioncited by 5,451
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2019Benchmarking Neural Network Robustness to Common Corruptions and Perturbationscited by 2,857
- 2022Locating and Editing Factual Associations in GPTcited by 2,609
- 2019Language Models as Knowledge Bases?cited by 1,922
- 2014FitNets: Hints for Thin Deep Netscited by 1,619
- 2020MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformerscited by 1,414
- 2016Understanding intermediate layers using linear classifier probescited by 1,279
- 2023Inference-Time Intervention: Eliciting Truthful Answers from a Language Modelcited by 1,259
- 2021DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharingcited by 1,219
17 more
- 2016Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfercited by 1,059
- 2019BERT Rediscovers the Classical NLP Pipelinecited by 1,046
- 2019When Does Label Smoothing Help?cited by 985
- 2019Relational Knowledge Distillationcited by 873
- 2019Contrastive Representation Distillationcited by 734
- 2023Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizescited by 710
- 2017Program Induction by Rationale Generation : Learning to Solve and Explain Algebraic Word Problemscited by 642
- 2019Designing and Interpreting Probes with Control Taskscited by 608
- 2023Eliciting Latent Predictions from Transformers with the Tuned Lenscited by 579
- 2019Patient Knowledge Distillation for BERT Model Compressioncited by 502
- 2018Analysis Methods in Neural Language Processing: A Surveycited by 276
- 2023Implicit Chain of Thought Reasoning via Knowledge Distillationcited by 165
- 2022Distilling Reasoning Capabilities into Smaller Language Modelscited by 155
- 2021Rethinking Soft Labels for Knowledge Distillation: A Bias-Variance Tradeoff Perspectivecited by 116
- 2023A Mechanistic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysiscited by 88
- 2024A Theoretical Analysis of Soft-Label vs Hard-Label Training in Neural Networkscited by 5
- 2025InnerThoughts: Disentangling Representations and Predictions in Large Language Modelscited by 1
Built on by
Citing papers per year: 2026: 2.
2 papers.