Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

本論文は、教師モデルからの報酬を柔軟に活用する「REOPOLD」というフレームワークを提案し、オンポリシー蒸留の不安定さを解消するとともに、数学や視覚、エージェントタスクにおいて従来の強化学習手法を凌駕するサンプル効率と推論時のスケーリング性能を実現することを示しています。

要約

小さな頭脳を「天才」の思考法で育てる：REOPOLD の仕組み

この論文は、**「小さな AI（学生）に、巨大な AI（先生）の『考える力』をいかに効率的に、かつ安定して教えるか」**という問題を解決する新しい方法「REOPOLD」について書かれています。

まるで、**「天才的な数学者の思考プロセスを、小学生に教える」**ような場面を想像してください。

---

1. 従来の問題点：「真似しすぎ」の罠

これまでの方法（オンポリシー蒸留）では、学生 AI は先生 AI の答えを**「ただひたすら真似」**しようとしていました。
しかし、これには 2 つの大きな問題がありました。

• 問題 A：「間違っているのに、無理やり真似させられる」

  • アナロジー： 先生が「この問題は 3 だ」と言っているのに、学生が「いや、これは 2 じゃないかな？」と少し違う答えを出したとします。従来の方法では、先生が「バカ！」と激しく怒り（強い罰）、学生は「2」という考え方を完全に捨ててしまいます。
  • 結果： 学生は「2」という可能性を失い、思考の幅が狭まってしまいます。これを**「エントロピーの崩壊（思考の多様性の消失）」**と呼びます。

• 問題 B：「無駄な練習に時間を浪費」

  • アナロジー： 先生と学生が「1+1=2」という簡単な部分で意見が一致している場合、従来の方法では「正解だ！」と何度も練習させられます。でも、これは学生にとってすでにわかっていることなので、**「学習効果はゼロ」**なのに、計算リソースだけ消費してしまいます。

2. 新手法「REOPOLD」の 3 つの魔法

REOPOLD は、この「ただの真似」を「賢い指導」に変える 3 つの工夫をしています。

① 罰の「上限」を決める（Reward Clipping）

• 仕組み： 先生が学生を怒るとき、**「怒りの強さに上限」**を設けます。
• アナロジー： 先生が「お前、バカだ！」と叫んでも、その叫び声の音量（罰の強さ）が一定以上にならないように制限します。
• 効果： 学生が「2」という答えを出したとき、先生が「100 点満点で 0 点！」と極端に怒るのを防ぎます。これにより、学生は「2」という可能性を完全に捨てずに済みます。思考の多様性が守られ、安定して学習できます。

② 「重要な瞬間」だけ教える（Entropy-Guided Sampling）

• 仕組み： 先生と学生の意見が**「ぶつかり合う（迷っている）瞬間」**だけを学習のターゲットにします。
• アナロジー：
  • 「1+1=2」のような簡単な問題は、**「スキップ（飛ばし）」**します。
  • 「この複雑な図形、どう解くんだ？」と学生が迷い、先生も「うーん、ここは難しいな」と考えるような**「高エントロピー（混乱度が高い）」な瞬間**だけ、先生は「こう考えたらいいよ」と指導します。
• 効果： 無駄な練習を省き、「本当に必要な部分」に集中して学習できるため、「サンプル効率（少ないデータで学ぶ力）」が 6〜12 倍に向上しました。

③ 「探索」と「洗練」の 2 段階トレーニング

• 仕組み： 学習を 2 つのフェーズに分けます。
  1. 探索フェーズ（冒険）： 最初は「どんな答えが出ても OK」として、学生に自由に考えさせます。先生は「間違っても怒らない」ようにします。
  2. 洗練フェーズ（磨き上げ）： 慣れてきたら、先生は「ここは間違っているよ」と厳しく指導し、正解への道筋を絞り込みます。
• アナロジー：
  • 最初は**「自由な発想でアイデア出し」**（アイデアが枯渇しないように）。
  • 後から**「プロのチェック」**（質を高める）。
• 効果： 最初から厳しくしすぎると思考が止まってしまうのを防ぎ、最終的に高い精度を達成します。

---

3. 実際の成果：小さなモデルが巨人に追いつく

この方法を使うと、驚くべき結果が生まれました。

• 数学の問題： 従来の方法より6.7 倍〜12 倍少ないデータで、同じレベルの性能を達成しました。
• 視覚推理（図形やグラフ）： 320 億パラメータという「巨人」の先生 AI に匹敵する性能を、70 億パラメータの「小柄な」学生 AIが達成しました。
  • スピードアップ： 答えを出すまでの時間が、従来の方法より3.3 倍速くなりました。
  • アナロジー： 「巨人の先生と同じくらい賢い答え」を、「小さな子供が、巨人より速く、少ないエネルギーで」導き出せるようになったのです。

まとめ

REOPOLD は、**「AI に『思考』を教える際、ただ真似させるのではなく、先生と学生の『思考のズレ』を賢く利用する」**という新しいアプローチです。

• 怒りすぎない（罰の制限）
• 無駄な練習を省く（重要な瞬間だけ教える）
• まずは自由に考えさせる（探索フェーズ）

この 3 つの工夫により、小さな AI でも、大きな AI に負けない「考える力」を、効率的に身につけることができるようになりました。これは、高性能な AI を安価で高速に使えるようになるための重要な一歩です。

技術概要

論文「Scaling Reasoning Efficiently via Relaxed On-Policy Distillation」の技術的サマリー

本論文は、大規模言語モデル（LLM）の推論能力を、制約のある小規模モデル（SLM）へ効率的に転移させるための新しいフレームワーク**「REOPOLD (Relaxed On-Policy Distillation)」**を提案するものです。従来のオンポリシー蒸留（On-Policy Distillation）が抱える不安定性や負の転移（Negative Transfer）の問題を、強化学習（RL）の視点から分析し、緩和された蒸馏アプローチによって解決しています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

---

1. 背景と問題定義

• 小規模モデルにおける推論の課題:
  大規模モデル（Teacher）は強化学習（RL）やテスト時のスケーリングにより高度な推論能力を獲得していますが、小規模モデル（Student）はリソース制約により直接の RL 最適化が困難です。そのため、Teacher の軌跡から学習する「オンポリシー蒸留」が注目されています。
• 既存手法の限界:
  従来のオンポリシー蒸留（特に RKL 最小化に基づくもの）は、以下の問題に直面しています。
  1. 不安定性と負の転移: 学生モデルが初期化モデルよりも性能が低下する現象。
  2. エントロピーの崩壊: 学習の早期に多様性が失われ、狭い解空間に収束してしまう。
  3. 信号の非効率性: Teacher と Student の分布が一致しているトークン（低エントロピー）では学習信号がゼロに近く、一方で分布が乖離しているトークン（高エントロピー）では負の報酬が極端に大きくなり（Heavy-tailed）、勾配爆発や不安定化を招く。
  4. 厳密な模倣の制約: 学生が Teacher を過度に模倣しようとするため、探索能力が阻害され、RL 的な最適化の恩恵を受けられない。

2. 提案手法：REOPOLD

REOPOLD は、オンポリシー蒸留を「方策最適化（Policy Optimization）」の一種として再解釈し、Teacher-Student の対数尤度比を「トークンレベルの報酬」と見なすことで、現代の RL 技術（クリッピング、動的サンプリングなど）を統合したフレームワークです。

2.1 理論的基盤

• 停止勾配（Stop-Gradient）の導入:
  従来の RKL 最小化では、報酬項自体が方策パラメータに依存し、勾配計算が不安定になります。REOPOLD は、Teacher-Student の対数尤度比 $R_{i,t}$ に対して停止勾配演算子を適用し、これを「固定された外部的報酬」として扱います。これにより、勾配の分散を抑制し、安定した最適化の基盤を確立します。

2.2 主要な技術的構成要素

REOPOLD は、以下の 3 つのメカニズムを組み合わせて学習信号を「緩和（Relaxed）」し、制御します。

1. ミックスベースの報酬クリッピング (Mixture-Based Reward Clipping):

   • 問題: Teacher が確率をほぼゼロとみなすトークンで、対数尤度比が $-\infty$ に発散し、勾配を破壊する（Heavy-tailed negative rewards）。
   • 解決: Teacher と Student の分布の凸結合（Convex Mixture）に基づく理論的限界値 $\log \frac{\lambda}{1-\lambda}$ を用いて、負の報酬をクリップします。これにより、極端な罰則を回避し、学習の安定性を保ちます。

2. エントロピーに基づくトークンレベルの動的サンプリング (Entropy-Guided Token-Level Dynamic Sampling):

   • 問題: 低エントロピー（確定的）なトークンでは Teacher と Student が一致しており学習信号が希薄（Gradient Dilution）になる一方、高エントロピーなトークン（推論の分岐点）で重要な学習信号が含まれます。
   • 解決: トークンのエントロピーに基づき、学習対象を「高エントロピーのトークン」に絞り込みます。これにより、バッチ内の情報密度を高め、効率的な学習を実現します。

3. 探索から洗練へのマルチステージ学習 (Exploration-to-Refinement Multi-Stage Training):

   • フェーズ 1（探索）: 学習初期段階では、過度な負の報酬をフィルタリングし、多様な解の探索を促します（SFT に近い挙動）。これによりエントロピーの早期崩壊を防ぎます。
   • フェーズ 2（洗練）: 後期段階では、エントロピーに基づくマスクを導入し、Teacher と Student の乖離が大きい重要なトークンに対して厳密なフィードバックを行い、推論経路を精緻化します。

3. 主要な貢献

1. オンポリシー蒸留の RL 的再解釈と診断:
   蒸留を方策勾配法として定式化し、停止勾配が分散低減の制御変数として機能することを理論的に示しました。これにより、不安定性の主要原因（重尾分布の負の報酬、信号の非効率性）を特定しました。
2. REOPOLD フレームワークの提案:
   上記の洞察に基づき、報酬クリッピング、動的サンプリング、マルチステージ学習を統合した新しい手法を提案しました。これにより、従来の蒸留手法が失敗する場面で安定した最適化を可能にします。
3. SOTA な効率性とスケーラビリティ:
   数学、視覚、エージェントツール使用など多様なタスクにおいて、トレーニングサンプル効率とテスト時のスケーリング性能で既存手法を凌駕することを実証しました。

4. 実験結果

• 数学推論タスク (AIME, MATH 等):
  • サンプル効率: REOPOLD は、ProRL や DeepScaleR などの最新 RL 手法と比較して、6.7 倍〜12 倍のサンプル効率を達成しました（例：1.5B モデルで 600 ステップで同等の性能）。
  • 安定性: 7B モデルを用いた実験では、従来の RKL が初期に性能低下を起こすのに対し、REOPOLD は安定して性能を向上させました。
• 視覚推論タスク (Geometry3K, MathVerse 等):
  • 教師モデルとの比較: 32B の Teacher モデルに対して、7B の Student モデルが視覚推論タスクで同等の性能を達成し、推論速度は約 3.3 倍高速化しました。
  • テスト時スケーリング: 生成サンプル数（Pass@K）を増やしても、REOPOLD は RKL ベースラインよりも優れたスケーリング曲線を示しました。
• エージェント推論 (Visual Tool-Use):
  • 画像操作ツールを使用するタスクにおいても、複雑な報酬設計を必要とせず、REOPOLD が GRPO や RKL を上回る性能とサンプル効率を示しました。

5. 意義と結論

本論文は、小規模モデルの推論能力を拡張する上で、「厳密な模倣（Strict Imitation）」から「緩和された学習（Relaxed Learning）」へのパラダイムシフトの重要性を明らかにしました。

• 理論的意義: 蒸留と RL の境界を曖昧にし、蒸留を RL の文脈で最適化技術（クリッピング、エントロピー制御）を適用可能な領域として再定義しました。
• 実用的意義: 計算リソースが限られる環境でも、大規模 Teacher モデルの推論能力を効率的に引き出すことを可能にしました。特に、テスト時のスケーリング（推論時の計算量増大による性能向上）を小規模モデルでも実現できる点は、実運用において極めて重要です。

結論として、REOPOLD は、オンポリシー蒸留の根本的な不安定性を解決し、数学、視覚、エージェントタスクにおいて最先端の性能と効率性を両立させる、小規模モデル向け推論学習の新しい基準を確立しました。