🤖 AI

Trajectory-Refined Distillation

本論文は、教師の指導の下で生徒のロールアウトに対して軌跡レベルの修正を適用することにより、構造的な「プレフィックス失敗（prefix failure）」問題を軽減し、様々なベンチマークおよびモデルスケールにおける推論の正確性と網羅性を向上させる、新しいオンポリシー蒸留手法であるTrajectory-Refined Distillation (TRD) を提案する。

原著者： Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

公開日 2026-06-09

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、ある複雑な数学の問題を解く方法を学生（AI）に教えていると想像してください。あなたには、正解を知っている非常に優秀な教師（より高度なAI）がいます。

論文で説明されている標準的な手法である**オンポリシー蒸留（On-Policy Distillation: OPD）**では、プロセスは次のように進みます：

学生が自力で問題を解こうと試みます。
学生が手順を書き進めるにつれて、教師はそれを見守り、「今あなたが書いたこの単語に対して、次に続くべき言葉の確率はこれくらいですよ」と伝えます。
学生は、すべての単語における教師の確率に一致するように、自分の脳を調整しようと努めます。

問題：「プレフィックス失敗（Prefix Failure）」（誤った方向への転換）

この論文は、標準的な手法における重大な欠陥として「プレフィックス失敗（Prefix Failure）」と呼ばれる現象を指摘しています。

学生が問題を解くための道を歩いていると想像してください。ステップ3で、学生は道を間違えてしまいました。今や、行き止まりの通りに迷い込んでいます。

教師のジレンマ： 教師は学生が迷っているのを見ています。教師は、正しい道は「あ、実は……」という言葉から始まり、全く別の方向へ進むものであることを知っています。
混乱： しかし、教師は（現在の学生の場所である）行き止まりを見続けなければなりません。その結果、教師の助言は、「すでにここにいるのだから、この行き止まりに留まりなさい」という指示と、「すぐに引き返せ」という指示が混ざり合った、混乱したものになってしまいます。
結果： 学生は混乱した信号を受け取ります。学生は教師から学ぼうとしますが、学生が間違った道に捕まっているために、教師の助言が断片化されてしまうのです。学生は「正しい道に戻る方法」を学ぶのではなく、「迷っている状態において上手くなる方法」を学んでしまいます。論文ではこれを「バイモーダルな混合（bimodal mixture）」と呼んでいます。つまり、教師が同時に二つの相反することを教えようとしているため、学生が混乱してしまう状態のことです。

既存の解決策は、特定の単語に対する「成績（損失関数）」を微調整しようとしましたが、学生が依然として間違った道に捕まっているという根本的な事実は解決できませんでした。

解決策：軌跡洗練蒸留（Trajectory-Refined Distillation: TRD）

著者らは、新しい手法である**軌跡洗練蒸illation（TRD）**を提案しています。学生の支離滅裂なドラフトを単語ごとに採点するのではなく、ワークフロー自体を変更します。

ドラフト： 学生はまず、問題に対して最初の一撃（「生のロールアウト」）を与えます。
書き直し： 採点を行う前に、教師はその支離滅裂なドラフトを受け取り、それを書き直します。教師は間違った方向を修正し、論理を正し、学生の出発点に基づいた上で、クリーンで完璧な解決策を作成します。
レッスン： 学生は、この支離滅裂なオリジナルではなく、このクリーンに修正されたバージョンから学びます。

比喩：

従来の方法： あなたが物語を書いていて、途中にプロットの穴（矛盾）があるとします。編集者はそれを読み、「第3パラグラフの『猫』という単語については、『犬』と書くべきでした」と言います。しかし、あなたはなぜ自分が「猫」と書いたのかを理解しようと、依然として第3パラグラフで立ち往生したままです。これは混乱を招きます。
TRDの方法： あなたがプロットの穴がある物語を書きます。編集者はあなたのドラフトを受け取り、そのプロットの穴を修正し、物語が完璧に流れるように中盤をすべて書き直します。その後、あなたは次に物語をより良く書けるようになるために、編集者の「書き直されたバージョン」を研究します。

なぜこれがより優れているのか

根本原因の解決： 学生が学習を始める前に「誤った方向」を修正することで、教師の助言が混乱することを防ぎます。これにより、明確で単一の道筋が得られます。
新しい経路の探索： たとえ学生が最初から正解に辿り着いていたとしても、教師はより短く、あるいはよりスマートな解法を示すことができます。これにより、学生はより柔軟になれます。
効率性： 論文では、「書き直された」経路は学生の元のとりとめもない試行よりもはるかに短く明快であることが多く、それが実際にはトレーニングを高速化させたことが示されています。

結果

研究者らは、この手法を難易度の高い数学コンテスト（AIMEやHMMTなど）やコーディングの課題でテストしました。その結果、以下のことが判明しました：

新しい手法（TRD）は、従来の手法を一貫して上回りました。
特に、AIが通常行き詰まってしまうような最も困難な問題を解く際に、非常に効果的でした。
教師が（外部の）より大きなAIであっても、あるいは同じAIが「特権情報（privileged information）」を用いることで自身の「賢い」教師として機能する場合であっても、同様に機能しました。

要約すると、この論文は、AIを効果的に教えるためには、単語ごとに間違いを採点するのではなく、まず教師に経路全体を修正させ、それから学生にその「修正された旅路」から学ばせるべきであると主張しています。

技術要約：軌跡精緻化蒸留 (Trajectory-Refined Distillation: TRD)

1. 問題提起：オンポリシー蒸留におけるプレフィックス失敗 (Prefix Failure)

本論文は、大規模言語モデル（LLM）におけるオンポリシー蒸留（OPD）およびオンポリシー自己蒸留（OPSD）における根本的な構造的限界として、「プレフィックス失敗（prefix failure）」と名付けられた現象を特定している。

標準的なOPD/OPSDでは、生徒モデルがロールアウト（軌跡） $y_o$ を生成し、教師モデルがその特定のロールアウトに基づいた高密度なトークン単位の教師信号を提供する。著者らは、生徒のロールアウトが誤った推論経路（「失敗したプレフィックス」）を辿った場合、教師の監督信号が構造的に損なわれると主張している。

二峰性の教師混合 (Bimodal Teacher Mixture): プレフィックスが誤っている場合、教師の分布は2つのモードの混合体となる。一つは（シーケンスの一貫性のために）失敗したプレフィックスを継続するもの、もう一つは正しい解へと転換するものである。
断片化された勾配 (Fragmented Gradients): 標準的なトークンレベルの損失関数（KLダイバージェンス）は、凍結された失敗した軌跡に沿って監督を行う。その結果、勾配更新は断片化される。教師は、展開されるべき修正経路に条件付けるのではなく、誤ったコンテキスト（失敗したプレフィックス）に条件付けられた状態で、修正トークン（例：「Wait」や「Actually」）を繰り返し推奨してしまう。
トークンレベルの修正の無効性: 既存の救済策（損失クリッピング、Top-K 截断、重み付けなど）は、トークンレベルの損失に対して作用する。これらは「情報量の多い」トークンを選択したり「ノイズの多い」トークンを抑制したりしようとするが、根本的な原因である失敗したプレフィックス自体が構造的に維持されているという問題に対処できていない。生徒に届けられる監督ペアは理想的な修正経路から逸脱しており、モデルがエラーから効果的に回復することを妨げている。

2. 手法：軌跡精緻化蒸留 (Trajectory-Refined Distillation: TRD)

プレフィックス失敗を軽減するために、著者らは介入の対象をトークンレベルから軌跡レベルへとシフトさせる手法である**軌跡精緻化蒸留（TRD）**を提案している。

コアメカニズム

生の生徒のロールアウト $y_o$ を直接蒸留に使用する代わりに、TRDは損失計算の前に精緻化ステップを導入する。

サンプリング: 生徒モデル $\pi_\theta$ が生のオンポリシー・ロールアウト $y_o \sim \pi_\theta(\cdot | x)$ を生成する。
精緻化 (Refinement): 教師モデル $\pi_T$ $π_{T}$ （または、正解 $y^*$ $y^{*}$ のような特権情報に基づき条件付けられた生徒自身（OPSDの場合））が、精緻化された軌跡 $y_r$ $y_{r}$ を生成する。
- OPDの場合: $y_r \sim \pi_T(\cdot | x, y_o)$
- OPSDの場合: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$
- 決定的なのは、教師が参照解（reference solution）に導かれながら生徒のロールアウトを修正するように促されることであり、これにより、生徒の推論パターンが有効な箇所は維持しつつ、誤った部分のプレフィックスを実質的に「書き換える」ことが可能になる。
蒸留: 精緻化された軌跡 $y_r$ が、標準的な分布一致損失（フル語彙を用いたForward KL）の監督ターゲットとして使用される。

理論的正当性

TRDは、監督のコンテキストが失敗したプレフィックスではなく、修正経路に沿って成長するようにすることで、「断片化された勾配」の問題に対処する。

理想的な勾配の回復: $y_r$ に沿って監督を行うことで、各トークンのコンテキストに修正されたプレフィックスが含まれるようになり、教師が（コンテキストの不一致による反復的な信号 $g_{frag}$ ではなく）一貫した多段階の修正信号（ $g_{ideal}$ ）を提供できるようになる。
探索 (Exploration): $y_o$ が正しい場合であっても、TRDは教師によって示唆される代替の有効な導出過程を生徒に提示し、生徒が学習する正しい推論経路の分布を広げる。

アルゴリズム

プロセスはアルゴリズム1に要約される：

$y_o \sim \pi_\theta(\cdot | x)$ をサンプリングする。
$y_r \sim \pi_T(\cdot | x, y_o)$ を生成する。
$\pi_\theta$ と $y_r$ 上の教師の分布との間のKLダイバージェンスを用いて $\theta$ を更新する。

3. 主な貢献

プレフィックス失敗の特定: 本論文は、オンポリシー蒸留における劣化の構造的な原因として「プレフィックス失敗」を定式化し、なぜトークンレベルの介入が推論エラーからの回復に失敗するのかを説明した。
軌跡レベルの修正: TRDは、過度な計算コストをかけることなく、オンポリシーのサポート制約を尊重しながら、失敗したプレフィックスの回復を明示的にターゲットとした最初の軌跡レベルの最適化設計として提案された。
デュアルレジームにおける改善: TRDは以下の2つのレジームにおいて性能を向上させる：
- エラー修正: 蒸留前に問題のあるプレフィックスを修正することで、プレフィックス失敗を軽減する。
- 探索: 正しいロールアウトであっても、代替の有効な導出過程を生徒に提示することで、推論のカバー範囲を広げる。
自己蒸留への適用可能性: 本手法は自然にOPSDへと拡張可能である。そこでは、生徒が特権情報に基づき自身の教師として機能するため、別途の大きな教師モデルを必要としない。

4. 実験結果

著者らは、Qwen3モデル（1.7B, 4B, 8B）を用い、5つの競技数学ベンチマーク（AIME24/25, HMMT25, BeyondAIME, AMOBench）およびコード生成ベンチマーク（HumanEval+, MBPP+, LiveCodeBench）にわたってTRDを評価した。

主な知見

優れた性能: TRDは、OPDおよびOPSDの両方の設定において、従来のベースライン（Forward/Reverse KL、およびクリッピング/截断の有無）を一貫して上回る性能を示した。
難易度の高いタスクでの大幅な向上: 改善は最も困難なベンチマークで最も顕著であった。
- AMOBench: OPSD設定において、Qwen3-8BのPass@16で、ベースラインモデルと比較して約50%の相対的な向上を達成した。
- AIME24/25: TRDは一貫した利得を示し、特に小規模モデル（例：Qwen3-1.7BにおけるAIME24での+4.6%）において顕著であった。
安定性: 特定のベンチマークで性能が低下するベースライン（例：AIME24で性能が低下したReverse KL）とは異なり、TRDがベースラインを下回ることはなかった。
効率性: OPSDの設定において、精緻化された軌跡 $y_r$ は生のロールアウト $y_o$ よりも大幅に短く（中央値で約9倍の圧縮）、追加のサンプリングステップがあるにもかかわらず、トレーニングの実行時間を約60%削減した。
コード生成: TRDはHumanEval+およびMBPP+においてベースラインと同等またはそれを上回ったが、LiveCodeBenchでの利得は限定的であり、教師の精緻化能力がタスクの難易度によって制限される可能性を示唆している。

5. 重要性と主張

本論文は、TRDがLLMのポストトレーニング戦略における必要な転換を象徴していると主張している。著者らは以下のように述べている：

トークンレベルの介入は不十分である: プレフィックス失敗の構造的な性質には、軌跡レベルの修正が必要である。
オンポリシーのサポートが維持される: TRDは、固定されたエキスパート・ポリシーを使用することなく（これとは対照的に）、これらの修正を達成しており、計算効率の高い標準的な学習パラダイムと互換性がある。
自己改善は実行可能である: OPSD設定におけるTRDの成功は、モデルが別途のより大きな教師モデルを必要とせず、特権情報（グラウンドトゥルース）を用いて自身の推論経路を効果的に精緻化できることを示しており、軽量な自己改善への道筋を示している。

この研究は、将来の蒸留手法は、個々のトークンの統計的な整合性だけでなく、推論の軌跡の構造的な完全性を考慮しなければならないことを示唆している。