研究提出 Switch Distillation:中训练阶段知识蒸馏更利于推理而非事实记忆

HuggingFace Daily Papers(社区热门论文)·2026-09-01 08:00·2天前
AI 导读

研究发现前向 KL 蒸馏在 pre-training 阶段同时提升推理与事实记忆,但在 mid-training 阶段会拖慢事实记忆习得,同时推理仍持续提升。作者提出 Switch Distillation,以教师预测熵作为路由信号,仅在教师置信的 token 上蒸馏,其余回退到交叉熵。

HuggingFace Daily Papers(社区热门论文)
43AI 编辑部评分,满分 100

研究提出 Switch Distillation:中训练阶段知识蒸馏更利于推理而非事实记忆

2026-09-01 08:00· 2天前
AI 导读

研究发现前向 KL 蒸馏在 pre-training 阶段同时提升推理与事实记忆,但在 mid-training 阶段会拖慢事实记忆习得,同时推理仍持续提升。作者提出 Switch Distillation,以教师预测熵作为路由信号,仅在教师置信的 token 上蒸馏,其余回退到交叉熵。

Logit-based knowledge distillation (KD) is used to train smaller language models (LMs) via supervision from stronger teachers, but whether its benefits are consistent across training stages remains unclear. Through controlled experiments, we find that forward Kullback-Leibler (KL) distillation--the standard KD formulation--with post-trained teachers behaves fundamentally differently during mid-training, an intermediate phase of self-supervised learning on curated corpora. Surprisingly, while forward KD simultaneously improves reasoning and factual recall during pre-training relative to standard next-token prediction (NTP), it instead slows factual recall acquisition during mid-training despite continued reasoning gains. We trace this stage dependence to an asymmetry in teacher confidence across data domains and the student's evolving knowledge state: teachers are more confident on procedural than knowledge-intensive data, while students acquire low-entropy factual knowledge earlier in training. To mitigate this imbalance, we propose Switch Distillation, a simple mid-training objective that distills on tokens where the teacher is confident, using teacher predictive entropy as a lightweight routing signal, and otherwise falls back to cross-entropy. Switch Distillation consistently outperforms existing distillation objectives across teacher sizes. Relative to standard NTP, it achieves 1.61-1.71x the reasoning performance and 1.13-1.19x the knowledge and commonsense performance while preserving 96.7-96.8% of factual recall. Crucially, these benefits persist after post-training: Switch Distillation closes the factual recall gap while maintaining 1.25-1.32x and 1.13-1.20x gains in reasoning and knowledge and commonsense, respectively.

来源:HuggingFace Daily Papers(社区热门论文)· arxiv.org