jiangxinke/Agentic-RAG-R1

Agentic RAG R1 Framework via Reinforcement Learning

解決する課題

Agentic RAG‑R1 は、Retrieval‑Augmented Generation（RAG）システムにおける言語モデルの推論と検索能力を向上させることを目的としています。モデルが「いつ情報を検索すべきか」「何を取得すべきか」「取得した証拠を最終回答にどう統合すべきか」といった判断に苦慮するという制限に対処します。

仕組み

本プロジェクトは強化学習、特に GRPO（Generalized Relevance Policy Optimization）アルゴリズムを使用して、ベース言語モデルを自律的に検索・推論ループを処理できるように訓練します。システムはエージェントメモリスタックを採用し、モデルが推論、バックトラッキング、要約、検索ツール（ArtSearch 経由の Wikipedia など）の使用といったアクションを実行できるようにします。正確性、フォーマット、RAG 固有の正確性（RAGAS フレームワーク経由）を組み合わせた報酬モデルを用いて、成功した行動を強化します。

対象者

自律的な AI エージェントを構築したい開発者・研究者向けです。特に医療などの専門領域での複雑なマルチステップ検索・推論タスク（MedQA テストセットの結果参照）に関心のある方に適しています。

ハイライト

  • 強化学習統合：GRPO を使用して検索深度と回答品質を向上。
  • マルチステップ推論：熟考ループ内で推論、バックトラッキング、要約をサポート。
  • カスタムツール統合：ユーザーが独自ツールや個人 RAG データセットを組み込める。
  • リソース効率：LoRA と量子化（nf4）により、2 枚の A100 GPU だけで最大 32B パラメータのモデルをサポート。
  • 分散学習：DeepSpeed Zero 2 および Zero 3 ステージに対応。