技術要約: $\beta$-OPSD

1. 問題提起

オンポリシー自己蒸留（OPSD）は、大規模言語モデル（LLM）の推論能力を向上させるための有望な手法として台頭しています。標準的なOPSDでは、学生モデルは自身がサンプリングした軌跡から学習する一方で、訓練中のみ利用可能な「特権的な教師（privileged teacher）」（正解の解または検証済みのトレースにアクセス可能な、同一アーキテクチャのモデル）から監督を受けます。

しかし、著者らは、標準的なOPSDは実際には極めて脆弱であり、信頼性を確保するために多大なエンジニアリングの努力を必要とすることを指摘しています。核心的な問題は、標準的なOPSDが教師を唯一の最適化対象として扱っており、学生を初期または参照ポリシーに繋ぎ止めるための正則化強度を暗黙的に固定してしまっていることです。この硬直した定式化には、学生が参照ポリシーから教師へとどの程度積極的に離れていくかを制御するメカニメントが欠けており、それがしばしば不安定性や不十分な収束を招いています。

2. 手法

本論文は、オンポリシー自己蒸留を、より広範なKL正則化方策最適化問題のファミリーとして再定義するフレームワークである $\beta$-OPSD を提案しています。その手法は、理論的なターゲットの導出、それを実用的な蒸留目的関数への変換、そしてクレジット割り当ての洗練という3つの主要な段階で進行します。

2.1. 理論的導出: 方策最適化としての $\beta$-OPSD

著者らは、標準的なOPSDが、まさに $\beta = 1$ の場合のKL正則化方策最適化ファミリーであることを証明しています。彼らは目的関数 $J_\beta(\theta)$ を以下のように定義します：
$$J_\beta(\theta) = \mathbb{E}_{y \sim \pi_\theta} \left[ \log \frac{p_T(y|x,c)}{\pi_{\text{ref}}(y|x)} \right] - \beta D_{KL}(\pi_\theta(\cdot|x) \parallel \pi_{\text{ref}}(\cdot|x))$$
ここで：

• $\pi_\theta$: 学生の方策（student policy）。
• $p_T$: 特権的な教師の方策（問題 $x$ と特権的情報 $c$ に条件付けられたもの）。
• $\pi_{\text{ref}}$: 参照方策（例：初期の学生モデル）。
• $\beta$: 制御可能な正則化パラメータ。

重要な洞察: $\beta=1$ のとき、参照項は相殺され、教師に対する標準的な逆KL最小化が回収されます。$\beta > 1$ の場合、目的関数はより保守的になり、参照方策からの逸脱に対してより重いペナルティを課します。

2.2. 最適な方策と幾何学的補間

著者らは、この目的関数に対する閉形式の最適方策 $\pi^\star_\beta$ を導出しています：
$$\pi^\star_\beta(y|x,c) = \frac{\pi_{\text{ref}}(y|x)^{1-1/\beta} p_T(y|x,c)^{1/\beta}}{Z_\beta(x,c)}$$
この解は、参照方策と特権的な教師との間の幾何学的補間を表しています。パラメータ $\beta$（具体的にはその逆数 $1/\beta$）は、この経路上の位置を決定する重みとして機能します。

• $\beta = 1 \implies$ ターゲットは教師。
• $\beta \to \infty \implies$ ターゲットは参照方策に近づく。

2.3. 実践的な訓練: スケジュールされたロジット補間

シーケンスレベルの正規化因子 $Z_\beta$ が計算困難であり、勾配の分散が高いことから、RL（強化学習）の目的関数を直接最適化することは計算コストが高くなります。代わりに、$\beta$-OPSDは、この理論的な解を効率的な蒸留ターゲットへと変換します。

• 局所近似: シーケンスレベルの正規化因子は計算困難であるため、本手法は各デコーディング・プレフィックス $h_t$ において、参照と教師のロジットを線形補間することで、最適方策を近似します：
  $$\tilde{p}_{\beta_k}(\cdot | h_t, c) = \text{softmax}\left( \left(1 - \frac{1}{\beta_k}\right) z_{\text{ref}} + \frac{1}{\beta_k} z_T \right)$$
• カリキュラム・スケジューリング: $\beta$ を固定するのではなく、本手法は、教師の重み $w_k = 1/\beta_k$ が訓練ステップに従って $w_{\text{start}}$ から $w_{\text{end}}$ まで線形に増加するスケジュールを採用しています。これにより、最初は学生に近いターゲットから始まり、徐々に強力な教師のガイダンスを取り入れていくという、スムーズなカリキュラムが作成されます。

2.4. Return-to-go によるクレジット割り当て

標準的なトークンレベルのOPSD更新は「近視眼的」であり、各トークンの損失を独立して扱います。これは、初期のトークンの選択が将来のプレフィックスや、その後のミスマッチにどのように影響するかを無視しています。
これに対処するため、$\beta$-OPSDはreturn-to-go (RTG) クレジット割り当てを導入しています。トークン $t$ の損失は、軌跡に沿った累積的な将来のミスマッチ（割引リターン）によって重み付けされます：
$$G_{t, \gamma} = \sum_{s=t}^{T} \gamma^{s-t} \rho_s$$
ここで $\rho_s$ はステップ $s$ における対数比のミスマッチです。これにより、初期のトークンが、後にシーケンス内で引き起こす分布的ミスマッチに対して、クレジット（または責め）を受け取ることが保証され、トークンレベルの更新がシーケンスレベルの目的関数と整合するようにします。

3. 主な貢献

1. 理論的統一: 本論文は、標準的なOPSDがより広範なKL正則化方策最適化ファミリーの $\beta=1$ のケースであることを確立し、暗黙的な正則化の選択を明示的かつ制御可能にしました。
2. $\beta$-OPSD フレームワーク: 参照方策と教師との間の幾何学的補間として最適な方策を定義し、それをスケジュールされたロジット補間を通じて実用的に実現する、一般化された目的関数を導入しました。
3. 効率的な訓練アルゴリズム: アドバンテージ推定のような複雑なRLメカニズムを必要とせず、安価なトークンレベルの蒸留を用いて、高コストなRL由来の解を近似する方法を提案しました。
4. シーケンスレベルの整合性: return-to-go 重みを用いたクレジット割り当てを導入することで、標準的なトークンKL更新の近視眼的な性質を修正し、シーケンスレベルの目的関数に対する不偏推定量を提供しました。

4. 実験結果

本手法は、Qwen3 モデル（パラメータ数 1.7B から 8B まで）を用い、数学的推論ベンチマーク（AIME 2024, AIME 2025, HMT 2025）で評価されました。

• 性能向上: $\beta$-OPSD は、一貫して vanilla OPSD、SFT、および GRPO ベースラインを上回りました。
  • Qwen3-1.7B において、$\beta$-OPSD は vanilla OPSD に対して平均性能 (avg@12) を 5.74 パーセントポイント向上させ、AIME 2024 では 9.16 ポイントの具体的な向上を実現しました。
  • より大きなモデル（4B および 8B）においても、競争力のある、あるいは優れた性能を維持し、全体的な平均をそれぞれ 1.76 および 1.66 ポイント向上させました。
• アブレーション研究:
  • ターゲット分布: 固定された教師ターゲットを、スケジュールされたロジット補間ターゲットに置き換えることで、大幅な性能向上（例：AIME 2024 で +6.03 ポイント）が得られ、スムーズな遷移の有用性が確認されました。
  • クレジット割り当て: ローカルなトークングラジエントの代わりに return-to-go 重みを使用することで、性能が向上しました（例：AIME 2025 で +5.55 ポイント）。これにより、将来を見据えた更新の重要性が検証されました。
  • 参照の構築: 最良の性能は、ダイナミックな学生（現在のモデル）と固定された教師（特権的な教師）を使用した場合に得られました。これは、安定した監督のアンカーを維持しながら、進化する学生の分布を追跡することが最適であることを示唆しています。

5. 意義と主張

本論文は、「自己蒸留から方策最適化へ、そして再び戻るための原理的なルート」を提供すると主張しています。その主な意義は以下の通りです：

• 安定性: $\beta$ を制御可能なパラメータとして扱うことで、本手法は vanilla OPSD の脆弱性に対処し、参照方策への近接性と教師のガイダンスとのトレードオフを行う、よりスムーズな最適化を可能にします。
• 効率性: 完全な RL の計算コストを伴わずに、方策最適化（シーケンスレベルの目的関数、KL 正則化）の利点を得ることができ、標準的な蒸留の効率性を維持しています。
• 汎用性: このアプローチは、学生と教師の間の分布間のギャップを縮小するメカニズムを提供します。これは、このギャップが大きい小規模モデルにとって特に有益です。

著者らは、スムーズなターゲットパスと軌跡レベルのクレジット割り当てが、オンポリシー蒸留のための効果的な原則であり、直接的な教師模倣に代わる堅牢な選択肢となる、と結論付けています。