DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

この論文は、強化学習におけるサンプル効率の低さと多様性の欠如という課題を解決するため、動的な参照分布を用いて時間的に関連するサンプルを保持し、分布制約を通じて多様性を維持する新しい正則化フレームワーク「DyJR」を提案し、数学的推論や Text-to-SQL などのベンチマークで GRPO や既存手法を上回る性能と訓練効率の両立を実証しています。

要約

🧠 問題：AI は「勉強しすぎ」で頭が固くなってしまう

まず、背景にある問題から説明します。

最近の AI（大規模言語モデル）は、数学の問題や複雑な論理パズルを解くために、**「試行錯誤（ロールアウト）」**を繰り返して学習しています。

• 従来の方法（GRPO など）： AI が正解を見つけると、「よし、この解き方を覚えよう！」と、その瞬間のデータを使って学習します。しかし、一度学習が終わると、そのデータは捨ててしまいます。
• 過去の失敗： 効率を上げるために「過去の正解データ」を再利用しようとした他の研究では、AI が**「特定の解き方しか考えられなくなる（多様性の欠如）」**という問題が起きました。
  • 例え話： 料理のレシピ本を勉強している生徒が、一度「完璧なパスタの作り方」を教えてもらったら、「パスタはこれしかない！」と固執し、他のどんな料理（他の解き方）も考えられなくなってしまうような状態です。これを「モード崩壊（Mode Collapse）」と呼びます。

💡 解決策：DyJR（ダイナミック・ジェンセン・シャノン・リプレイ）

この論文の著者たちは、「過去のデータは『正解を覚えるため』ではなく、『多様な考え方を保つため』に使おう」と考えました。そこで提案されたのがDyJRという仕組みです。

この仕組みには、2 つの大きな工夫があります。

1. 🗓️ 「時限付きの思い出のアルバム」を作ろう（動的なバッファ）

AI が学習する初期段階では、まだ試行錯誤が多く、いろいろな「面白い解き方」が見つかります。しかし、学習が進むと、AI はすぐに「一番確実な解き方」に固執し始めます。

• 従来の方法： 過去のデータをすべて保存して、何年も前のデータまで使おうとする（メモリを圧迫し、古い情報が邪魔になる）。
• DyJR の方法： 「最近の思い出」だけを残すようにします。
  • 例え話： 勉強中の生徒に、**「1 週間前のノートだけ」**を常に手元に置いておき、それより前のノートは破棄するルールを作ります。
  • さらに、**「勉強の最初の 20 分間は、ノートを 4 倍も集めておこう！」**というルールもあります。なぜなら、勉強の最初はアイデアが溢れているから。でも、落ち着いてきたら、ノートの量は減らして、最新の「新鮮なアイデア」だけを残すのです。
  • これにより、AI は「古い固定観念」に縛られず、常に「今一番近い未来の多様な考え方」を参考にできます。

2. 🧭 「正解」ではなく「バラエティ」を基準にしよう（JS 発散の正則化）

ここが最も重要な部分です。過去のデータを使って AI を更新する際、どうアプローチするか？

• 従来の方法： 「過去の正解データに、AI の答えを近づけろ！」と強制的に修正する。
  • 結果： AI は「あの正解と同じこと」しか言わなくなり、創造性が失われます。
• DyJR の方法： 「過去の正解データの**『集まり』**から、AI が大きく逸脱しないように優しく誘導する」ことにします。
  • 例え話： 料理のコンテストで、過去の優勝レシピ（多様な正解の集まり）があります。
    • 従来の方法は、「優勝レシピと一模一样（同じ味）にしろ！」と命令します。
    • DyJR は、「優勝レシピの**『雰囲気』や『バリエーション』から大きく外れないようにしなさい。でも、新しい味も探していいよ」という「道しるべ（コンパス）」**として使います。
  • これにより、AI は「正解」を追求しつつも、「他の面白い解き方」も忘れずに探求し続けることができます。

🏆 結果：どうなった？

この新しい方法（DyJR）を試したところ、以下の素晴らしい結果が出ました。

1. 数学の問題や SQL（データベース検索）の作成などで、他の方法より高い正解率を達成。
2. AI が「一つの解き方に固執する」のを防ぎ、多様なアプローチを維持できた。
3. 計算コストはほとんど増えず、従来の方法と同じくらい速く学習できた。

🌟 まとめ

この論文が伝えたかったことは、**「AI に勉強させる時、過去の『正解』をただコピーさせるのではなく、過去の『多様な成功体験』を『地図』として使って、AI が迷子にならないように優しく導いてあげることが重要だ」**ということです。

DyJR は、AI が「賢く」なるだけでなく、「柔軟で創造的」であり続けるための、とても賢い学習のルールブックなのです。

技術概要

DyJR: 検証可能報酬を伴う強化学習における多様性の維持に関する技術的サマリー

本論文「DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay」は、大規模言語モデル（LLM）の推論能力向上における強化学習（RL）の効率性と多様性の維持に関する新たなアプローチを提案しています。以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

1. 背景と問題定義

背景

推論能力は、現在の LLM における汎用人工知能（AGI）の核心です。OpenAI o1 のようなモデルは、検証可能な報酬（Verifiable Rewards）を用いた強化学習（RLVR）を通じて、長い Chain-of-Thought（CoT）探索により論理能力を飛躍的に向上させています。

既存手法の課題

現在の主流であるオンポリシーアルゴリズム（例：GRPO）は、サンプル効率が悪く、過去のロールアウト（生成履歴）を一度の更新後に破棄してしまいます。これを解決するため、経験再生（Experience Replay）を用いた手法（ReMix, RLEP, Ex-GRPO など）が提案されました。これらは過去の正確なサンプルを再利用して方策を更新します。

しかし、既存の経験再生手法には以下の2 つの根本的な誤解と課題があります。

1. モード崩壊（Mode Collapse）の悪化: 過去の軌跡の尤度を直接最大化する更新を行うと、モデルは特定の解経路に過剰適合（Overfitting）し、探索能力が急速に失われます。
2. 計算コストとメモリ負荷: 従来の手法は膨大な履歴データを保存・再利用する必要があり、GPU メモリへの過大な負荷がかかります。また、RL 訓練の初期段階ではモデルのエントロピーが急激に低下するため、すべての履歴を均等に保存することは非効率的です。

2. 提案手法：DyJR (Dynamic Jensen-Shannon Replay)

著者らは、経験再生の目的を「正確性の強化」から「多様性の維持（正則化）」へと転換することを提案し、DyJRを開発しました。

主要な革新点

(1) 時間感応型ダイナミックバッファ (Time-Sensitive Dynamic Buffer)

モデルの進化に合わせて、保存するデータの量と範囲を動的に調整します。

• FIFO と適応的サイズ調整: バッファは先入れ先出し（FIFO）方式で更新され、現在のモデルから時間的に最も近いサンプルのみを保持します。
• 動的容量メカニズム: 訓練の初期段階（急激な遷移フェーズ）では、高エントロピーの推論パターンを捕捉するためにバッファ容量を拡大し、モデルが安定するにつれて縮小させます。
• バイアス対応適応データ選択: 難易度に応じたデータ選択戦略を採用し、簡単なタスクでは高信頼度のサンプルを、難しいタスクでは稀な解を優先的に取り込みます。これにより、初期段階での多様性の崩壊を防ぎます。

(2) ジェンセン・シャノン発散正則化 (Jensen-Shannon Divergence Regularization)

過去の履歴を直接の方策勾配更新に使うのではなく、分布制約として利用します。

• ダイナミックな分布アンカー: バッファに保存された過去の成功軌跡の混合分布を「動的な分布アンカー」として定義します。
• JS 発散の最小化: 現在の方策とこの混合分布との間の Jensen-Shannon (JS) 発散を最小化することで、モデルが多様な成功経路から乖離するのを防ぎます。
• 利点: 従来の KL 発散（特に Forward KL）はモードカバリング（平均化）傾向があり、分布のシフトに対して不安定になりがちですが、JS 発散は対称性があり有界であるため、よりロバストな正則化信号を提供します。

最適化目的関数

DyJR の総損失関数は、オンポリシーの GRPO 損失と JS 正則化項の和として定義されます。
$$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$$
ここで、$\alpha_{JS}$ は正則化係数です。

3. 主要な貢献

1. 再生パラダイムの再定義: 経験再生の価値を「正確性の強化」から「多様性の維持」へと転換し、直接勾配更新ではなく分布制約（JS 発散）を導入しました。
2. 時間的近接性に基づく動的データ構築戦略: 訓練の初期フェーズでのみ大規模な保存を行い、その後は最小限のフットプリントで運用する非一様なバッファ戦略を提案しました。
3. 広範な実験と詳細な分析: 数学的推論と Text-to-SQL タスクにおいて、既存の手法（GRPO, RLEP, Ex-GRPO, DPH-RL など）を上回る性能を示し、Rank-k トークンの確率進化の観点から訓練ダイナミクスを解明しました。

4. 実験結果

評価設定

• タスク: 数学的推論（Reinforce-Ada-Hard, AIME25, HMMT25 等）および Text-to-SQL（BIRD, Spider）。
• モデル: Qwen3-4B-Base, Llama-3.1-8B-Instruct。

主要な結果

• 数学的推論: DyJR は、ベースラインの GRPO（平均 29.8%）に対し、**34.1%の平均精度を達成し、+4.3%**の大幅な改善を示しました。特に AMC23 では +7.4%、HMMT25 では +2.9% の向上が見られました。
• Text-to-SQL: BIRD データセットで Pass@1 が +3.3%、Spider で +5.0% 向上し、Pass@16 でも同様に優れた性能を示しました。
• 効率性: 従来の RLEP が大量のメモリを必要とするのに対し、DyJR は必要なメモリを劇的に削減（理論上 2k ペア程度）しつつ、GRPO と同等のトレーニング効率を維持しました。
• JS 発散の有効性: Forward KL 発散を用いた変種と比較しても、JS 発散を用いた DyJR の方が性能が上回りました（34.1% vs 32.5%）。これは、JS 発散が非定常な混合分布に対してより安定した正則化を提供するためです。

多様性の分析

• エントロピーと Rank-k 確率: GRPO は訓練初期にエントロピーが急激に低下し、Rank-1 トークンの確率が 90% 以上で固定される「モード崩壊」を起こしました。一方、DyJR は Rank-1 の支配を緩和し、Rank-2 や Rank-3 への確率を再分配することで、多様性を維持しつつ性能を向上させました。
• スケーラビリティ: 大規模なサンプリング（Pass@k, k=1024）においても、DyJR は性能が継続的に向上し、他の手法を凌駕しました。

5. 意義と結論

DyJR は、RLVR における経験再生の役割を根本的に再定義しました。歴史的データは、単に「正解」を反復させるためではなく、訓練初期段階に内在する「豊かな探索パターン」を維持し、モデルが狭い解経路に収束するのを防ぐために利用すべきであると示しました。

このアプローチにより、従来の経験再生手法が抱えていた計算コストの増大やモード崩壊の問題を解決しつつ、推論タスクにおける LLM の性能を大幅に向上させることに成功しました。特に、JS 発散を正則化項として用いることで、モデルの探索能力と安定性のバランスを最適化できる点が、本研究の最大の技術的貢献と言えます。