Home 論文解説: RAG-Gym — MDP定式化とプロセス監督によるAgentic RAG最適化
投稿
キャンセル

📄 論文解説: RAG-Gym — MDP定式化とプロセス監督によるAgentic RAG最適化

本記事は RAG-Gym: Optimizing Reasoning and Search Agents with Process Supervision (arXiv:2502.13957) の解説記事です。

論文概要(Abstract)

RAG-Gymは、推論エージェントが反復的に情報を検索して複雑な質問に回答するAgentic RAGを、マルコフ決定過程(MDP)として定式化し、プロセス監督(ステップ単位のフィードバック)で訓練する統合最適化フレームワークである。さらに、検索クエリ発行前に不足情報を明示的にヒントとして生成するHintSearch戦略を提案している。著者らの実験では、プロセス監督が結果監督(最終回答のみの報酬)をBRIGHTベンチマークで8-15%上回り、HintSearchがさらに3-7%の追加改善を達成したことが報告されている。

この記事は Zenn記事: LangGraph×Claude 4.6で推論精度と応答速度を両立するマルチエージェントRAG の深掘りです。

情報源

  • arXiv ID: 2502.13957
  • URL: https://arxiv.org/abs/2502.13957
  • 著者: Guangzhi Xiong, Qiao Jin, Xiao Gu, Huaiyuan Ying, Xiaoxing Wang, Siru Ouyang et al.
  • 発表年: 2025
  • 分野: cs.AI, cs.CL, cs.IR

背景と動機(Background & Motivation)

Agentic RAG(推論エージェントが反復的に検索と推論を行うRAGシステム)は、知識集約型タスクにおいて高い性能を示している。しかし、著者らはこれらのエージェントの最適化が困難である3つの理由を指摘している。

  1. 探索空間の広さ: 可能な検索クエリと推論ステップの組み合わせが膨大
  2. スパースな訓練信号: 結果監督(最終回答の正誤のみを報酬とする)では、中間ステップの良し悪しが学習されにくい
  3. 中間ステップの評価困難: 検索結果の関連度や推論ステップの品質を定量評価する標準的手法がない

この問題はZenn記事で解説しているマルチエージェントRAGの設計課題と対応する。特に、grader_agent(文書の関連度評価)やquery_rewriter(クエリ書き換え)の最適化は、まさにRAG-Gymが取り組む問題である。

主要な貢献(Key Contributions)

  • 貢献1: Agentic RAGをマルコフ決定過程(MDP)として定式化し、検索・推論・回答生成を統一的に扱うフレームワーク
  • 貢献2: ステップ単位のプロセス監督(PRM: Process Reward Model)による訓練手法。結果監督(ORM: Outcome Reward Model)より8-15%改善
  • 貢献3: HintSearch戦略の提案。検索前に「何が不足しているか」を明示的に生成し、より的を絞った検索を実現
  • 貢献4: BRIGHT, MuSiQue, HotpotQA, FEVERにおけるSOTA達成

技術的詳細(Technical Details)

MDP定式化

著者らはAgentic RAGを以下のMDPとして定式化している。

\[\mathcal{M} = (S, A, T, R, \gamma)\]

ここで、

  • $S$(状態空間): 現在のクエリ $q$、推論コンテキスト $c$、取得済み文書集合 $D$の組 $s = (q, c, D)$
  • $A$(行動空間): 検索クエリの発行 $a_{\text{search}}(q’)$、推論ステップの生成 $a_{\text{reason}}(r)$、最終回答の出力 $a_{\text{answer}}(y)$
  • $T$(遷移関数): 検索行動は検索結果 $D’$ を返し状態を更新、推論行動はコンテキスト $c’$ を更新
  • $R$(報酬関数): プロセス報酬 $R_{\text{step}}(s, a)$(ステップ単位)+ 結果報酬 $R_{\text{outcome}}(y, y^*)$(最終回答)
  • $\gamma$: 割引率

プロセス報酬と結果報酬の比較:

\[R_{\text{process}} = \sum_{t=0}^{T} \gamma^t \cdot R_{\text{step}}(s_t, a_t) + R_{\text{outcome}}(y, y^*)\] \[R_{\text{outcome-only}} = R_{\text{outcome}}(y, y^*)\]

プロセス報酬はステップごとに報酬を与えるため、密なフィードバック信号となり、より効率的な学習が可能になる。

プロセス報酬モデル(PRM)

著者らはステップ単位の報酬を以下の3軸で定義している。

  1. 検索品質: 取得文書が現在の推論ステップに関連しているか
  2. 推論品質: 推論ステップが回答に向かって一貫して進んでいるか
  3. クエリ品質: 発行された検索クエリが的を絞ったものか

PRMはこれらの軸で各ステップにスコアを付与し、GRPO(Group Relative Policy Optimization)でエージェントを訓練する。

HintSearch戦略

HintSearchは、RAG-Gymフレームワーク内の検索戦略である。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
async def hint_search(
    query: str,
    current_context: list[dict],
    llm,
) -> list[dict]:
    """HintSearch: 情報ギャップ駆動の検索戦略

    検索クエリ発行前に「何が不足しているか」のヒントを生成し、
    より的を絞った検索を実行する

    Args:
        query: 元のユーザークエリ
        current_context: 既に取得済みの文書リスト
        llm: ヒント生成とクエリ生成に使用するLLM

    Returns:
        検索で取得した新たな文書リスト
    """
    context_summary = summarize_context(current_context)

    # Step 1: ヒント生成(不足情報の明示化)
    hint = await llm.ainvoke(
        f"以下の質問に回答するために、現在のコンテキストで不足している"
        f"情報を具体的に特定してください。\n\n"
        f"質問: {query}\n"
        f"現在のコンテキスト: {context_summary}\n\n"
        f"不足している情報(ヒント):"
    )

    # Step 2: ヒントに基づく検索クエリ生成
    search_query = await llm.ainvoke(
        f"以下のヒントに基づいて、効果的な検索クエリを生成してください。\n\n"
        f"元の質問: {query}\n"
        f"不足情報のヒント: {hint}\n\n"
        f"検索クエリ:"
    )

    # Step 3: 的を絞った検索実行
    results = await retriever.search(search_query, k=5)

    return results

HintSearchが有効な理由(著者らの分析):

  1. 明示的なギャップ特定: 検索前に「何が分かっていないか」を言語化することで、推論プロセスが構造化される
  2. 的確なクエリ生成: ギャップに基づくクエリは、元のクエリをそのまま使うよりも関連度の高い文書を取得する
  3. 推論の方向付け: ヒント生成が推論の「次のステップ」を明確にし、探索の効率が向上する

Zenn記事のinterleaved thinkingとの対応

HintSearchの「不足情報の明示化→的を絞った検索」というパターンは、Zenn記事のinterleaved thinking(tool call間にClaude自身の推論を挟む)と概念的に対応する。

HintSearchinterleaved thinking
ヒント生成(不足情報特定)thinking block(検索結果の評価)
ヒントベースのクエリ生成次のtool call(修正クエリ)
的を絞った検索実行search_documents tool call

Claude Sonnet 4.6のadaptive thinkingを有効にしてツールを定義すると、Claudeは各tool call間に思考ブロックを自動生成し、HintSearchと同様の「ギャップ特定→的確な検索」のサイクルを暗黙的に実行する。

実装のポイント(Implementation)

PRMの訓練データ収集

プロセス報酬モデルの訓練には、ラベル付き推論ステップデータが必要である。著者らは以下のアプローチを採用している。

  1. 初期エージェントで大量の推論トレースを生成
  2. 最終回答の正誤と各ステップの関連度を自動ラベリング
  3. ラベル付きデータでPRMを訓練
  4. PRMを使ってエージェントをGRPOで訓練

このパイプラインはコストが高いが、結果監督と比較して8-15%の改善を著者らは報告している。

human-in-the-loopとの統合

著者らは、PRMの自動監督とhuman-in-the-loop(HITL)が補完的であると述べている。

  • Phase 1: 人間フィードバックでPRMを訓練(高品質だが少量のデータ)
  • Phase 2: 訓練されたPRMで自動プロセス監督(大量データ、低コスト)

この段階的アプローチは、Zenn記事のinterrupt()によるhuman-in-the-loop設計と整合する。幻覚チェックで不合格の場合にinterrupt()で人間判断を要求するパターンは、RAG-Gymの「自動評価→人間フォールバック」のアーキテクチャと同じ構造を持つ。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# Zenn記事のHITL + RAG-Gymのプロセス監督を組み合わせた例
from langgraph.types import interrupt

async def production_rag_with_process_supervision(query: str) -> str:
    """プロセス監督 + HITL の統合パイプライン"""
    docs = await retrieve_with_retry(query)
    relevant_docs = await grade_documents(query, docs)

    # プロセス監督: 検索品質のスコアリング
    retrieval_quality = evaluate_retrieval_quality(query, relevant_docs)

    if retrieval_quality < 0.5:
        # 品質が低い場合: HintSearchスタイルの再検索
        hint = await generate_hint(query, relevant_docs)
        additional_docs = await retrieve_with_hint(hint)
        relevant_docs.extend(additional_docs)

    result = await generate_with_hallucination_check(query, relevant_docs)

    if not result["is_grounded"]:
        # 幻覚検出: human-in-the-loopで判断
        human_decision = interrupt({
            "query": query,
            "answer": result["answer"],
            "reason": "幻覚チェック不合格",
        })
        return human_decision.get("revised_answer", result["answer"])

    return result["answer"]

Production Deployment Guide

AWS実装パターン(コスト最適化重視)

RAG-GymのHintSearch戦略をプロダクションに適用する場合、ヒント生成のLLM呼び出しが追加されるため、コスト設計が重要になる。

トラフィック量別の推奨構成:

規模月間リクエスト推奨構成月額コスト主要サービス
Small~3,000 (100/日)Serverless$80-200Lambda + Bedrock + DynamoDB
Medium~30,000 (1,000/日)Hybrid$400-1,000Lambda + ECS + ElastiCache
Large300,000+ (10,000/日)Container$2,500-6,000EKS + Karpenter + Spot

HintSearchの追加LLM呼び出し(ヒント生成 + クエリ生成)により、通常のRAGと比較して1リクエストあたり20-40%のBedrock APIコスト増加が見込まれる。ただし、的を絞った検索により再検索の回数が減少するため、total costでは同等以下になる場合もある。

コスト試算の注意事項: 上記は2026年2月時点のAWS ap-northeast-1料金に基づく概算値です。最新料金は AWS料金計算ツール で確認してください。

Terraformインフラコード

Small構成: Lambda + Bedrock + DynamoDB

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# --- Lambda: HintSearch Agent ---
resource "aws_lambda_function" "hint_search" {
  filename      = "hint_search.zip"
  function_name = "rag-gym-hint-search"
  role          = aws_iam_role.lambda_bedrock.arn
  handler       = "index.handler"
  runtime       = "python3.12"
  timeout       = 120  # ヒント生成+検索で長め
  memory_size   = 1024

  environment {
    variables = {
      BEDROCK_HINT_MODEL   = "anthropic.claude-3-5-haiku-20241022-v1:0"
      BEDROCK_ANSWER_MODEL = "anthropic.claude-sonnet-4-6-20260514-v1:0"
      DYNAMODB_TABLE       = aws_dynamodb_table.hint_cache.name
    }
  }
}

# --- DynamoDB: ヒント+検索結果キャッシュ ---
resource "aws_dynamodb_table" "hint_cache" {
  name         = "rag-gym-hint-cache"
  billing_mode = "PAY_PER_REQUEST"
  hash_key     = "query_hash"

  attribute {
    name = "query_hash"
    type = "S"
  }

  ttl {
    attribute_name = "expire_at"
    enabled        = true
  }
}

運用・監視設定

CloudWatch: HintSearch効果の監視

1
2
3
4
fields @timestamp, query, hint_generated, retrieval_quality_before, retrieval_quality_after
| stats avg(retrieval_quality_after - retrieval_quality_before) as avg_improvement
  by bin(1h)
| filter hint_generated = true

コスト最適化チェックリスト

  • ヒント生成はHaiku($0.25/MTok)で実行、回答生成はSonnet
  • ヒントキャッシュ: 類似クエリのヒントをDynamoDB TTL 24hで再利用
  • HintSearchはretrieval_quality < 0.5の場合のみ発動(条件付き)
  • Bedrock Batch API: 非リアルタイムのヒント生成に50%割引適用
  • Prompt Caching: ヒント生成用システムプロンプト固定で30-90%削減
  • AWS Budgets: 月額予算設定
  • CloudWatch: ヒント生成回数の監視
  • HintSearch発動率のアラーム(品質劣化の早期検知)

実験結果(Results)

著者らは4種のベンチマークで評価を行っている。

主要結果:

ベンチマーク結果監督 (ORM)プロセス監督 (PRM)PRM + HintSearch改善率 (ORM→PRM+Hint)
BRIGHTベースライン+8-15%+11-22%11-22%
MuSiQueベースライン改善さらに改善-
HotpotQAベースライン改善さらに改善-
FEVERベースライン改善さらに改善-

(出典: 論文Section Experimental Results)

著者らの報告による主要な知見:

  1. プロセス監督 > 結果監督: BRIGHTベンチマークで8-15%の改善。ステップ単位の密なフィードバックが中間推論ステップの品質を向上
  2. HintSearchの追加効果: プロセス監督に加えてさらに3-7%の改善。情報ギャップの明示化が検索の的確さを向上
  3. サンプル効率: プロセス監督は結果監督より少ない訓練データで収束
  4. GPT-4o / Claude / Llamaとの比較: ベースラインとして使用。RAG-Gymで訓練されたエージェントがpromptingのみのベースラインを上回る

制約: PRMの訓練データ収集にコストがかかる。HintSearchは1リクエストあたりのLLM呼び出し回数が増加する。評価は英語タスクに限定されている。

実運用への応用(Practical Applications)

RAG-Gymの知見は、Zenn記事のマルチエージェントRAGに以下の形で適用できる。

  • HintSearchのpromptingベース適用: fine-tuningなしでも、interleaved thinkingで「不足情報の特定→的を絞った検索」のパターンをプロンプトで誘導可能
  • HITL設計の指針: 自動品質評価(PRM的機能)で大半をフィルタし、低品質ケースのみ人間にエスカレーションする2段階設計
  • grader_agentの改善: 文書の関連度評価をプロセス監督の観点で設計し、「回答に向かって進んでいるか」を評価軸に追加

関連研究(Related Work)

  • MARA (Singh et al., 2025, arXiv:2504.04603): LangGraphベースのマルチエージェントRAG。RAG-GymはMARAのような既存アーキテクチャを最適化するフレームワークと位置付けられる
  • Adaptive-RAG (Jeong et al., 2024, arXiv:2404.19756): クエリ複雑度による戦略選択。RAG-Gymは選択された戦略内のステップ最適化に焦点を当てる
  • Self-RAG (Asai et al., 2023, arXiv:2310.11511): 反省トークンによる自己批評。RAG-Gymのプロセス監督は外部PRMを使用する点で異なる

まとめと今後の展望

RAG-Gymは、Agentic RAGをMDPとして定式化し、プロセス監督で訓練する統合フレームワークを提供している。著者らの実験では、プロセス監督が結果監督を8-15%上回り、HintSearchがさらに3-7%の改善をもたらすことが報告されている。

Zenn記事のinterleaved thinking(tool call間の段階的推論)は、HintSearchの「情報ギャップ特定→的確な検索」パターンと同様の効果をプロンプティングレベルで実現するアプローチと解釈できる。production RAGの設計においては、自動品質評価(プロセス監督)とhuman-in-the-loop(interrupt)の組み合わせが、著者らの提案する最適化フレームワークと整合した設計パターンとなる。

参考文献


:::message 本記事はAI(Claude Code)により自動生成されました。内容は論文に基づいていますが、実際の利用時は原論文もご確認ください。 :::

この投稿は CC BY 4.0 でライセンスされています。