Bootstrapped Mixed Rewards for RL Post-Training: Injecting Canonical Action Order

本論文は、Zebra パズルにおける強化学習による事後学習において、単一のタスク報酬に加え、解の順序を正解の順序に近づける「順序報酬」を混合して用いることで、モデルの性能向上と正解軌道への誘導が可能であることを示しています。

要約

この論文は、**「AI に『正解』だけでなく、『考え方の順序』も教えてあげると、もっと賢くなるのか？」**という面白い実験について書かれています。

専門用語を抜きにして、わかりやすい例え話で解説しますね。

🧩 物語の舞台：「ジブラルタル（ゼブラ）パズル」

まず、実験に使われたのは「ジブラルタル（ゼブラ）パズル」という論理パズルです。
「A さんは赤い家に住んでいる」「B さんはコーヒーを飲んでいる」などのヒントから、誰がどこに住んでいて何が好きか、すべてを推測して埋め尽くすゲームです。

AI はこのパズルを解くために、マス目（セル）を一つずつ埋めていきます。

🎭 実験のセットアップ：2 つのトレーニング方法

研究者たちは、AI（ここでは「GPT-2」というタイプの頭脳）に 2 つの段階で学習させました。

1. 最初のトレーニング（ランダムな順序）：
   AI にパズルの解き方を教える際、**「正解のマス目を、でたらめな順番で」**見せました。

   • 例え話: 料理のレシピを教えるとき、「卵を割る→フライパンを熱する→卵を焼く」という正しい手順ではなく、「卵を焼く→卵を割る→フライパンを熱する」といった**「でたらめな手順」**で教えたようなものです。
   • 結果：AI は「正解」は覚えたけれど、「どうやって解いたか」という論理的な流れは混乱していました。

2. 2 回目のトレーニング（リハーサル）：
   ここで、AI に**「正解かどうか」だけでなく、「考え方の順序」**も褒めるルールを追加しました。

   • 報酬 A（正解ボーナス）： パズルが完全に解けたら「おめでとう！」と褒める（1 点）。
   • 報酬 B（順序ボーナス）： 正解かどうかは一旦置いといて、「賢い人が解くときと同じ順番でマス目を埋めたら」、少しだけ「いいね！」をもらう（0.01 点程度）。

🚀 驚きの発見：小さな「順序」のヒントが大きな力に

実験の結果、「正解」だけを褒める場合よりも、「正解＋少しだけ順序を褒める」場合の方が、AI の成績が劇的に上がりました。

• 最も効果的だった組み合わせ：
  「正解」を 99% 褒めて、「順序」をたった**1%**だけ褒める設定でした。
  • 例え話: 料理の味付けにおいて、「塩（正解）」を大量に入れるのは当然ですが、そこに**「ほんの少しのハーブ（順序）」**を加えるだけで、料理の味が格段に美味しくなるようなものです。

💡 なぜこうなるの？（重要なポイント）

ここがこの論文の一番面白いところです。

• AI は「正しい順序」を一度も見たことがありません。
  最初のトレーニングでは「でたらめな順序」しか教えていませんでした。
• でも、AI は「順序のヒント」だけで、正しい手順を勝手に見つけました。
  「正解」だけを目指すと、AI は「どうせ正解なら、適当にマス目を埋めても OK」と考えて、理屈に合わない手順で解こうとします。
  しかし、「順序」に少しだけ報酬をくっつけると、AI は**「あ、この順番で進めると先生（報酬）が喜ぶんだな」**と気づき、自然と論理的な思考プロセス（世界モデル）を身につけるようになったのです。

🛠️ この研究がすごい理由

1. データ作りが楽： 最初から「正しい順序のデータ」を大量に集めて教える必要がありません。既存の「でたらめなデータ」に、小さな「順序のヒント」を混ぜるだけで良くなります。
2. 仕組みを変えなくていい： AI の頭脳（アーキテクチャ）自体を変える必要もありません。ただ「褒め方（報酬）」を少し工夫するだけで、AI の思考力が向上します。
3. 世界モデルの理解： AI が単に答えを暗記するのではなく、「物事をどう順序立てて考えるか」という**「思考の地図（世界モデル）」**を自分で描き始めたことを示しています。

📝 まとめ

この論文は、**「AI に『正解』だけを教えるのではなく、『考え方の流れ』を少しだけ褒めてあげると、AI は自分で論理的な思考を身につけ、もっと賢くなる」**ということを証明しました。

まるで、子供に「宿題を終わらせたらご褒美」だけでなく、「勉強するときは『まず教科書を開いて、次に問題を解く』という順番で頑張ると、もっとご褒美がもらえるよ」と教えてあげたようなものです。その小さなヒントが、AI の頭の中を整理整頓させ、劇的な性能向上をもたらしたのです。

技術概要

論文要約：強化学習（RL）ポストトレーニングにおける「正解順序」の注入

1. 問題設定 (Problem)

強化学習（RL）を用いた言語モデルのポストトレーニング（微調整後のさらなる学習）では、通常、タスクの成否（スカラー値の報酬）のみを最適化し、解が生成される過程における**構造的な情報（例：中間アクションの順序）**を無視する傾向があります。
しかし、論理パズルなどのタスクにおいて、モデルが「正解への論理的な推論順序（ソルバー順序）」に従って生成を行うことは、性能向上に寄与する可能性があります。
本研究の核心的な問いは以下の通りです：

「微調整データがランダムな順序で構成されている場合でも、RL ポストトレーニングの段階でのみ『正解順序』に関するスカラーヒント（報酬）を与えることで、モデルの性能を向上させることができるか？」

2. 手法 (Methodology)

2.1 データセットとタスク

• タスク: ゼブラパズル（アインシュタインのパズル）。これは、与えられた手がかりから、エンティティと属性のグリッドを埋める論理的推論タスクです。
• データ: Shah et al. (2024) のデータセットを使用。各パズルの完全な解は 9 段階のアクション（行、列、値のトリプル）で構成されます。
• 順序のバリエーション:
  1. ソルバー順序: 決定論的なソルバーが論理的に導き出す、正しいステップバイステップの順序。
  2. ランダム順序: ソルバー順序をシャッフルしたもの。
• 実験設定: モデルはまずランダム順序のデータで教師あり微調整（SFT）されます。その後、RL ポストトレーニングを行います。

2.2 モデルアーキテクチャ

• GPT-2 風のトランスフォーマー（4 レイヤー、アテンションヘッド 4 個、隠れ層サイズ 256）。
• 事前学習済みチェックポイントではなく、ゼロから学習（Scratch）を開始。

2.3 報酬設計 (Reward Design)
GRPO（Group Relative Policy Optimization）アルゴリズムを用いて、以下の 2 つの報酬を混合して学習させます。

1. 解決報酬 (Solved Reward, $R_{solve}$):
   • モデルがパズルを完全に正しく解いた場合のみ 1、それ以外は 0 のスパースな報酬。
   • 順序は考慮せず、結果の正しさのみを評価。
2. 順序報酬 (Ordering Reward, $R_{order}$):
   • モデルが生成したセルの順序が、ソルバー順序とどの程度一致するかを評価。
   • 値の正誤に関わらず、生成順序の一致度に基づいて報酬を付与（$1 / (1 + |\pi^* - \hat{\pi}|)$ の形式）。
   • これは「報酬形状付け（Reward Shaping）」の一種であり、モデルをソルバーのような軌道へ誘導します。

2.4 ブートストラップド・スケーリング (Bootstrapped Scaling)

• 課題: $R_{solve}$ と $R_{order}$ の絶対的な値のスケールが異なる場合、単純な重み付け混合では一方が支配的になり、意図した混合比率が保てない。
• 解決策: GRPO 開始前に、微調整済みモデルを検証用データで評価し、各報酬の平均値（$\bar{R}_{solve}, \bar{R}_{order}$）を計算。
• スケーリング: 目標の混合比率 $\alpha$ に対して、以下のスケーリング因子を固定して適用する。
  $$SOLVESCALE = \frac{\alpha}{\bar{R}_{solve}}, \quad ORDERSCALE = \frac{1-\alpha}{\bar{R}_{order}}$$
  これにより、初期化時点で各成分が意図した比率で寄与するように正規化され、異なる混合比率間の比較を可能にします。

3. 主要な貢献 (Key Contributions)

1. スカラー報酬による順序注入: 微調整データを変更したり、アーキテクチャを改造したりすることなく、RL ポストトレーニング段階でのみ「ソルバー順序」へのスカラーヒントを導入する手法を提案。
2. ブートストラップド・スケーリング手法: 異質な報酬の大きさを正規化し、制御された混合実験を可能にする簡易なスケーリング手順を確立。
3. 実証的証拠: ゼブラパズルにおいて、正解性（Correctness）と粗い順序信号（Coarse ordering signals）を混合することで、RL ポストトレーニングの精度が向上することを示した。

4. 結果 (Results)

• ベースライン: ランダム順序で微調整されたモデルのテスト精度は 0.279。
• タスクのみ最適化 (1 : 0): 解決報酬のみで GRPO 学習させた場合、精度は 0.326 まで向上。
• 混合報酬の効果: 順序報酬をわずかに混ぜることで、さらに精度が向上。
  • 最適結果: 解決：順序 = 0.99 : 0.01 の混合比率で、精度 0.363 を達成。
  • 洞察: 順序報酬の重みが非常に小さい（1% 未満）場合でも、タスクのみ最適化よりも明確な改善が見られました。
  • 広範な混合: 0.75 : 0.25 や 0.9 : 0.1 のような幅広い混合比率でも、ベースライン（0.279）を上回る結果（0.355 程度）が得られました。

結論: モデルはトレーニング中にソルバー順序のシーケンスを見たことがなくても、RL 段階でのみ与えられたスカラーな順序ヒントによって、推論時にソルバーに似た軌道（Canonical trajectories）を生成するようバイアスされ、性能が向上します。

5. 意義と限界 (Significance & Limitations)

意義:

• 構造的バイアスの低コスト注入: 新しい教師データを集めたり、ゼロから学習し直したりすることなく、既存の RL ポストトレーニングパイプラインに「構造的なバイアス」を注入する安価でモジュール化された方法を提供します。
• 世界モデルの理解: 言語モデルが、明示的な順序データがなくても、報酬信号を通じて環境の動的構造（状態遷移の順序）を暗黙的に学習・利用できる可能性を示唆しています。

限界と今後の課題:

• 実験はゼブラパズルという単一のタスクと、GPT-2 風の単一モデルに限定されています。
• 報酬スケーリング因子は固定されており、学習中に報酬成分の改善速度が異なる場合、後期に較正が崩れる可能性があります。
• 今後の課題として、スケーリング因子の定期的な更新や、他のタスク・スケール・アーキテクチャへの一般化が挙げられます。

---

総括:
この論文は、RL によるポストトレーニングにおいて、単なる「正解/不正解」だけでなく、「解く過程の順序」に関する弱いシグナルを報酬に含めることで、モデルの推論能力を大幅に向上させうることを実証しました。特に、その順序情報が微調整データに含まれていなくても、RL 段階でのみ提供されるスカラー報酬として機能しうる点は、効率的なモデル改善手法として注目すべき点です。