🔥

rStar2-Agent：エージェント型強化学習で数学推論を革新する14Bモデル~中野哲平

に公開
強化学習
生成 AI
中野哲平
zennfes2025free
生成aiと強化学習
tech

はじめに：「より長く考える」から「より賢く考える」へのパラダイムシフト

AI分野において数学推論は長年の難題でした。従来のアプローチは、Chain-of-Thought（CoT）による「より長い思考プロセス」で精度向上を図ってきましたが、推論過程で微細なエラーが発生すると、それが累積して最終的に間違った答えに導いてしまう根本的な限界がありました。

このような背景の中、Microsoftの研究チームが発表したrStar2-Agentは、単に「長く考える」のではなく「賢く考える」という革新的なアプローチで、この問題に挑戦しています。わずか14億パラメータのモデルでありながら、671億パラメータの巨大モデルDeepSeek-R1と同等の性能を達成し、しかもより簡潔な回答で同じ結果を実現しています。

本記事では、この画期的な研究について、その技術的背景から実装の詳細まで、丁寧に解説していきます。

背景知識：従来手法の限界とエージェント型強化学習の必要性

Chain-of-Thoughtの限界

従来の数学推論では、長いCoT（Chain-of-Thought）による段階的思考プロセスが主流でした。しかし、この手法には以下の問題がありました：

  1. エラーの累積問題：推論チェーンの途中でエラーが発生すると、後続の推論もすべて影響を受ける
  2. 自己修正の困難性：内部的な自己省察だけでは、根本的に間違ったアプローチを修正できない
  3. 検証機能の欠如：生成された推論が正しいかどうかを客観的に検証する仕組みがない

エージェント型強化学習とは

エージェント型強化学習（Agentic Reinforcement Learning）は、AI モデルが外部環境（この場合はPythonコード実行環境）と積極的に相互作用しながら学習する手法です。従来の静的なテキスト生成とは異なり、以下の特徴があります：

  • 動的な問題解決：コードを書いて実行し、結果を分析して次のアクションを決定
  • 客観的検証：実際のコード実行結果による仮説の検証
  • 反復的改善：実行結果を受けて推論プロセスを修正・改善

rStar2-Agentの革新的アーキテクチャ

核心となる3つのイノベーション

rStar2-Agentは、大規模エージェント型強化学習を可能にする3つの主要な技術革新により実現されています：

1. 高効率RL実行基盤

技術的特徴：

  • 同時に45,000件のツール呼び出しを処理可能
  • 平均レイテンシ0.3秒（スケジューリング+実行時間）
  • 64個のMI300X GPUという限られたリソースで1週間での訓練完了

アーキテクチャの工夫：

┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│ 動的ロールアウト │    │ 負荷分散システム │    │ 分離型コード実行 │
│ スケジューラ    │────│                  │────│ 環境            │
└─────────────────┘    └──────────────────┘    └─────────────────┘

従来のシステムでは、不均等なワークロード分散によりGPUのアイドル時間が発生していましたが、rStar2-AgentはリアルタイムGPUキャッシュ可用性に基づく動的ワークロード分散により、この問題を解決しています。

2. GRPO-RoC：環境ノイズに対応する新アルゴリズム

**Group Relative Policy Optimization with Resampling on Correct（GRPO-RoC）**は、従来のGRPOアルゴリズムを拡張し、コード実行環境特有のノイズに対処します。

従来の問題：

  • 最終的な答えが正解であれば報酬が与えられるため、途中の誤ったツール利用が修正されない
  • フォーマット違反や実行エラーによる学習の不安定性

GRPO-RoCの解決策：

  1. 生成軌跡の分割：

    • 正解集合：最終答えが正しい軌跡
    • 不正解集合：最終答えが間違っている軌跡

  2. 品質による選別：

    • 正解集合内で、ツール呼び出しの正確性とフォーマット違反の少なさで評価
    • 高品質な半分のみを残し、残りはダウンサンプリング

  3. 非対称サンプリング：

    • ポジティブ軌跡：高品質なもののみ強化
    • ネガティブ軌跡：多様性保持のため一律ダウンサンプリング

数学的定式化：

L = E[r(s,a) - baseline] - β * KL(π_θ || π_ref)

ここで、βは削除され（β=0）、低確率探索の促進と学習の安定化を図っています。

3. 効率的な段階的学習レシピ

3段階のRL訓練プロセス：

第1段階：基礎推論能力の獲得

  • 応答長：8K トークン制限
  • 目的：基本的なツール使用と短い推論の習得

第2段階：推論能力の拡張

  • 応答長：12K トークン制限
  • 目的：より複雑な問題への対応能力向上

第3段階：難問特化訓練

  • 応答長：制限なし
  • 目的：最高難度問題での最適化

事前準備：

  • 非推論SFT：指示追従、JSONフォーマット、ツール使用のみ
  • 推論能力は意図的に除外：強化学習で自然に獲得させるため

データ品質管理とキュレーション

学習データの厳格な品質管理

学習データは、Qwen3-32Bで2回正解した問題のみを使用という厳しい基準を設けています。

データセット構成：

  • OpenMathReasoning：基礎的な数学問題
  • Project Euler：937問の高難度プログラミング数学問題
  • 総計42K問：厳選された高品質な問題-解答ペア

品質管理プロセス：

  1. 検証不可能な回答の除去：「極限は存在しない」などの曖昧な回答
  2. 複雑すぎるフォーマットの除去：過度に複雑な数式表現
  3. 実行タイムアウトの原因となる大きな数値の除去

データキュレーションの革新

整数回答フィルタリング：

# Qwen3-32Bで16回生成し、整数回答が2回以上一致する問題のみ採用
for problem in dataset:
    responses = qwen3_32b.generate(problem, n=16)
    integer_answers = [r for r in responses if is_integer(r)]
    if count_matches(integer_answers) >= 2:
        dataset_filtered.append(problem)

GRPO-RoCアルゴリズムの詳細解析

基本的なGRPOの理論

Group Relative Policy Optimization（GRPO）は、価値関数を必要とせず、グループベースの優位性推定により、LLMの推論能力を向上させる強化学習アルゴリズムです。

従来のPPOとの違い：

PPO: 価値関数ネットワークが必要 → メモリ使用量大
GRPO: 同一プロンプトの複数回答の平均報酬をベースライン → メモリ効率的

グループベース優位性計算：

Advantage = Reward - Group_Average_Reward

GRPO-RoCの拡張機能

Resample-on-Correct（RoC）戦略：

  1. オーバーサンプリング：標準バッチサイズより多くのロールアウトを生成
  2. 品質評価：ツールエラーとフォーマット違反で評価
  3. 選択的ダウンサンプリング：
    • 正解軌跡：最高品質のもののみ保持
    • 不正解軌跡：多様性保持のため均等にダウンサンプリング

アルゴリズムの疑似コード：

def grpo_roc_step(prompts, model, reward_function):
    # 1. オーバーサンプリング
    all_responses = []
    for prompt in prompts:
        responses = model.generate(prompt, n=oversample_ratio)
        all_responses.extend(responses)
    
    # 2. 報酬計算と分類
    correct_responses = []
    incorrect_responses = []
    for response in all_responses:
        reward = reward_function(response)
        if reward > threshold:
            correct_responses.append((response, reward, tool_error_count(response)))
        else:
            incorrect_responses.append((response, reward))
    
    # 3. 品質ベース選別
    correct_responses.sort(key=lambda x: x[2])  # ツールエラー数でソート
    high_quality_correct = correct_responses[:len(correct_responses)//2]
    
    # 4. 均等ダウンサンプリング
    sampled_incorrect = random.sample(incorrect_responses, target_size)
    
    return high_quality_correct + sampled_incorrect

実験結果と性能分析

数学推論ベンチマークでの成果

rStar2-Agent-14Bは、AIME24で80.6%、AIME25で69.8%の正答率を達成し、671BパラメータのDeepSeek-R1を上回る性能を実現しました。

詳細な性能比較：

モデル パラメータ数 AIME24 AIME25 平均応答長
rStar2-Agent-14B 14B 80.6% 69.8% ~10,000 tokens
DeepSeek-R1 671B 79.8% - ~17,000 tokens
o3-mini (medium) - 79.6% - -
Claude Opus 4.0 - 77.0% - -

効率性の証明：

  • 50%短い応答で同等以上の性能
  • 47倍少ないパラメータで最先端性能

一般化能力の実証

数学以外の分野での性能：

数学のみの強化学習にも関わらず、科学推論（GPQA-Diamond）や一般タスク（IFEval、Arena-Hard）でも優秀な性能を発揮：

  • GPQA-Diamond：DeepSeek-V3を上回る
  • BFCL v3：エージェント型ツール使用で優秀
  • IFEval：一般的なアライメントタスクで競争力ある結果

技術的発見と考察

リフレクション・トークンの発見

従来の推論トークン：

  • フォーキング・トークン：自己省察と探索のトリガー

新発見のリフレクション・トークン：
ツールからのフィードバックに対する応答として現れる新しいカテゴリーのトークン

  • 環境駆動型推論：コード実行結果の慎重な分析
  • エラー診断：実行失敗の原因分析
  • 適応的アプローチ：結果に基づく戦略調整

計算効率の革命

従来の課題：

  • モデル学習用GPU
  • 推論用GPU
  • ツール実行用CPU
    これらの統合による複雑度の増大

rStar2-Agentの解決：

  • 統合インフラによる効率的リソース利用
  • 動的負荷分散によるボトルネック解消
  • 最適化されたスケジューリングによる高スループット達成

実装上の技術的課題と解決策

分散コード実行環境の構築

セキュリティ考慮事項：

# 必須のセキュリティ対策
- Dockerによる完全な環境分離
- 外部ネットワークへの接続遮断
- 実行時間制限（MAX_EXECUTION_TIME）
- ワーカープロセスの監視

スケーラビリティ管理：

# 動的ワーカー調整
MAX_WORKERS = adjust_based_on_cpu_cores()
TIMEOUT_SETTINGS = {
    'simple_computation': 5,
    'complex_computation': 30,
    'symbolic_math': 60
}

KL発散項の削除と探索の促進

理論的背景：
最近の研究（Open-Reasoner-Zero、Understanding R1-Zero-Like Training等）により、GRPOにおけるKL発散項は必須ではないことが示されている

実装上の利点：

# β = 0.0 による効果
- 低確率な有効解の探索が可能
- 学習の安定性向上
- 計算コストの削減

エントロピー損失項の調整

目的：

  • 探索と活用のバランス最適化
  • 過度に保守的な生成の防止
  • 創造的な解法の発見促進

限界と今後の課題

現在の制約事項

  1. データ品質依存性：

    • Qwen3-32Bの能力を超える問題は扱えない
    • より高次の数学概念への拡張には限界

  2. 計算資源要件：

    • GPU、CPU、ストレージの統合管理が必要
    • 大規模展開時のインフラコスト

  3. ドメイン特化性：

    • 現在は数学に特化
    • 他分野への一般化には追加研究が必要

今後の発展方向

技術的改善：

  • より効率的なサンプリング戦略
  • 適応的グループサイズの決定
  • マルチドメインへの拡張

応用分野の拡大：

  • 科学研究における仮説検証
  • エンジニアリング問題の自動解決
  • 教育分野での個別指導システム

実践的な示唆と影響

AI開発への示唆

  1. 効率性重視の設計：巨大なモデルよりも効率的な学習手法
  2. エージェント型アプローチ：静的生成から動的問題解決へ
  3. 品質重視のデータキュレーション：量より質の重要性

産業への影響

教育分野：

  • 個別化された数学指導システム
  • リアルタイムの誤解検出と修正

研究開発：

  • 自動化された仮説検証
  • 科学計算の効率化

金融・エンジニアリング：

  • 複雑な数値計算の自動化
  • リスク分析の精度向上

結論：AI推論の新たな地平

rStar2-Agentは、AI の数学推論において「より長く考える」から「より賢く考える」への根本的なパラダイムシフトを実現しました。14億パラメータという中規模モデルでありながら、洗練された訓練手法により最先端の推論能力を達成し、より持続可能なAI能力向上の道筋を示していることは、今後のAI研究開発に大きな影響を与えるでしょう。

主要な貢献：

  1. 技術的革新：GRPO-RoCアルゴリズムと効率的インフラ
  2. 実証的成果：最先端性能の達成と一般化能力の確認
  3. 方向性の提示：効率性と性能を両立する新しいアプローチ

今後の展望：
この研究は、AIシステムが単なるテキスト生成を超えて、動的で相互作用的な問題解決能力を持つ真のエージェントへと進化する道筋を明確に示しています。静的な知識処理から、ツールを活用した能動的な探索・検証・学習へのシフトは、AGI（汎用人工知能）への重要な一歩と言えるでしょう。

rStar2-Agentの成功は、AI研究者だけでなく、実用的なAIアプリケーションを開発する企業や教育機関にとっても、新たな可能性を切り開く画期的な成果として記憶されることになるでしょう。

この研究の詳細な実装コードとトレーニングレシピは、GitHubで公開されており、研究コミュニティのさらなる発展に貢献することが期待されています。

Discussion