論文要約：EXPLAIN IN YOUR OWN WORDS: IMPROVING REASONING VIA TOKEN-SELECTIVE DUAL KNOWLEDGE DISTILLATION (TSD-KD)

この論文は、大規模言語モデル（教師モデル）の推論能力を小規模モデル（学生モデル）に転移させるための新しい知識蒸留（Knowledge Distillation: KD）手法、TSD-KD (Token-Selective Dual Knowledge Distillation) を提案しています。特に、複雑な推論タスク（Chain-of-Thought など）において、学生モデルが教師モデルの出力分布を盲目的に模倣するのではなく、自らの推論プロセスを強化し、場合によっては教師モデルを上回る性能を発揮することを目的としています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細に解説します。

1. 背景と問題定義

従来の知識蒸留（KD）は、教師モデルの出力分布（logits や確率分布）を学生モデルに一致させることで行われます。しかし、複雑な推論タスクにおいてこのアプローチには以下の課題があります。

• 分布の不一致（Distribution Mismatch）: 学生モデルの容量が限られている場合、教師モデルの全トークンにわたる厳密な分布追従は学生を圧倒し、分布の不一致を引き起こします。
• 推論プロセスの非対称性: 複雑な問題解決において、教師と学生の推論プロセスは本質的に異なる可能性があります（例：大学生と中学生の思考プロセスの違い）。すべてのトークンで教師の分布を強制することは、学生独自の推論能力の育成を阻害します。
• オフポリシー KD の限界: 教師が生成した固定データで学習するオフポリシー KD は、推論時の分布シフト（Distribution Shift）を招き、汎化性能を低下させます。

これらの課題に対し、著者は**「学生中心（Student-Centric）」**のアプローチを提案し、重要なトークンに焦点を当て、学生が自らの言葉で推論を説明することを促す手法を開発しました。

2. 提案手法：TSD-KD

TSD-KD は、オンポリシー（学生が生成したデータで学習） の枠組みにおいて、以下の 3 つの要素を組み合わせた「トークン選択型」の二重蒸留フレームワークです。

2.1 間接蒸留：教師の選好に基づくランキング（Indirect Distillation）

従来の KD が教師の分布を直接模倣するのに対し、間接蒸留は**「選好（Preference）」**という弱いフィードバックを利用します。

• プロセス: 学生モデルが各ステップでトップ-k の候補トークンを生成し、教師モデルがそれらを再ランキングします。
• 損失関数: プラケット・ルー（Plackett-Luce）モデルを用いて、学生が生成した候補の順序が教師の選好と一致するように学習します。
• トークン選択性: 推論の初期段階（「Opener」と呼ばれる部分）に集中します。図 1 の分析から、推論の分岐点となる重要なトークンは初期に高エントロピー（不確実性）を示すことが判明しており、この部分のみを蒸留対象とします。これにより、学生は残りの推論プロセスを自由に展開できます。

2.2 直接蒸留：不確実性のギャップに基づく選択（Direct Distillation）

教師の分布を直接追従する KL 発散（または一般化 Jensen-Shannon 発散：JSD）を使用しますが、すべてのトークンではなく**「学生が不確実で、教師が確信しているトークン」**に限定して適用します。

• ゲート機構: 学生のエントロピー $H_t(p_S)$ と教師のエントロピー $H_t(p_T)$ の差に基づき、シグモイド関数で蒸留の強度を動的に制御します。
• 効果: 学生が難しいと感じる重要なステップでのみ教師の知識を取り込み、学生が自信を持っている部分では干渉を最小限に抑えます。これにより、分布のシフトを抑制しつつ、重要な推論ステップの精度を向上させます。

2.3 エントロピー正則化（Entropy Regularization）

推論の信頼性を高めるため、学生モデルのエントロピー最小化を適用します。

• トークン選択性: すべてのトークンではなく、学生が最も不確実（エントロピーが高い）なトップ 10% のトークンにのみ適用します。
• 目的: 重要な推論ステップにおける学生の不確実性を低減し、推論経路の信頼性を高めます。

最終的な損失関数:
$$\min_{\theta} \mathbb{E}_{x \sim X} [\alpha L_{\text{Indirect}} + L_{\text{Direct}} + L_{\text{EM}}]$$
ここで、$\alpha$ は間接蒸留の重みです。

3. 主要な貢献

1. 学生中心の蒸留フレームワークの提案: 教師の分布を強制するのではなく、学生が自らの推論プロセスを強化するための「間接的」かつ「選択的」なフィードバックを提供する TSD-KD を提案しました。
2. トークン選択メカニズム: 推論の初期段階（Opener）や、教師と学生の不確実性のギャップが大きいトークンなど、重要なトークンにのみ蒸留を集中させることで、学習効率と汎化性能を向上させました。
3. 教師を超える性能の実証: 複数のベンチマークにおいて、学生モデルが教師モデル（14B → 1.5B など）を凌駕する結果を達成しました。これは、単なる模倣ではなく、学生モデル自身の推論能力が向上したことを示唆しています。

4. 実験結果

著者は Qwen2.5 (14B→1.5B, 9B→2B) および Gemma2 (9B→2B) などのモデルを用いて、数学推論、科学推論、コード生成、指示追従など 10 種類のベンチマークで評価を行いました。

• 主要な結果:

  • TSD-KD は、既存のオンポリシー KD 手法（GKD, MiniLLM, DistiLLM など）およびオフポリシー手法をすべて上回りました。
  • MATH ベンチマーク: 学生モデルの精度が 26.1 となり、ベースライン学生（16.9）に対して**54.4%の改善、2 位手法に対して40.3%**の改善を達成。
  • 教師モデルの凌駕: 4 つのタスクにおいて、TSD-KD で学習した学生モデルが教師モデルを最大**20.3%**上回りました（例：MATH で 21.7 → 26.1）。
  • Gemma2 での検証: 9B→2B の設定でも同様の傾向が見られ、IFEval や MBPP で教師を上回る結果を得ました。

• アブレーション研究:

  • 間接蒸留（$L_{\text{Indirect}}$）の導入が最も大きな性能向上（最大 15.5% 改善）をもたらしました。
  • 直接蒸留のトークン選択（ゲート機構）とエントロピー正則化もそれぞれ重要な役割を果たし、これらを組み合わせた完全な TSD-KD が最高性能を示しました。
  • 「Opener」の長さパラメータ $c$ について、推論の初期 10% 程度（$c=10\%$）に限定することが最適であり、全トークン（$c=100\%$）を蒸留すると性能が低下することが確認されました。

5. 意義と結論

この論文は、大規模モデルの推論能力を小規模モデルに転移させる際、**「すべての出力を模倣する」のではなく、「重要な推論ステップに焦点を当て、学生が自らの論理を構築する余地を残す」**というパラダイムシフトを示しています。

• 計算コストの削減: 小規模モデルでも複雑な推論タスクを高精度に実行可能にするため、推論コストを大幅に削減できます。
• 一般化性能の向上: 分布シフトを抑制し、学生モデルが未知のタスクに対しても柔軟に推論できる能力を育成します。
• 実用性: 教師モデルを凌駕する性能を達成できる可能性は、モデル圧縮の限界を突破する新たな道を開くものであり、リソース制約のある環境での高度な AI 応用にとって極めて重要です。

要約すれば、TSD-KD は「教える（蒸留する）」ことと「学ぶ（推論する）」ことのバランスを最適化し、学生モデルが自らの言葉で論理的思考を深めることを可能にする画期的な手法です。