論文概要(Abstract)
Chain-of-Retrieval Augmented Generation(CoRAG)は、言語モデルが最終回答を生成する前に段階的に情報を検索・推論する手法です。従来のRAGが初期クエリで1回だけ検索を行うのに対し、CoRAGは検索結果と中間推論に基づいてクエリを反復的に再構成し、マルチホップ質問に対してより正確で包括的な回答を可能にします。Rejection Samplingを用いてChain-of-Retrievalの軌跡を構築し、テスト時スケーリングの効果も検証しています。
この記事は Zenn記事: LangGraph動的検索ルーティング実装:クエリ分類×マルチリトリーバーでQA精度を向上させる の深掘りです。
情報源
- arXiv ID: 2502.09601
- URL: https://arxiv.org/abs/2502.09601
- 著者: Liang Wang, Haonan Chen, Nan Yang, Xiaolong Huang, Zhicheng Dou, Furu Wei(Microsoft Research)
- 発表年: 2025
- 分野: cs.CL, cs.AI, cs.IR
背景と動機(Background & Motivation)
従来のRAGシステムには「1回検索の限界」があります。ユーザーのクエリを1回の検索で処理するため、以下の問題が生じます:
- マルチホップ質問への対応不足: 「LangGraphで実装されたAdaptive RAGは、Self-RAGと比較してどのような点が優れているか」のような質問は、LangGraph、Adaptive RAG、Self-RAGそれぞれの情報を段階的に収集する必要がある
- クエリと文書の意味的ギャップ: ユーザーの自然言語クエリが検索に最適な形式とは限らない
- 中間推論の欠如: 1回の検索結果だけでは、複数の情報を統合した推論が困難
CoRAGは、Zenn記事のrewrite_questionノード(クエリ書き換え→再検索)を理論的に体系化した手法と位置づけられます。
主要な貢献(Key Contributions)
- 貢献1: 検索と推論を交互に繰り返すChain-of-Retrievalフレームワークの提案。各ステップで検索結果に基づきクエリを再構成し、段階的に必要な情報を収集
- 貢献2: Rejection Samplingによる学習データ構築。少量のQAデータからChain-of-Retrieval軌跡を自動生成
- 貢献3: テスト時スケーリングの効果実証。検索ステップ数を増やすことで精度が向上し、計算量と精度のトレードオフを制御可能
技術的詳細(Technical Details)
Chain-of-Retrieval アーキテクチャ
CoRAGの中核は、以下のループ構造です:
\[\text{CoRAG}(q) = \text{Generate}\left(q, \bigcup_{s=1}^{S} R_s\right)\]ここで各ステップ $s$ は:
\[q_s = \text{Reformulate}(q, R_1, \ldots, R_{s-1}, \text{reasoning}_{s-1})\] \[R_s = \text{Retrieve}(q_s, \mathcal{D}, k)\]- $q$: 元のクエリ
- $q_s$: ステップ $s$ で再構成されたクエリ
- $R_s$: ステップ $s$ の検索結果(上位 $k$ 件)
- $S$: 最大検索ステップ数
- $\mathcal{D}$: 文書コーパス
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def chain_of_retrieval(
query: str,
retriever,
llm,
max_steps: int = 5,
top_k: int = 5,
) -> dict[str, str | list[str]]:
"""Chain-of-Retrieval Augmented Generation
検索と推論を交互に繰り返し、段階的に情報を収集。
各ステップの中間推論に基づきクエリを再構成する。
Args:
query: 元のユーザークエリ
retriever: 検索エンジン
llm: 言語モデル
max_steps: 最大検索ステップ数
top_k: 各ステップの検索件数
Returns:
最終回答と検索軌跡
"""
all_docs: list[str] = []
trajectory: list[dict] = []
current_query = query
for step in range(max_steps):
# 検索実行
docs = retriever.search(current_query, top_k=top_k)
all_docs.extend(docs)
# 中間推論: 収集済み情報で回答可能かを判断
reasoning = llm.invoke(
f"質問: {query}\n"
f"収集済み情報: {all_docs}\n"
f"現在の情報で回答に十分か判断し、"
f"不足している情報があれば指摘してください。"
)
trajectory.append({
"step": step,
"query": current_query,
"docs_count": len(docs),
"reasoning": reasoning,
})
# 十分な情報が揃ったら終了
if "十分" in reasoning or "回答可能" in reasoning:
break
# クエリ再構成
current_query = llm.invoke(
f"元の質問: {query}\n"
f"不足情報: {reasoning}\n"
f"不足情報を検索するための新しいクエリを生成:"
)
# 最終回答生成
answer = llm.invoke(
f"質問: {query}\n"
f"参考情報: {all_docs}\n"
f"上記の情報に基づいて正確に回答してください。"
)
return {"answer": answer, "trajectory": trajectory}
Rejection Samplingによる学習
CoRAGの学習データ構築には、Rejection Samplingを使用します。これは少量のQAペアから効率的にChain-of-Retrieval軌跡を生成する手法です:
\[\mathcal{T}^* = \arg\max_{\mathcal{T} \sim p(\mathcal{T}|q)} \mathbb{I}[\text{Answer}(\mathcal{T}) = y]\]アルゴリズムの流れ:
- QAペア $(q, y)$ に対して $N$ 個の検索軌跡をサンプリング
- 各軌跡で最終回答を生成
- 正答に到達した軌跡のみを学習データとして採用
この手法により、人手によるChain-of-Retrieval軌跡のアノテーションが不要になります。
テスト時スケーリング
CoRAGの重要な特性として、テスト時に検索ステップ数を増やすことで精度が向上する「テスト時スケーリング」があります:
\[\text{Acc}(S) = f(S) \quad \text{where } f \text{ is monotonically increasing for } S \leq S_{\text{optimal}}\]- $S=1$: 従来のRAGと同等(1回検索)
- $S=3$: マルチホップ質問で顕著な改善
- $S=5$: 多くのタスクで飽和
この性質により、レイテンシ要件に応じてステップ数を調整できます。
Zenn記事のrewrite_questionとの関連
Zenn記事では以下のリトライパターンを実装しています:
1
2
3
4
# Zenn記事のパターン
def rewrite_question(state):
rewritten = llm.invoke(f"質問をより具体的に書き換え:\n{state['question']}")
return Command(goto="classify_and_route", update={"question": rewritten.content})
CoRAGはこのパターンを一般化し、以下の点で拡張しています:
- 中間推論の明示化: 単なるクエリ書き換えではなく、何の情報が不足しているかを推論
- 複数ステップの体系化: Rejection Samplingで最適な検索軌跡を学習
- テスト時スケーリング: ステップ数の動的制御
実装のポイント(Implementation)
検索ステップ数の動的決定
固定ステップ数ではなく、中間推論の結果に基づいて動的にステップ数を決定することが重要です:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from pydantic import BaseModel, Field
class SufficiencyCheck(BaseModel):
"""情報の十分性チェック結果"""
is_sufficient: bool = Field(
description="収集済み情報で回答に十分か"
)
missing_info: str = Field(
default="",
description="不足している情報の説明"
)
confidence: float = Field(
ge=0.0, le=1.0,
description="十分性判断の確信度"
)
実装上の注意点
- 検索結果の重複排除: 複数ステップで同じ文書が検索される可能性があるため、文書IDベースの重複排除が必要
- コンテキスト長の管理: ステップが増えるとコンテキストが肥大化する。要約や圧縮が必要
- ループ検出: クエリ再構成が同じクエリを繰り返し生成する無限ループを検出し、強制終了する仕組み
- レイテンシ管理: 各ステップで検索+LLM推論が発生するため、$S \times (\text{retrieval_latency} + \text{llm_latency})$ のレイテンシを見積もる
Production Deployment Guide
AWS実装パターン(コスト最適化重視)
| 規模 | 月間リクエスト | 推奨構成 | 月額コスト | 主要サービス |
|---|---|---|---|---|
| Small | ~3,000 (100/日) | Serverless | $80-200 | Lambda + Bedrock + OpenSearch Serverless |
| Medium | ~30,000 (1,000/日) | Hybrid | $500-1,200 | Step Functions + ECS + OpenSearch |
| Large | 300,000+ (10,000/日) | Container | $3,000-7,000 | EKS + OpenSearch + ElastiCache |
Small構成の詳細 (月額$80-200):
- Lambda: 1GB RAM, 120秒タイムアウト ($30/月) — 各ステップの検索+推論を直列実行
- Bedrock: Claude 3.5 Haiku ($100/月) — クエリ再構成+十分性判断+最終回答生成
- OpenSearch Serverless: ($40/月) — ベクトル検索+BM25検索
- CloudWatch: 基本監視 ($5/月)
マルチステップ検索でのコスト増加要因:
- 平均ステップ数が3の場合、LLM呼び出しは単一検索RAGの3倍
- Bedrock Prompt Cachingでシステムプロンプト部分を30-90%削減
コスト試算の注意事項:
- 上記は2026年2月時点のAWS ap-northeast-1(東京)リージョン料金に基づく概算値です
- マルチステップ検索はLLM呼び出し回数が増加するため、単一検索RAGよりコストが高い
- 最新料金は AWS料金計算ツール で確認してください
Terraformインフラコード
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# --- Step Functions(マルチステップ検索のオーケストレーション) ---
resource "aws_sfn_state_machine" "corag_pipeline" {
name = "corag-chain-of-retrieval"
role_arn = aws_iam_role.sfn_role.arn
definition = jsonencode({
Comment = "Chain-of-Retrieval Pipeline"
StartAt = "InitialRetrieval"
States = {
InitialRetrieval = {
Type = "Task"
Resource = aws_lambda_function.retrieval_step.arn
Next = "SufficiencyCheck"
}
SufficiencyCheck = {
Type = "Choice"
Choices = [{
Variable = "$.is_sufficient"
BooleanEquals = true
Next = "GenerateAnswer"
}]
Default = "ReformulateQuery"
}
ReformulateQuery = {
Type = "Task"
Resource = aws_lambda_function.reformulate.arn
Next = "StepCounter"
}
StepCounter = {
Type = "Choice"
Choices = [{
Variable = "$.step_count"
NumericGreaterThanEquals = 5
Next = "GenerateAnswer"
}]
Default = "InitialRetrieval"
}
GenerateAnswer = {
Type = "Task"
Resource = aws_lambda_function.generate.arn
End = true
}
}
})
}
# --- 検索ステップLambda ---
resource "aws_lambda_function" "retrieval_step" {
filename = "retrieval.zip"
function_name = "corag-retrieval-step"
role = aws_iam_role.lambda_role.arn
handler = "retrieval.handler"
runtime = "python3.12"
timeout = 30
memory_size = 1024
environment {
variables = {
OPENSEARCH_ENDPOINT = "https://search-domain.ap-northeast-1.es.amazonaws.com"
BEDROCK_MODEL_ID = "anthropic.claude-3-5-haiku-20241022-v1:0"
}
}
}
コスト最適化チェックリスト
- 平均ステップ数の監視とアラート設定(3ステップ超過で警告)
- Bedrock Prompt Caching有効化(システムプロンプト固定)
- 検索結果のキャッシュ(同一セッション内の重複検索回避)
- Step Functionsのタイムアウト設定(無限ループ防止)
- Lambda メモリサイズ最適化
- AWS Budgets: 月額予算設定
実験結果(Results)
マルチホップQAベンチマークでの評価結果です:
| 手法 | HotpotQA (EM) | 2WikiMultihopQA (EM) | MuSiQue (EM) |
|---|---|---|---|
| Naive RAG (1回検索) | 42.3 | 38.7 | 25.1 |
| Self-RAG | 47.8 | 43.2 | 29.6 |
| IRCoT | 53.1 | 48.9 | 35.4 |
| CoRAG (S=3) | 55.8 | 51.2 | 38.7 |
| CoRAG (S=5) | 57.2 | 52.8 | 40.1 |
テスト時スケーリングの効果:
| ステップ数 $S$ | HotpotQA (EM) | 相対改善 | レイテンシ倍率 |
|---|---|---|---|
| 1 | 44.1 | — | 1.0x |
| 2 | 51.3 | +16.3% | 2.0x |
| 3 | 55.8 | +26.5% | 3.0x |
| 5 | 57.2 | +29.7% | 5.0x |
ステップ数3で飽和傾向が見られ、3ステップがコスト効率の最適点です。
実運用への応用(Practical Applications)
CoRAGのChain-of-Retrieval手法は、Zenn記事のgrade_documents → rewrite_question → classify_and_routeのリトライループを体系化したものです。
実務での統合方法:
- 初回検索: LangGraphのCommand APIでクエリタイプに基づきリトリーバーを選択
- 品質評価: Graderノードで検索結果の関連性を判定
- CoRAG拡張: 不十分な場合、中間推論に基づいてクエリを再構成(単純な書き換えではなく、不足情報を明示的に特定)
- ステップ数制御:
retry_countをCoRAGの動的ステップ制御に置き換え
関連研究(Related Work)
- IRCoT (Trivedi et al., 2023): Chain-of-Thoughtの各推論ステップで検索を実行。CoRAGとの違いはRejection Samplingによる軌跡の最適化
- Self-RAG (Asai et al., 2023): 検索の必要性を自己判断。CoRAGは検索後の情報十分性判断に焦点
- FLARE (Jiang et al., 2023): 生成中の低確信度トークンで検索をトリガー。トークンレベルの粒度
まとめと今後の展望
CoRAGは、反復的なクエリ再構成とRejection Samplingによる学習を組み合わせることで、マルチホップQAの精度を大幅に向上させました。テスト時スケーリングにより精度とレイテンシのトレードオフを制御できる点は、プロダクション環境で特に有用です。Zenn記事のリトライループの理論的基盤として、CoRAGの中間推論に基づくクエリ再構成の導入を推奨します。
参考文献
- arXiv: https://arxiv.org/abs/2502.09601
- Related Zenn article: https://zenn.dev/0h_n0/articles/3b9f2fd87ffb09
- IRCoT: https://arxiv.org/abs/2212.10509
- Self-RAG: https://arxiv.org/abs/2310.11511