知識蒸留は中間訓練中に事実想起よりも推論を優先する

要旨

ロジットベースの知識蒸留（KD）は、より強力な教師モデルからの教師信号を用いて小規模な言語モデル（LM）を学習させるために用いられるが、その利点が学習段階を通じて一貫しているかは依然として不明である。我々は制御実験を通じて、標準的なKD定式化であるフォワードKullback-Leibler（KL）蒸留が、事後学習済みの教師モデルを用いた場合、中間学習（すなわち厳選コーパス上での自己教師あり学習からなる中間段階）中には根本的に異なる挙動を示すことを見いだした。驚くべきことに、フォワードKDは、事前学習中には標準的な次トークン予測（NTP）と比較して推論と事実想起性能の両方を同時に改善する一方、中間学習中には、推論性能の向上が続いているにもかかわらず、事実想起能力の獲得を遅らせる。我々は、この段階依存性が、データ領域間における教師モデルの確信度の非対称性と、生徒モデルの変化する知識状態に起因することを明らかにする。すなわち、教師モデルは知識集約型データよりも手続き型データに対して高い確信度を示す一方、生徒モデルは学習の早い段階で低エントロピーな事実知識を獲得する。この不均衡を緩和するため、我々はSwitch Distillationを提案する。これは、教師モデルの予測エントロピーを軽量なルーティング信号として用いて、教師モデルが高い確信度を持つトークンに対して蒸留を行い、それ以外の場合には交差エントロピーにフォールバックする、シンプルな中間学習用の目的関数である。Switch Distillationは、教師モデルの規模にかかわらず、既存の蒸留目的関数を一貫して上回る。標準的なNTPと比較すると、推論性能では1.61〜1.71倍、知識・常識性能では1.13〜1.19倍を達成し、事実想起性能の96.7〜96.8%を維持する。重要なことに、これらの利点は事後学習後も持続する。Switch Distillationは、事実想起性能のギャップを解消しつつ、推論では1.25〜1.32倍、知識・常識性能では1.13〜1.20倍の向上をそれぞれ維持する。

English

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.