近年、Large Language Model (LLM) の推論能力は飛躍的に向上しているが、そのトレーニング手法においては、効率性と汎化性能のバランスという古典的なトレードオフが依然として課題となっている。
従来のSupervised Fine-Tuning (SFT) やReinforcement Learning (RL) に代わる新たなアプローチとして、MetaとUCLAの研究チームが発表した「Self-Distilled Reasoner: On-Policy Self-Distillation (OPSD)」が注目を集めている。本稿では、このOPSDの概念、数学的定式化、およびその実験的成果について詳細に解説する。
既存のトレーニングパラダイムの課題
LLMの推論能力を向上させるための既存手法には、それぞれ固有の限界が存在する。
Supervised Fine-Tuning (SFT): 専門家によるデモンストレーション(正解データ)を用いて学習する手法。しかし、モデルは学習中に自分自身のエラーを見ることはなく、常に「正しいtrajectory」のみを与えられる。これにより、推論時に自身の生成したトークン列が学習分布から外れた際に、誤差が累積して回復不能になる「Exposure Bias(露光バイアス)」の問題を抱えている。
Reinforcement Learning (RL): GRPO (Group Relative Policy Optimization) などの手法は、オンポリシー(On-Policy)での探索を行うため汎化性能に優れる。しかし、これらは通常、最終的な答えが合っているかどうかという「疎(Sparse)」な報酬信号に依存している。中間の推論ステップが適切であったかどうかを評価するのは難しく、また学習を安定させるために多くのサンプル(ロールアウト)を必要とするため、計算コストが非常に高いという欠点がある。
これに対し、On-Policy Self-Distillation (OPSD) は、これらの課題を「自己蒸留」というエレガントな枠組みで解決しようとする試みである。
OPSDの核心:自己教師としてのLLM
OPSDの基本的なアイデアは非常にシンプルである。「正解を知っているモデルは、正解を知らない自分自身を指導できるか?」という問いに対する肯定的な回答が、この手法の根幹を成している。
人間が学習する際、問題を解けなかったとしても、正解と解説を見れば「なぜその答えになるのか」を理解し、思考プロセスを修正することができる。OPSDはこのプロセスを模倣する。
教師ポリシーと生徒ポリシー
OPSDでは、単一のモデルが教師(Teacher)と生徒(Student)の二役を演じる。重要なのは、両者が同じパラメータを共有している点である。違いは、入力されるコンテキスト(条件付け)のみにある。
- 生徒ポリシー \(p_S(\cdot \mid x)\): 推論時と同様に、問題文 \(x\) のみを与えられる。正解は知らない状態で回答を生成する。
- 教師ポリシー \(p_T(\cdot \mid x, y^*)\): 問題文 \(x\) に加えて、正解(Ground Truth)\(y^*\) を「特権情報(Privileged Information)」として与えられる。
正解を一から生成する(Generation)のは難しいが、正解を与えられた状態でその論理的整合性を説明する(Rationalization)のは、LLMにとって相対的に容易である。OPSDはこの性質を利用し、正解を知っている「教師としての自分」の確率分布に、正解を知らない「生徒としての自分」を近づけることで学習を行う。
メソドロジー:密なフィードバックによる蒸留
OPSDのトレーニングプロセスは以下の3ステップで構成される。
1. 生徒によるオンポリシーサンプリング
まず、生徒ポリシーが生徒自身の確率分布に従って回答候補(推論trajectory)\(\hat{y}\) を生成する。 \[\hat{y} \sim p_S(\cdot \mid x)\] これはSFTとは異なり、モデル自身が生成したデータを用いるため、Exposure Biasの問題を回避できる(On-Policy性)。
2. 教師と生徒の分布計算
次に、生徒が生成したtrajectory \(\hat{y}\) の各トークンステップ \(n\) において、教師と生徒それぞれの次トークン予測分布を計算する。ここで教師は正解 \(y^*\) を参照できるため、より適切な推論ステップに対して高い確率を割り当てることができる。
3. トークンレベルの分布マッチング
教師の分布 \(p_T\) をターゲットとして、生徒の分布 \(p_S\) を近づけるように学習を行う。具体的には、生成されたtrajectory上の各トークンにおいて、両者の分布間のダイバージェンス(乖離)を最小化する。
損失関数 \(\mathcal{L}(\theta)\) は以下のように定義される(一般化されたJensen-Shannonダイバージェンスなどを用いる)。
\[\mathcal{L}(\theta) = \mathbb{E}_{(x,y^*)\sim \mathcal{S}} \left[ \mathbb{E}_{\hat{y}\sim p_S(\cdot|x)} \left[ \frac{1}{|\hat{y}|} \sum_{n=1}^{|\hat{y}|} D\left(p_T(\cdot \mid x, y^*, \hat{y}_{<n}) \,\|\, p_S(\cdot \mid x, \hat{y}_{<n})\right) \right] \right]\]
ここで重要なのは、RLのような「正解/不正解」というバイナリな報酬ではなく、トークンごとの全語彙に対する確率分布(Full-vocabulary distribution)を通じて、教師から生徒へ「密(Dense)」な指導が行われる点である。これにより、どのステップでの推論が確からしいかという詳細なシグナルが伝達される。
実験結果と考察
論文では、Qwen3ファミリー(1.7B, 4B, 8B)を用いた実験が行われており、数学推論ベンチマーク(AIME, HMMT, AMO-Bench)において顕著な成果が報告されている。
圧倒的なトークン効率
OPSDの最大の利点はその効率性にある。GRPOと比較して、4〜8倍少ないトレーニングトークン数で同等以上の性能を達成している。 GRPOが報酬を得るために多数のサンプリング(例:1問あたり8〜16回の生成)を必要とするのに対し、OPSDは各ステップで教師からの密なフィードバックを得られるため、1回の生成(および短い生成長)でも十分に学習が進む。
モデル規模と「自己指導」の有効性
興味深い知見として、モデルの規模による効果の違いが挙げられる。 * 小規模モデル(1.7B): OPSDによる改善は限定的。 * 大規模モデル(4B, 8B): 顕著な性能向上が見られる。
これは、自己蒸留が成立するためには、教師役となるモデル自身に、正解 \(y^*\) を解釈して適切な推論プロセスを導き出すだけの十分な能力(Capacity)が必要であることを示唆している。モデルが未熟な場合、正解を見せられても適切な指導ができないのである。
初期のトークンが重要
また、蒸留においては推論の後半よりも前半のトークン(Early Tokens)が重要であるという仮説も提示されている。推論の初期段階は論理の分岐点(Branching Points)を含んでおり、ここで正しい方向へ舵を切れるかどうかが最終的な正解率を左右するため、教師からのガイダンスが特に効果を発揮すると考えられる。
今後の展望と課題
OPSDは、外部の教師モデルや複雑な報酬設計を必要とせず、LLMが自律的に推論能力を高めるための強力なフレームワークである。しかし、いくつかの課題や拡張の余地も残されている。
検証シグナルの統合: 現在のOPSDは分布マッチングに特化しており、生成された回答が最終的に正解したかどうかの検証シグナル(Verification Signal)を直接的には利用していない。これを組み合わせることで、さらなる精度向上が期待される。
カリキュラム学習: モデルの能力を超える難問に対しては、正解を与えられても教師ポリシーが機能しない可能性がある。モデルの成長に合わせて問題の難易度を調整するカリキュラム学習の導入が有効かもしれない。
結論
On-Policy Self-Distillation (OPSD) は、「正解を与えられたLLMは、正解を知らない自分自身よりも賢い」という前提に基づき、計算コストの高いRLやExposure Biasを持つSFTの代替となりうる有望な手法である。 モデルの大規模化が進む中、外部からの監督データに依存せず、自らの能力を使って自らを律するこのパラダイムは、AIの自律的な進化(Self-Improvement)に向けた重要なステップとなるだろう。
参考文献 * Zhao, S., et al. “Self-Distilled Reasoner: On-Policy Self-Distillation”. Siyan Zhao Blog