技術的概要：推論のための位置重み付きオンポリシー自己蒸留

問題定義

オンポリシー自己蒸留（OPSD）は、数学的推論を蒸留するための実用的な手法であり、学生モデルが自身の生成したロールアウト上で訓練され、privileged teacher（正解に条件付けられたモデルのコピー）がトークンレベルのターゲットを提供する。しかし、標準的な OPSD 目的関数は、生成されたすべてのトークンに一様前方 KL 監督重みを適用する。これは、学生が訪れるすべての接頭辞において、privileged teacher のターゲットが同様に信頼性が高いと暗黙的に仮定している。

最近の適応的蒸留手法は、教師エントロピー、教師 - 学生確率比、またはコンセンサスなどの局所シグナルに基づいて監督を調整することで、この一様性を緩和しようとする。しかし、著者らは、これらの局所的な曖昧さシグナルが教師の信頼性の十分な代理指標ではないと主張する。高い教師エントロピーは、2 つの異なる現象を示唆し得る：

1. 実行不可能な不確実性：教師は複数の次のトークンに確率を割り当てるが、そのうち 1 つ（あるいは 0 つ）のみが正しい解に至り、他は「行き止まり」である。
2. 良性の多様性：教師は複数の実行可能な解経路に確率を割り当て、それらすべてが正しさを維持する。

既存の手法はこれらを混同しており、良性の分岐点では監督を過剰に減衰させたり、教師が逸脱する実行不可能な分岐点では監督を減衰させられなかったりする可能性がある。核心的な問題は、実行不可能な不確実性と実行可能な多様性を区別し、トークンを適切に重み付けするためのシグナルを特定することである。

手法

1. 分岐実行可能性診断

教師の信頼性を実証的に調査するために、著者は分岐実行可能性診断を導入する。

• 設定：単一モデル（Qwen3-4B）を使用し、「正解を含む」教師プロンプトと「通常の問題」である学生プロンプトを定義する。
• プロセス：
  1. 学生が正しい「スピン（脊髄）」ロールアウトを生成する。
  2. 教師がこのスピン上の高曖昧さ位置（高い切り捨てエントロピー）を特定する。
  3. これらの位置において、教師は代替の次のトークンを提案する。
  4. モデルは、学生プロンプト＋スピン接頭辞＋強制された代替トークンから、privileged な情報なしに続行を強制される。
  5. ラベリング：強制された続行が正しい答えを回復できない場合、その分岐は実質的不確実（real-uncertain）（低信頼性）とラベル付けされる。答えを回復できる場合は、多様性（diversity）（良性）とラベル付けされる。
• 発見：この診断は、強い位置構造を明らかにする。実質的不確実な分岐はシーケンスの早期に集中するのに対し、実行可能な多様性の分岐は後期に多く見られる。
• 予測能力：方向付けられたシーケンス内位置スコア（早期位置＝低信頼性）は、実質的不確実候補と多様性候補を区別する際に AUROC 0.83を達成する。これに対し、標準的な局所不確実性メトリクス（予測エントロピー、MC-dropout 相互情報量、ディリクレ精度）は、偶然に近い性能（AUROC $\le$ 0.57）しか示さない。

2. 位置重み付きオンポリシー自己蒸留（PW-OPSD）

早期の高曖昧さ位置が構造的に信頼性が低いという発見に触発され、著者はPW-OPSDを提案する。

• 中核メカニズム：PW-OPSD は、標準的な OPSD 構成（学生ロールアウト、privileged teacher パス、語彙ごとのクリップ前方 KL）を維持しつつ、外部トークン集約を修正する。
• 重み付けスケジュール：一様重みの代わりに、シーケンス内の正規化された位置に基づいた増加シグモイド関数を適用する：
  $$w_{i,t} = w_{min} + (1 - w_{min})\sigma\left(\frac{r_{i,t} - \tau}{s}\right)$$
  ここで、$r_{i,t}$は正規化された位置である。
• 設計根拠：このスケジュールは、診断が高信頼性を示す後期のトークンには高い重みを、高信頼性の低い（診断が示す）早期のトークンには低い重みを割り当てる。フロア（$w_{min}$）は、早期のトークンが完全に無視されないことを保証する。
• 削減：この手法は、バッチ全体での平均化に先立ち、各シーケンス内のトークン損失を平均化し、トークン数のみによる勾配の支配を防ぐ。
• 効率性：追加の教師パス、補助検証器、複雑なサンプリングを必要としない。重みは位置の決定論的関数である。

主要な貢献

1. 信頼性の診断：本論文は、推論における教師トークンの信頼性が、局所的（エントロピー依存）ではなく、軌道構造的（位置依存）であることを実証的に示す分岐実行可能性診断を導入する。
2. エントロピーに対する位置の優位性：著者らは、位置が局所不確実性スコア（AUROC $\le$ 0.57）よりも教師の信頼性の予測因子として著しく強力（AUROC 0.83）であることを示し、エントロピーベースの適応的重み付けが最適であるという仮定に挑戦する。
3. PW-OPSD アルゴリズム：前方 KL 損失に位置依存重みを適用する、シンプルでプラグアンドプレイ可能な OPSD の修正版であり、追加の計算オーバーヘッドなしに推論性能を向上させる。
4. 汎化性：この手法は、モデルごとのハイパーパラメータ調整なしに、異なるモデルファミリー（Qwen、Llama、OLMo）およびスケール（4B から 8B パラメータ）で検証されている。

実験結果

主要評価（Qwen3-4B）

最大生成長 38,912 トークンで MATH-500、AIME 2024、AIME 2025、HMMT 2025 上で評価：

• AIME 2024：PW-OPSD（Moderate）は、標準 OPSD に対して Avg@12 を**+1.0 ポイント**改善（76.20 対 75.19）。
• AIME 2025：PW-OPSD は Avg@12 を**+1.1 ポイント**改善（67.78 対 66.67）。
• HMMT 2025：「Aggressive」スケジュール変種は、OPSD に対して Avg@12 を**+1.48 ポイント**改善。
• ベースライン：エントロピー（EOPD）または方策勾配報酬（REOPOLD）を使用する適応的ベースラインは、AIME ベンチマークにおいて一般的に OPSD よりも劣り、位置診断によって特定された信頼性パターンを捉えられなかった。
• MATH-500：性能は OPSD と同等（ほぼ飽和状態）であり、この手法がより簡単なタスクで性能を低下させないことを示している。

クロスモデル転送

Qwen3-4B の診断から導出された「Moderate」スケジュールを、再調整なしに適用した：

• DeepSeek-R1-Distill-Llama-8B：Avg@12 を**+0.35 ポイント**改善。
• Olmo-3-7B-Think：Avg@12 を**+0.50 ポイント**改善。
  これは、位置信頼性の事前分布がモデルファミリーやアーキテクチャを超えて頑健であることを示唆する。

除去実験

• スケジュール感度：Mild、Sharp、Aggressive など、さまざまなスケジュール設定がすべて平均して OPSD を上回り、位置重み付け概念の頑健性を確認した。
• 削減対位置付け：著者らは、改善が位置重み付けとシーケンスごとの損失削減の併用に依存していることを発見した。一方の軸のみを適用しても、完全な性能向上は得られなかった。

意義と主張

本論文は、推論蒸留における教師トークンの信頼性が軌道構造的であると主張する。主な意義は、適応的蒸留において、局所的な曖昧さ（エントロピー）から構造的な信頼性（位置）へと焦点を移行させる点にある。

• 控えめな主張：著者らは明示的に、PW-OPSD が性能を向上させるものの、その利益は「控えめ」である（例：AIME で +1.0 から +1.1 ポイント）と述べている。PW-OPSD は「プラグアンドプレイ」的な近似であり、将来的には位置に基づいて前方 KL と逆 KL を混合するなど、より洗練された目的関数を設計する余地があると認めている。
• 追加コストなし：重要な主張として、これらの改善は追加の教師計算や補助検証ステップなしで達成されており、既存の蒸留パイプラインに対する非常に効率的な修正である。
• 理論的洞察：この研究は、蒸留を信頼性重み付きリスクとして解釈する理論的枠組みを提供する。その目標は、教師のターゲットが正しい解を維持するに役立つ事後確率に基づいてトークンを重み付けることである。診断は、位置がこの事後確率の強力な代理指標であるという実証的証拠を提供する。

コードは https://github.com/SaFo-Lab/PW-OPSD で利用可能。