arXiv:2603.12270cs.CLcs.AI2026-03被引 1

用中间层探针提升大模型推理知识蒸馏效果

Task-Specific Knowledge Distillation via Intermediate Probes

  • 在教师模型中间层训练轻量探针,用探针输出代替原始输出做监督
  • 在四个推理基准上均获提升,小数据下增益更明显
  • 无需修改模型结构,计算开销极低,适合资源受限场景

大型语言模型的知识蒸馏通常假设教师输出分布是高质量的训练信号。但在推理任务中,这一假设常不成立:模型的中间表示可能已包含正确答案,但经过词汇投影后,提示格式和答案词选择会导致输出脆弱且噪声大。本文提出一种新蒸馏框架 method{},通过在冻结的教师隐藏状态上训练轻量探针,并用探针预测结果而非输出逻辑作为学生模型的监督信号,绕过该瓶颈。该方法在四个推理基准(AQuA-RAT、ARC Easy/Challenge、MMLU)上实现一致提升,尤其在数据有限时效果更显著。中间表示训练的探针能提供比教师自身输出更清晰的标签,有效降噪蒸馏信号。 method{} 不需修改师生模型架构,具备架构无关性,探针训练成本低,教师表示可缓存复用。通过利用内部表示,该方法使从业者无需额外数据或复杂结构即可从大模型中提取更高价值。

原文摘要 · Abstract (English)

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.

知识蒸馏推理增强中间表示轻量探针

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。