技術的概要：OISD – 言語モデルのオンポリシー内部自己蒸留

1. 問題定義

大規模言語モデル（LLM）に対する最近の強化学習（RL）ポストトレーニング手法、特に検証可能な報酬を用いたグループ相対方策最適化（GRPO）を利用する手法は、数学的推論において著しい進歩を遂げました。しかし、これらの手法は主に最終出力中心です。これらは、最終回答の正誤などのスパースな結果レベルの報酬を用いて最終動作方策を最適化しますが、中間層に符号化された予測信号、推論プロセス、内部表現を大きく見落としています。

知識蒸留（KD）は、密なトークンレベルの教師信号を提供するメカニズムを提供しますが、既存のパラダイムは特定の限界に苦しんでいます。

• オフポリシー蒸留: 学生が固定された教師によって生成された軌跡で訓練されるため、分布の不一致が生じます。
• オンポリシー蒸留: 不一致を軽減しますが、通常は外部の教師モデルに依存します。
• オンポリシー自己蒸留: 最近の手法は外部教師を排除しますが、モデルを監督するために外部から注入された特権情報（特権プロンプト、真の解答、または補助的フィードバックなど）に依存しています。

外部の特権データや複数の動作方策を必要とせず、モデルの自身の内部計算を教師信号の源として活用するギャップが存在します。さらに、先行研究は、推論表現が層を超えて体系的に変化すること（初期層では広範な候補予測から、深い層では収束した意思決定へ）を示しており、中間層を深い層によって導くことで推論の整合性を向上できることを示唆しています。

2. 手法：OISD フレームワーク

著者は、モデルの最終層を「分離された内部教師」とし、選択された中間層を「補助的内部学生」として扱う**OISD（オンポリシー内部自己蒸留）**を提案します。これは、ロールアウトおよび GRPO 最適化プロセス中の単一のフォワードパス内で発生します。

コアアーキテクチャ

• 単一動作方策: 最終層（$L$）は、ロールアウトの生成と GRPO 目的関数の最適化を行う唯一の方策のままです。
• 内部教師 - 学生: 同じフォワードパスにおいて、選択された中間層（$l < L$）の隠れ状態は、アンエンベディング行列を用いてトークン空間に投影されます（logit lensフレームワークに従う）。これにより、「教師」（最終層）と同じコンテキストから導出された仮想的な「学生」分布が作成されます。
• 外部特権の不在: 教師信号は、同じオンポリシーロールアウトにおけるモデル自身の内部状態から完全に導出され、外部の真の解答や特権プロンプトを必要としません。

整合性目的

OISD は、最終層から中間層へ予測信号を転送するために、2 つの相補的な内部整合メカニズムを採用します。

1. Logit 整合（「どのように考えるか」）:

   • 目的: 高レベルの推論行動と予測信念を転送する。
   • メカニズム: 中間層（$p^l_t$）と最終層（$q^L_t$）の温度スケーリングされたトークン分布間のジェンセン・シャノン（JS）ダイバージェンスを計算する。
   • 重み付け: 損失は符号付き GRPO 優位性（$A_t$）で重み付けされます。高い優位性を持つトークンは中間層を最終層の分布へ引き寄せ、低い優位性を持つトークンは悪い継続を強化しないよう押し戻します。
   • 数式: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Attention 整合（「どこを見るか」）:

   • 目的: 一貫した注意パターンと証拠ルーティング行動を強制する。
   • メカニズム: 中間層と最終層の注意重み間の JS ダイバージェンスを、メモリ管理のためにローカルウィンドウとグローバルなストライド位置を組み合わせたサンプリングされた因果キーの集合に対して計算する。
   • 重み付け: Logit 整合と同様に、損失は符号付き優位性で重み付けされます。正の優位性は、中間層が最終層と同じ証拠に注意を向けるよう促し、負の優位性は低報酬パターンとの合意を減らします。
   • 数式: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

訓練目的

総損失は、標準的な GRPO 損失と 2 つの内部整合損失を組み合わせます。
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
ここで、$\lambda_{\text{think}}$ と $\lambda_{\text{attn}}$ は蒸留の強度を制御します。

3. 主な貢献

1. 新しいパラダイム: 単一の動作方策を維持しつつ、モデル自身の内部計算を教師信号として活用する推論 RL パラダイムであるオンポリシー内部自己蒸留の導入。
2. OISD フレームワーク: 最終層を分離された教師、中間層を学生として使用し、外部の特権情報なしでオンポリシー予測信号を転送する具体的な実装。
3. 二重整合メカニズム: 2 つの相補的な目的関数の開発。
   • Logit 整合: 推論プロセス（「どのように考えるか」）を監督する。
   • Attention 整合: 証拠ルーティング（「どこを見るか」）を監督する。
4. 実証的検証: 4 つの数学的推論ベンチマークにおいて、強力な推論 RL ベースライン（GRPO および BuPO を含む）に対して大幅かつ一貫した改善を実証。

4. 実験結果

著者は、Qwen3-4B、Qwen3-8B、およびOctoThinkerモデルを、AMC23、MATH500、AIME24、AIME25の 4 つのベンチマークで評価しました。

• 性能向上:
  • Qwen3-4Bにおいて、OISD は平均スコア64.75を達成し、GRPO（55.08）を9.67 ポイント、BuPO（58.51）を6.24 ポイント上回りました。
  • 改善は特に AIME ベンチマークで顕著でした（例：AIME24 において GRPO より +14.31 ポイント）。
  • Qwen3-8Bにおいて、OISD は66.72という最高の全体平均を達成し、GRPO（64.23）および BuPO（66.36）を上回りました。
• Pass@K 分析: OISD は、すべての $K$ 値において GRPO よりも一貫して高い Pass@K 率を達成しました。これは、平均精度だけでなく、サンプリングされた解集合の全体的な品質の向上を示しています。
• アブレーション研究:
  • Logit 整合と Attention 整合の両方が正の貢献をしました。Logit 整合は AIME24 で最大の改善を示し、Attention 整合は平均的にわずかに良いパフォーマンスを示しました。
  • 両者を組み合わせることで最良の結果が得られ、それらの相補性が確認されました。
• OPSD との比較: OISD（64.75）は、特権的な真の解答に依存する最先端のオンポリシー自己蒸留手法である OPSD（56.63）を大幅に上回りました。これは、OISD の純粋な内部教師信号の有効性を浮き彫りにしています。
• 定性的分析:
  • 訓練ダイナミクス: OISD は、安定する前に、安定した報酬の改善とエントロピーの初期増加（より広範な探索を示唆）を示しました。
  • Attention 可視化: OISD の中間層は、訓練中に重要な証拠トークン（端点、上限など）に対する最終層の注意と徐々に整合しました。
  • Logit-Lens 可視化: OISD は、GRPO（しばしば断片的な中間予測を生み出す）と比較して、推論に重要なトークンにおいて中間予測と最終予測の早期の整合を示しました。

5. 意義と主張

本論文は、オンポリシー内部自己蒸留の調査を開始したと主張しています。その主な意義は以下の点にあります。

• 外部依存の排除: 訓練中に外部教師、特権プロンプト、または真の解答を必要とすることなく、密な内部教師信号を提供します。
• 方策の一貫性: ロールアウト生成と最適化の両方に単一の動作方策を維持し、オフポリシー手法の分布不一致の問題を回避します。
• 推論の強化: 中間表現を最終層の推論および注意パターンと整合させることで、OISD はモデルがより強く、よりグローバルに一貫した中間表現を学習することを促します。

著者は、検証可能な回答を持つ数学的推論に焦点を当てていること、および最適な層選択に関する深掘りした研究が欠けていることを限界として挙げています。今後の研究としては、より複雑な推論タスクへの一般化や、多層内部教師信号の探求が提案されています。