技術要約: OPD+: オンポリシー蒸留におけるアドバンテージ設計の再考

1. 問題提起

オンポリシー蒸留（On-Policy Distillation, OPD）は、学生モデルが生成したロールアウトから得られる高密度なトークンレベルの教師信号を利用して、有能な教師言語モデルからベースとなる学生モデルへ能力を転移させる手法である。既存のOPDの実装（例：Tinker）は、多くの場合、カルバック・ライブラー（KL）ダイバージェンスを主要な目的関数として用いるが、そこには特定の設計上の選択が含まれている。すなわち、報酬項（学生と教師の尤度の対数比の差）に対して**勾配を停止（stop the gradient）**させている点である。

この設計は、Deep Q-Network（DQN）におけるターゲットネットワークと同様に、数値的な安定性を目的としている。しかし、著者らはこのアプローチに決定的な理論的欠陥があることを指摘している：

1. 勾配推定のバイアス: 一般的な$f$-ダイバージェンスにおいて、報酬関数は学生のポリシーパラメータ$\theta$に依存する。この報酬項に対して勾配を停止することは、報酬の$\theta$への直接的な依存性を無視することになり、数学的に誤った勾配推定を導く。
2. 性能の低下: このバイアスは、Reverse KL以外のダイバージェンス関数（例：Forward KLやJensen-Shannon Divergence）において、破滅的な学習崩壊を引き起こし、標準的なOPD設定下では学習不能な状態にする。Reverse KLの場合であっても、スコア恒等式（score identity）によってバイアスは理論上消失するものの、実験結果からは、このストップグラディエント設計が依然として性能を最適化できていないことが示唆されている。

2. 手法: OPD+ フレームワーク

本論文は、一般の$f$-ダイバージェンスに基づくオンポリシー蒸留の、修正されたアドバンテージ設計を再評価する最適化フレームワークである**OPD+**を提案している。

2.1 理論的基礎

著者らは、学生のポリシー$p_\theta$と教師のポリシー$q$の間の$f$-ダイバージェンスを最小化する問題として、OPDの目的関数を定式化している。目的関数は以下のように定義される：
$$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$$
ここで、$u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ はトークンレベルの密度比であり、$f$は$f(1)=0$を満たす凸関数である。

2.2 修正

論文では、この目的関数の真の勾配を導出している。標準的なOPDは、報酬 $r = -f(u_{\theta,n})$ を固定として扱う「ストップグラディエント」型のポリシー勾配推定器を使用している。著者らは、真の勾配には以下の2つの項が含まれることを示している：

1. 標準的なスコア関数項: $\mathbb{E}[r \nabla_\theta \log p_\theta]$
2. 直接的な報酬勾配項: $\mathbb{E}[\nabla_\theta r]$

一般的な$f$-ダイバージェンスにおいて、直接的な報酬勾配項は消失しない。具体的には、$\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$ となる。
したがって、正しいアドバンテージ関数は単なる $-f(u)$ ではなく、重み付きの項となる：
$$w_f(u) = -f(u) + u f'(u)$$

2.3 特殊ケース: Reverse KL

Reverse KLの場合、$f(u) = -\log u$ である。修正項は以下のようになる：
$$w(u) = -(-\log u) + u(-1/u) = \log u + 1$$
追加された $+1$ 項は定数ベースラインである。$\mathbb{E}[\nabla_\theta \log p_\theta] = 0$（スコア恒等式）であるため、この項は理論上、期待値において消失する。しかし、著者らは実用面において、この定数ベースラインが分散減少メカニズムとして機能し、収束速度と安定性を向上させると主張している。

2.4 実装

OPD+の実装には、既存のOPD実装におけるわずか1行のコード変更、すなわち生の負のダイバージェンス報酬 $-f(u)$ を、修正された重み $w_f(u)$ に置き換えるだけで済む。このフレームワークは、Forward KL、Reverse KL、およびJensen-Shannon Divergence（JSD）を含む様々な$f$-ダイバージェンスをサポートしている。

3. 主な貢献

1. 汎用的な最適化フレームワーク: 一般的な$f$-ダイバージェンスに基づくオンポリシー蒸留のための統一されたフレームワークを提供し、報酬のポリシーパラメータへの依存性を考慮した正しい勾配と損失関数を導出した。
2. バイアスの理論的証明: 標準的なストップグラディエント操作が、一般的な$f$-ダイバージェンスにおいてバイアスのある勾配推定をもたらし、報酬を固定とする一般的な慣行を無効にすることを数学的に証明した。
3. OPD+ アルゴリズム: Forward KLやJSDで見られる学習崩壊を排除し、Reverse KLの性能をも向上させる修正済みアルゴリズムであるOPD+を導入した。
4. 実験的検証: 本研究は、他の（正しく最適化された場合には）AlternativeなダイバージェンスがReverse KLよりも優れた性能を発揮し得ることを示し、Reverse KLがオンポリシー蒸留において唯一実行可能な指標であるという支配的な見解に疑問を投げかけている。

4. 実験結果

著者らは、QWEN3モデルファミリー（1.7Bおよび8Bパラメータ）を用い、競技数学（DeepMath, AIME 24/25, HMMT 25）およびエージェンティックなツール利用（Agentic Tool Use）という2つの異なるドメインでOPD+を評価した。

4.1 「学習不能」なダイバージェンスの救済

標準的なストップグラディエント設計（OPD）の下では、Forward KLとJSDは破滅的な失敗を示し、ツール利用タスクでは精度が0%に低下し、数学タスクではSFT初期化のベースラインにすら及ばなかった。

• OPD+ による修正: 修正されたアドバンテージ $w(u)$ を適用することで、OPD+はこの失敗モードを完全に排除した。
• JSDの性能: JSD+は安定化しただけでなく、標準的なReverse KLのベースラインを上回った。困難な AIME 25 ベンチマークにおいて、JSD+は 50.42% の精度を達成し、これはReverse KLの 48.54% を上回る。ツール利用タスクにおいても、JSD+は 60.29% に達し、Reverse KLのベースラインである 59.97% を上回った。

4.2 Reverse KLにおける改善

従来、Reverse KLが優位であるとされてきた領域においても、OPD+は厳密な経験的改善をもたらした：

• AIME 24: Reverse KL+ は 64.27% から 64.69% へと向上した。
• ツール利用: Reverse KL+ は 59.97% から 61.44% へと向上した。
• 収束性: 学習ダイナミクス（図1）は、OPD+のバリアントが初期の収束を加速させ、標準的なベースラインで見られる性能の停滞や低下を防ぐことを示している。

4.3 自己蒸留 (OPSD)

本手法は、モデルが学生と教師の両方の役割を果たす（教師は特権的な正解解に条件付けられている）オンポリシー自己蒸留（On-Policy Self-Distillation, OPSD）の設定でも検証された。OPD+は、すべてのダイバージェンスタイプにおいて、バニラのストップグラディエント版を一貫して上回った。ここでは修正されたReverse KLが、ダイバージェンスの中で最高の性能を示した。

5. 意義と主張

本論文は、他の（代替的な）$f$-ダイバージェンス（Forward KL, JSD）がオンポリシー蒸留において歴史的に低性能であった理由は、それらダイバージェンス自体の本質的な限界ではなく、ストップグラディエント設計によって引き起こされた勾配係数の誤指定による副作用であると主張している。

アドバンテージ関数を、$f$-ダイバージェンスの目的に数学的に忠実なものへと修正することで、OPD+は以下の実現を果たした：

1. 設計空間の拡大: 業界標準であるReverse KLを超えて、より多様な統計的距離の使用を可能にした。
2. 性能の向上: 推論およびツール利用のベンチマークにおいて、新たな最高性能（SOTA）を確立した。
3. 慣習への挑戦: Reverse KLが安定性の面で本質的に優れているという一般的な信念に対し、その安定性はダイバージェンスの選択によるものではなく、最適化手法のアーティファクトであった可能性を示唆することで、疑問を投げかけた。

著者らは、本研究がより数学的に正しく、経験的にも効果的なアルゴリズムを提供し、アドバンテージ関数の設計に関するさらなる研究を誘発する可能性があると結論付けている。また、モデルのスケーリング分析や、シングルトークン尤度とロジット蒸留の比較に関する制限事項についても認め、これらを今後の課題としている。