技術要約：ステップレベル・オンポリシー蒸留 (SOPD)

1. 問題提起

本論文は、大規模言語モデル（LLM）、特に推論およびエージェントタスクの学習における、現在の**オンポリシー蒸留（On-Policy Distillation: OPD）および教師あり微調整（Supervised Fine-Tuning: SFT）**のパラダイムにおける限界に対処している。

• 標準的なOPDの限界: OPDは、学生モデルを教師のロジット分布に適合させることで、オフポリシー手法における状態分布の不一致を回避するが（学生が生成した軌跡に基づく）、これはトークンレベルで動作する。その結果、「断片的な修正」が生じる。学生がエラーを起こした場合、教師は誤ったプレフィックス（接頭辞）に基づいた単一トークンのガイダンスを提供する。その結果、教師は完全で一貫した修復パスを展開することができない。なぜなら、後続のターゲットは、修正された軌跡ではなく、元の誤った学生の軌跡に基づいたままとなるからである。
• 標準的なSFTの限界: SFTは、一貫した長期間の監督を提供するが、本質的にオフポリシーである。これは教師が生成した軌跡で学習を行うため、訓練時の状態分布と、推論中に学生が実際に訪問する状態との間に不一致が生じる。
• 既存のハイブリッド手法の限界: DAggerやOn-Policy Expert Correction (OEC) のような手法は、教師の応答を学生の軌跡の中に挿入することでこのギャップを埋めようとする。しかし、これらの手法は本質的に学生が生成するはずの軌跡を変更してしまい、状態分布を変え、オンポリシーな性質を弱めてしまう。

2. 手法：ステップレベル・オンポリシー蒸留 (SOPD)

著者らは、SFTの長期間の修正能力と、OPDのオンポリシーな利点を補完する手法として、**ステップレベル・オンポリシー蒸留（Step-Level On-Policy Distillation: SOPD）**を提案している。

コアメカニズム

SOPDは、学生が生成した完全な軌跡に対して動作するが、それを自然なステップ（例：数学における推論の段落、あるいはエージェントにおける環境相互作用のターン）に分割する。

1. 学生のロールアウト: 学生モデル $\pi_\theta$ が完全な軌跡 $\tau_S = (s_1, s_2, \dots, s_K)$ を生成する。ここで各 $s_k$ は自然なステップである。各ステップの前のプレフィックス $h_k$（ステップ $k$ の前の状態）を記録する。
2. 教師へのクエリ: 学生が軌跡を完了した後、教師モデルに対してステップごとに一度クエリを行う。記録された各プレフィックス $h_k$ に対して、教師は単一のターゲットステップ $\tilde{s}_k \sim q_T(\cdot | h_k)$ を生成する。
   • 決定的なのは、教師はこのターゲットを、学生のプレフィックス $h_k$ のみに基づいて独立に生成することである。教師は以前の教師の出力を見ず、またその出力が学生の軌跡を変更することもない。
   • これにより、学生のステート占有率（$h_k \sim d_{\pi_\theta}$）を維持しつつ、特定のステップに対して一貫したマルチトークンの修正パスを提供する。
3. 学習目的関数: 学生は、プレフィックス $h_k$ が与えられたときに、教師が生成したステップ $\tilde{s}_k$ を予測するように学習される。損失関数はステップバランス・クロスエントロピーである：
   $$L_{SOPD} = \frac{1}{|S_B|} \sum_{(h_k, \tilde{s}_k) \in S_B} \frac{1}{|\tilde{s}_k|} \sum_{j=1}^{|\tilde{s}_k|} -\log \pi_\theta(\tilde{s}_{k,j} | h_k, \tilde{s}_{k,<j})$$
   この定式化は、各ステップ内の損失を平均し、次に全ステップにわたって平均をとるものであり、長い教師ステップが勾配更新を支配することを防ぐ。

実装の詳細

• パッキング・アテンション・マスク: 効率性を確保するため、入力プロンプト、完全な学生の軌跡、およびすべての教師のターゲットステップを単一のシーケンスにパッキングする。アテンション・マスクにより、各教師のターゲットは、プロンプト、対応する学生のプレフィックス、および自身の因果的トークンのみに注目できる。将来の学生のステップや他の教師のターゲットにはマスクがかかる。
• ブラックボックスへの適合性: SOPDは教師のロジットを必要とせず、教師が生成したテキストのみを必要とする。これにより、ブラックボックス蒸留のシナリオにも適用可能となる。
• 並列生成: 教師へのクエリは条件付きで独立している（固定された学生のプレフィックスに基づく）ため、学生のロールアウト完了後に並列で生成することができ、逐次的な教師・環境間の相互作用によるレイテンシを回避できる。

理論的限界

著者らは、SOPDが2つの粒度の限界を接続することを証明している：

• ワンステップ限界: 全レスポンスを単一のステップとして扱う場合、SOPDは標準的なSequence SFTに帰着する。
• ワン・トークン限界: すべてのステップが単一のトークンである場合、SOPDはサンプリングされたForward-KL OPDの目的関数を近似する。

3. 主な貢献

1. SOPDの提案: OPDのオンポリシーな状態被覆性と、SFTの一貫したマルチトークンによる指導を組み合わせた、新しい蒸留パラダイムの提示。
2. ブラックボックスへの適用性: ロジットを必要とする標準的なOPDとは異なり、SOPDは教師のテキスト生成のみを必要とするため、ブラックボックスの教師に適している。
3. 理論的統一: SOPDがステップの長さに応じてSFTに帰着するか、あるいはForward-KL OPDを近似するかを数学的に示し、これらの蒸留パラダイムに対する統一的な視点を提供した。
4. 実証的検証: 推論およびエージェントタスクにおいて、SFTおよびVanilla OPDの両方に対して優位性を示した。

4. 実験結果

著者らは、ALFWorld（インタラクティブ・エージェント・タスク）と数学的推論の2つの異なるドメインでSOPDを評価した。

ALFWorld (エージェント相互作用)

• 設定: 7BのRL教師（Qwen2.5-7B-RL）から3Bの学生（Qwen2.5-3B）を蒸留。
• 結果: SOPDは「Valid Seen」タスクで84.29%、「Valid Unseen」タスクで**82.09%**の成功率を達成した。
  • これは、Vanilla OPDと比較して、それぞれ**+18.57および+21.64**ポイントの改善である。
  • また、SOPDは環境相互作用の平均ラウンド数をOPDよりも約3.5〜4.3回減少させており、より効率的なポリシーであることを示している。
  • 「Hard」タスクにおいて、SOPDは学生の手法の中で最短の相互作用軌跡を維持しながら、最高の成功率を達成した。

数学的推論

• 設定: DeepMathデータセットを用い、4Bの非思考型学生（non-thinking student）を4BのRL教師（Qwen3-4B-Non-Thinking-RL-Math）から蒸留。
• ベンチマーク: AIME24, AIME25, HMMT25 (Feb), HMMT25 (Nov)。
• 結果: SOPDはすべてのベンチマークにおいてVanilla OPDを大幅に上回った。
  • 平均精度: SOPDは**57.7%の平均精度を達成し、OPDの47.7%**と比較して（約10パーセントポイントの向上）大幅に改善した。
  • 具体的な利得としては、AIME24で+9.9ポイント、HMMT25-Novで+13.1ポイントであった。
• 学習ダイナミクス: 分析によれば、学習が進むにつれて、学生モデルはより自然な推論ステップを生成するようになり、教師はステップあたりのより豊かな指導（より多くのトークン）を提供するようになった。これは、モデルがより構造化された推論プロセスを構築することを学習していることを示している。

5. 意義と主張

本論文は、SOPDが、状態分布の整合（OPD）と一貫した指導（SFT）のトレードオフを解決することにより、蒸留手法に関する新しい視点を提供すると主張している。

• 一貫した修正: トークンレベルのOPDとは異なり、SOPDは教師が、学生の軌跡自体を変更することなく、特定のセグメントに対して完全で実行可能な修正パスを提供することを可能にする。
• 効率性: エージェントタスクにおける教師の生成を学生の環境相互作用から切り離し、並列クエリを可能にすることで、SOPDはブラックボックス蒸留の実用的なコストを標準的なSFTと同程度まで低減させる。
• 汎用的な有効性: 結果は、ステップレベルの指導が、推論能力とエージェント相互作用ポリシーの両方を向上させるための堅牢なメカニズムであり、既存の最先端の蒸留および微調整のベースラインを凌駕することを示唆している。

著者らは、SOPDが学生が訪問した中間状態をブラックボックスの教師へのクエリへと効果的に変換し、学生の状態分布の完全性を保ちながら、高密度で局所的な指導を提供すると結論付けている。