技術的サマリー：LLM 蒸留における KL と軌道の分離

問題定義

知識蒸留は、特に推論タスクにおいて、強化学習（RL）と併せて運用されることが多く、大規模言語モデル（LLM）のポストトレーニングの基盤となっています。しかし、蒸留の設計空間は未だ十分に理解されていません。現在のパイプラインは、通常、2 つのパラダイムのいずれかに従います：オフポリシー蒸留（教師生成のトレースを用いた学習、通常は教師付き微調整（SFT）を介して）またはオンポリシー蒸留（OPD）（学生生成のロールアウトを教師フィードバックとともに用いた学習）。

本論文は、これらの支配的なパラダイムが、2 つの直交する設計選択を暗黙的に結合していると主張します。

1. プレフィックスソース: 監督が、教師によって生成されたプレフィックス（オフポリシー）に適用されるか、学生によって生成されたプレフィックス（オンポリシー）に適用されるか。
2. トークンレベルの KL 方向: 目的関数が前方 KL（$KL(p_{teacher} \parallel q_{student})$）を最小化するのか、逆 KL（$KL(q_{student} \parallel p_{teacher})$）を最小化するのか。

従来の実践では、教師プレフィックスと前方 KL、学生プレフィックスと逆 KL が結合されています。著者らは、この結合がトークンレベルの本質的なものではなく、これらの軸を分離することで、それぞれが異なる学習レジームに対応する 4 つの有効な目的関数を持つ、より広範な設計空間が明らかになると主張しています。

手法

理論的枠組み：軸の分離

著者らは、自己回帰的な連鎖律を用いてシーケンスレベルの KL 発散を分解します。彼らは以下のことを示します。

• 前方 KL（$KL(p_T \parallel q_S)$）は、期待値が教師の軌道分布に対して取られるため、教師プレフィックス（オフポリシー）と自然にペアになります。
• 逆 KL（$KL(q_S \parallel p_T)$）は、期待値が学生の軌道分布に対して取られるため、学生プレフィックス（オンポリシー）と自然にペアになります。

これらを独立した軸として扱うことで、著者らは 4 つの目的関数を定義します。

1. 教師プレフィックス + 前方 KL: オフポリシー SFT（標準的な蒸留）に対応します。
2. 学生プレフィックス + 前方 KL: DAgger スタイルのオンポリシー SFT に対応します。
3. 教師プレフィックス + 逆 KL: オフライン RL スタイルの蒸留目的関数に対応します。
4. 学生プレフィックス + 逆 KL: 密報酬オンポリシー RL として捉えられるオンポリシー蒸留（OPD）に対応します。

本論文は、以下の勾配レベルの恒等式を確立します。

• 前方 KLは、教師のソフトターゲットを用いた SFT スタイルのクロスエントロピー勾配を生成します。
• 逆 KLは、教師 - 学生対数比（$r = \log p_T - \log q_S$）で定義された密報酬を用いた RL スタイルの方策勾配を生成します。

実験設定

著者らは、Qwen3 モデル（0.6B の学生、4B/8B の教師）を用いて数学的推論に関する制御された研究を実施しました。彼らは、2 つのレジームにおいて 4 つの目的関数を評価しました。

1. 単独蒸留: 各目的関数のみで学習されたモデルの最終性能を比較します。
2. 蒸留後 RL: 蒸留済みチェックポイントを Group Relative Policy Optimization（GRPO）の初期化として使用します。

彼らは、さまざまなシーケンス長さ（128 トークン対 4096 トークン）およびベンチマーク（MATH500、GSM8K、AIME24、AMC23）にわたって、精度、予測エントロピー、応答長を追跡しました。

主要な貢献

1. 3 つのトレードオフの特定

実証的分析により、蒸留設計を支配する 3 つの根本的なトレードオフが明らかになりました。

• 精度 - エントロピートレードオフ（KL 方向）: 逆 KL は一貫して高い単独精度（Avg@k）を達成しますが、予測エントロピーを崩壊させ、多様性を低下させ（Pass@k の低下）、その後の RL に対するモデルの安定性を損ないます。一方、前方 KL はエントロピーを保持し、より安定した RL 改善を支援しますが、初期の単独精度は低くなります。
• 品質 - 計算コストトレードオフ（プレフィックスソース）: 学生プレフィックス（オンポリシー）は、学生が実際に訪問する状態を監督することで、一致させたトレーニングステップにおいてより高い精度をもたらします。しかし、教師プレフィックス（オフポリシー）は、キャッシュされた教師ログオッズとオフライン軌道の再利用を可能にし、オンライン学生生成のコストを回避するため、一致させた FLOPs においてより計算効率的です。
• 精度 - 安定性トレードオフ（トレーニング長さ）: 長シーケンス蒸留（4096 トークン）は推論精度を向上させますが、特に逆 KL の下では、深刻なエントロピー崩壊と「長さインフレーション」（生成限界に達する応答）を引き起こします。短シーケンス学習はより安定していますが、精度は低くなります。

2. トレードオフをバランスさせる提案手法

これらの知見に触発され、著者らは 2 つの実践的な手法を提案します。

• KL 混合: 各プレフィックスにおける前方 KL と逆 KL の重み付け組み合わせです。著者らは、長シーケンス蒸留において前方に偏った混合が不可欠であることを発見しました。これは、逆 KL の精度の利点を維持しつつ、エントロピー崩壊と長さインフレーションを防ぎ、トレーニングダイナミクスを安定させます。
• エントロピーゲート付き長さカリキュラム: 短いホライズンから開始し、予測エントロピーが安定閾値を上回っている間のみシーケンス長を増加させるトレーニングスケジュールです。このアプローチは、固定された長ホライズントレーニングに伴うエントロピー崩壊を回避します。

結果

単独蒸留

• 逆 KLは、平均して +2.45 ポイントの gains を伴い、単独精度（Avg@k）において一般的に前方 KL を上回ります。
• 学生プレフィックスは、特に長シーケンス設定において（Avg@k で +3.55 の gains）、教師プレフィックスよりも精度において優れています。
• エントロピー崩壊: 長シーケンス設定における純粋な逆 KL は、エントロピーをゼロに近づけ、応答長を最大限界（例：8192 トークン）まで膨張させ、Pass@k を劣化させます。

蒸留後 RL（GRPO）

• 初期化の安定性: 前方 KLで初期化されたモデル（単独精度が低い場合でも）は、高いエントロピーを保持するため、GRPO 中により信頼性の高い改善を示します。
• 逆 KL の落とし穴: 逆 KL で初期化されたモデルは、しばしば高い精度で開始しますが、低エントロピーと探索能力の低下により、RL 中に劣化するかプラトーに陥ります。
• カリキュラムの成功: エントロピーゲート付き長さカリキュラムは、パフォーマンスを大幅に向上させます。固定された 4096 トークントレーニングと比較して、Avg@k を 3.6 ポイント、Pass@k を最大 5.8 ポイント向上させ、平均応答長を約3 倍（例：約 8185 トークンから約 2688 トークンへ）削減しました。

KL 混合の結果

長シーケンス設定において、高い前方 KL 重み（例：0.8 前方 / 0.2 逆）の混合は、エントロピーを安定させ長さインフレーションを防ぎつつ、逆 KL の精度の大部分を維持しました。

意義と主張

本論文は、SFT、DAgger、オフライン RL、OPD といった多様なトレーニングレシピを、2 つの直交する設計選択の具体的なインスタンスとして結びつける統合的な理論的枠組みを提供すると主張しています。プレフィックスソースと KL 方向を分離することで、著者らは、実践者が硬直的なパラダイムを超えて、以下を明示的にバランスさせる推論蒸留目的関数を設計できるようにすると論じています。

• 精度（逆 KL と長シーケンスを通じて）、
• 多様性/探索（前方 KL とエントロピー保持を通じて）、
• 計算効率（教師プレフィックス/キャッシングを通じて）、
• ダウンストリーム RL の学習可能性（安定したエントロピーを通じて）。

提案された手法（KL 混合とエントロピーゲート付きカリキュラム）は、特定されたトレードオフを navigat するための実用的なツールとして提示されており、さらにの方策最適化に対して正確かつ安定した、より小型の推論モデルのトレーニングを可能にします。著者らは、「最良の」単独目的関数が必ずしも RL パイプラインにとって最良の初期化とはならないことを強調し、ポストトレーニングプロセスの包括的な視点の必要性を指摘しています。