Zenn
☸️

Amazon Bedrock Knowledge Bases が GraphRAG をサポート（プレビュー）！だが…それがなんだというのだ…

2024/12/14に公開
AWS
neptune
Bedrock
Amazon Bedrock
GraphRAG
tech

大規模言語モデル (LLM) を企業のナレッジベースと組み合わせて活用する際、一般的なアプローチとしてRAG (Retrieval-Augmented Generation) が注目されています。RAG は、質問に関連する文書をベクトル検索などで取得し、その情報を LLM に提供することで、より正確な回答の生成を可能にします。

しかし、従来のベクトル検索ベースの RAG には、いくつかの課題があることが指摘されています。特に以下の点が重要な課題として挙げられます。

  • 関係性の分断: 文書をチャンク (断片) に分割する際に、エンティティ間の関係性が失われてしまい、「XX が YY にどのように影響したか」といった因果関係を説明する質問に対して十分な回答が困難になる。
  • 冗長な情報の取得: 十分な情報をコンテキストに含めるために大量のチャンクを取得すると、その中から重要な情報を抽出するのが困難になることがある (Lost in the middle ジレンマ)。
  • グローバルな文脈の欠如: 局所的な情報の取得に優れる一方で、文書全体を通じた広い文脈や関連性の把握が困難。

このような課題に対して、グラフベースのアプローチである GraphRAG が新たな解決策として注目されています。GraphRAG は、文書から抽出した情報をエンティティとその関係性として構造化されたグラフとして表現します。これにより、テキスト間の相互関係を保持しながら情報を圧縮し、より正確で文脈を考慮した回答の生成を実現します。


(出典) Graph Retrieval-Augmented Generation: A Survey

2024年12月に実施された AWS (Amazon Web Services) の最大の年次イベントである re:Invent 2024 において、高度な RAG 構築のマネージドサービスである Amazon Bedrock Knowledge Bases の GraphRAG サポート (プレビュー) が発表されました。こちらの機能では AWS が提供しているグラフデータベースである Amazon Neptune をバックエンドとして GraphRAG の仕組みを構築することができます。

本記事では、Amazon Bedrock Knowledge Bases と Amazon Neptune を活用した GraphRAG の実装方法について解説し、従来の RAG が抱える課題をどのように解決できるのかを具体的に示していきます。

グラフデータベースと GraphRAG の基礎

グラフデータベースとは

データベースの世界では基本的にテーブルという二次元の構造でデータを管理し、主キーや外部キーを用いてテーブル間の関係を表現します。しかし、現実世界のデータは必ずしもテーブル構造に適していません。例えば、「誰が誰を知っているか」「どの商品がどの商品と関連しているか」「あるアイデアが別のアイデアにどのように影響したか」といった複雑な関係性を、リレーショナルデータベースで表現しようとすると、多数の結合テーブルが必要になり、クエリも複雑化してしまいます。

グラフデータベースは、このような関係性を直感的に表現・管理するために設計されたデータベースです。その基本構造は以下の通りです。

  • ノード (頂点): エンティティを表現する。例えば、人物、会社、場所、商品など。各ノードには、名前や属性といったプロパティを持たせることができる。また、ノードは「Company」「City」「Employee」のようなラベルを持つことができ、これによってノードの種類を分類することができる。
  • リレーション（エッジ）: ノード間の関係性を表現し、エッジとも呼ばれる。例えば「HAS_CEO」「LOCATED_IN」「INFLUENCED」といった関係性を表す。リレーションは方向性を持ち、さらにプロパティ (開始日や終了日など) を付与することができる。

このような構造により、例えば「Company ラベルのノードが City ラベルのノードに対して LOCATED_IN リレーションを持つ」といった形で、現実世界の関係性を直感的にモデル化することができます。

このような構造を持つグラフデータベースは、後述する GraphRAG において、文書から抽出した知識を構造化し、関係性を保持したまま検索可能にするための重要な基盤技術となっています。Amazon Neptune は、このようなグラフデータベースをマネージドサービスとして提供し、スケーラブルで信頼性の高いグラフデータの管理を実現します。

GraphRAG の仕組みと利点

GraphRAG は、文書から抽出した知識をグラフ構造として表現し、それを活用して高精度な質問応答を実現する手法です。その処理は大きく分けてインデックス作成フェーズと検索・生成フェーズの2つで構成されています。


出典: LlamaIndex Webinar: Advanced RAG with Knowledge Graphs (with Tomaz from Neo4j)

インデックス作成フェーズでは、まず文書からエンティティとその関係性を抽出し、トリプレット (SVO: 主語-述語-目的語) の形式に変換します。このトリプレットはグラフデータベースに格納され、主語と目的語がノードとして、述語がリレーションとして表現されます。また、それぞれの要素に関する追加情報はプロパティとして保存されます。

検索・生成フェーズでは、ユーザーからの質問を基にグラフのクエリを生成し、グラフデータベースから関連する部分グラフを取得します。この構造化された情報が LLM に提供され、LLM はこれを基に回答を生成します。

GraphRAG の最大の特徴は、エンティティ間の関係性を明示的に保持できる点です。従来の RAG では失われがちだった「どのように影響したか」といった因果関係を、グラフ構造として保存することで、より正確な回答が可能になります。例えば、「印象派の技法が後世の芸術運動にどう影響したか」という質問に対して、影響の連鎖を適切に追跡できるようになります。

文書の内容をグラフ構造として圧縮することで、情報の冗長性も効果的に排除できます。同じエンティティに関する情報は自然に統合され、「Lost in the middle」ジレンマの問題が緩和されます。また、情報が構造化された形式で保存されることで、LLM がより確実に情報を理解できるようになります。

これらの機能は、複雑な技術文書の理解と質問応答、製品間の関連性を考慮したカスタマーサポート、研究論文の引用関係や影響の分析、法令や規制の関連性の把握など、幅広いユースケースで活用できます。特に、Amazon Bedrock Knowledge Bases と Amazon Neptune の統合により、これらの利点を簡単に活用できる環境が整いました。

Amazon Bedrock Knowledge Bases による GraphRAG 構築

  1. 現状、ドキュメントからのグラフの構築には Anthropic Claude 3 Haiku が用いられます。まずは Amazon Bedrock のダッシュボードからモデルアクセスの設定を変更し、Claude 3 Haiku のアクセスを有効化してください。

  1. 次に、Amazon Bedrock Knowledge Bases のダッシュボードの「Create (作成)」メニューを開き、「Knowledge Base with vector store」をクリックしてください。

  1. 「Step 1 Provide Knowledge Base detals」は変更せず、次に進みます。

  2. 「Step 2 Configure data source」で、検索対象のドキュメント (今回はホワイトペーパー) を配置した S3 URI を入力します。今回はその他の設定は変更せず、次に進みます。

  1. 「Step 3 Select embeddings model and configure vector store」で、ベクトル埋め込みに用いるモデルを選択します。今回は "Titan Embeddings G1 - Text v1.2" を選択します。

  1. ベクトルデータベースの設定では「Quick create a new vector store - Recommended」を選択し、ベクトルストアは「Amazon Neptune Analytics (GraphRAG) - Preview」を選択します。

  1. 10分ほど待機すると Neptune Analytics の作成が完了します。

  2. 作成されたナレッジベースの詳細画面に遷移したら、データソース一覧にあるデータソースにチェックを入れ、「Sync (同期)」ボタンをクリックします。

これでドキュメントのチャンク化、ベクトル化、グラフ構造の作成などの処理が実行されます。裏側で Titan Text Embeddings や Claude 3 Haiku が甲斐甲斐しく働いている間、ドキュメントのサイズによって数分から数時間待つ必要があります。

質問応答の実例

作成したナレッジベースの詳細画面の右側の「Test Knowledge Base」ペインで GraphRAG の質問応答をテストすることができます。ここでは回答生成モデルとして Anthropic Claude 3.5 Sonnet を選択し、「Amazon SageMakerのトレーニングワークフローにおいて、データ品質の問題がモデルパフォーマンスに与える影響と、それを緩和するためのベストプラクティスについて教えてください」と質問してみました。

回答は以下のようになりました。

Amazon SageMakerのトレーニングワークフローにおいて、データ品質の問題はモデルパフォーマンスに大きな影響を与える可能性があります。トレーニング時には通常、クリーンなデータを使用しますが、本番環境では異なるソースからデータが送られてくるため、モデルの性能が低下する可能性があります。

この問題を緩和するためのベストプラクティスとして、以下の2つのアプローチが推奨されます：

  1. 統計の比較：ソース入力データとMLモデルが実際にトレーニングに使用したデータを比較します。これには定量的評価（データ数、データ期間、入力データの精度の確認）と定性的評価が含まれます。
  2. クリーンでないデータ入力に対するモデルの検証：本番環境に近い状態でモデルの性能を評価します。[1][2]

さらに、データ品質の問題に対処するためのベストプラクティスには以下が含まれます：

  1. 小規模で統計的に有効なサンプルデータセットから始め、異なるデータ準備戦略を反復的に試す。
  2. データクリーニングプロセス中に、異常を警告するフィードバックループを実装する。
  3. データの整合性を継続的に強制する。
  4. 管理されたETL（抽出、変換、ロード）サービスを活用する。
  5. Amazon SageMaker Ground Truthを使用して、高品質のトレーニングデータセットを迅速かつ正確に構築する。
  6. データの可視化と分析を行い、パターンを識別する。

これらの実践により、データ品質の問題がモデルパフォーマンスに与える影響を最小限に抑え、より信頼性の高いMLモデルを開発することができます。[3][4]

いかがでしょうか。クエリとしては、(1) 複数のドキュメントにまたがる情報を統合する必要がある、(2) 文脈を理解する必要がある、(3) 因果関係の理解が必要、といった観点で考えてみましたが、GraphRAG によってよい回答が得られたのではないかと思います。

Neptune Analytics グラフの可視化

Amazon SageMaker AI Notebook インスタンスの作成

Neptune Analytics に構築されたグラフに対してクエリを叩いたり、可視化をするには Amazon SageMaker AI Notebook インスタンスを利用するのが簡単です。「Creating a new Neptune Analytics notebook using a AWS CloudFormation template」のドキュメントにそれぞれのリージョンにおける AWS CloudFormation スタックが用意されており、これを使うと Neptune Analytics の操作をするための graph-notebook などが設定済みのノートブックインスタンスが起動します。

CloufFormation スタックのパラメータでいくつか Neptune Analytics に接続するためのネットワークの設定を入力します。事前に Neptune Analytics のパブリック接続を有効化しましょう。

  1. まずは Neptune Analytics のグラフのダッシュボードを表示し、Amazon Bedrock Knowledge Bases によって作成されたグラフの ID をクリックします。

  1. 画面右上にある「Modify (修正)」ボタンをクリックします。

  1. Network > Enable public connectivity の下の「Allows your graph to be reachable over the Internet. All access to the graph requires IAM authentication. (インターネット経由のグラフへのアクセスを許可します。すべてのグラフへのアクセスは IAM 認証が必要です。)」にチェックを入れ、「Next (次へ)」をクリックします。

  1. 確認画面で「Submit (送信)」ボタンをクリックすれば変更が反映されます。

  2. 「Creating a new Neptune Analytics notebook using a AWS CloudFormation template」のドキュメントにある「Launch Stack (スタックの起動)」ボタンをクリックします。私は東京リージョンで作業しているので、「Asia Pacific (Tokyo)」の行にあるボタンをクリックします。

  1. 「Step 1 Create stack」は変更せず、「Step 2 Specify stack details」に進みます。
  2. (2) の手順でメモしたグラフエンドポイントを入力します。また、任意のノートブック名 (graphrag-notebook) を入力し、「Next (次へ)」ボタンをクリックします。

  1. 「Step 3 Configure stack options」で「I acknowledge that AWS CloudFormation might create IAM resources (AWS CloudFormation が IAM リソースを作成し得ることを承認します。)」にチェックを入れ、「Next (次へ)」ボタンをクリックします。

  1. 「Step 4 Review and create」の内容が問題なければ「Submit (送信)」ボタンをクリックします。CloudFormation スタックの作成が始まり、数分待つと完了します。
  2. 「Outputs (出力)」タブにある SageMaker Notebook インスタンスの URL をクリックし、SageMaker Notebook (マネージド Jupyter ノートブック) を開きます。

グラフのクエリと可視化

右側にある「New」→「Python 3」をクリックして、新しいノートブックを開きます。

今回 CloudFormation で作成した SageMaker AI Notebook は、Neptune Analytics の操作をするための graph-notebook が設定済みであり、便利なマジックコマンドを使うことができます。

まず、セルの中に %summary と入力し、Shift + Enter などで実行しましょう。

すると、JSON 形式でグラフの要約情報が取得できます。

  • ノード数は 8,066
  • エッジ (リレーション) の数は 20,905
  • ノードのラベルは "Entity"、"DocumentId"、"Chunk" の 3 種類
  • エッジのラベルは "FROM"、"CONTAINS" の 2 種類

グラフ DB に具体的にどのような値が入っているかを取得するにはクエリを実行します。Neptune Analytics は openCypher 形式のクエリ言語をサポートしており、マジックコマンド %%oc をセルの頭に書くことで実行できます。

ここでは、MATCH 文で「ノード (node) から任意のリレーション (relation) で繋がっている隣接ノード (peer)」を取得するクエリを実行します。全件取ってくると時間がかかりますし、可視化するとごちゃごちゃするので LIMIT 100 として件数を絞っています。

%%oc
MATCH (n)-[r]-(p)
RETURN n, r, p
LIMIT 100

実行結果の「Graph」セルをクリックすると、グラフを可視化することができます。要約情報にあったように、 "Entity" (黄色)、"DocumentId" (赤色)、"Chunk" (青色) のノードがあり、チャンクがエンティティを含む (CONTAINS)、もしくはチャンクがドキュメント ID に由来する (FROM) といったリレーションが貼られていることがわかります。

各ノードやリレーションの詳細を表示するには右上のメニューアイコンをクリックします。例えば以下の図のエンティティは "amazon s3" という ID (x-amz-bedrock-kb- はヘッダー) を持っています。

チャンクノードの詳細を表示すると、AMAZON_BEDROCK_TEXT というプロパティの中にテキストチャンクが含まれていることがわかります。

なお、チャンクの埋め込みベクトルは Neptune Analytics の neptune.algo.vectors.get 関数を使って取得することができます。

%%oc
MATCH (n:Chunk)
call neptune.algo.vectors.get(n)
yield node, embedding
RETURN n.AMAZON_BEDROCK_TEXT, embedding
LIMIT 1

このように埋め込みベクトルも格納されていることで、グラフクエリだけでなく、ベクトル検索も組み合わせたハイブリッド検索が実現されています。

まとめと今後の展望

本記事では、Amazon Bedrock Knowledge Bases の GraphRAG サポート (プレビュー) について、その基本的な概念から実装方法まで解説してきました。GraphRAG は、従来の RAG が抱えていた関係性の分断や冗長な情報取得、グローバルな文脈の欠如といった課題に対して、グラフ構造を活用することで効果的な解決策を提供します。Amazon Bedrock Knowledge Bases のマネージドな仕組みにより、これまで複雑だった GraphRAG の実装が大幅に簡素化され、より多くの開発者がこの技術を活用できるようになりました。

今後の展望として、以下のような点に期待が寄せられます。

  1. カスタマイズ可能なグラフスキーマ
    • 現状では、ノードとリレーションのラベルが "Entity"、"DocumentId"、"Chunk" および "FROM"、"CONTAINS" に固定されています。今後、ユースケースに応じたカスタムスキーマの定義が可能になるとより汎用性が高まるかと思います。
  2. 高度なグラフ分析機能の統合
    • グラフコミュニティ検出など、先進的な GraphRAG 実装で用いられている高度な分析機能は現時点では利用されていません。これらの機能が追加されることで、より広い文脈理解や関連性の分析が可能になると考えられます。

なお、GitHub 上で 1,900 以上のスターを獲得しているポピュラーな GraphRAG 実装である gusye1234/nano-graphrag において、Amazon Bedrock のサポートを追加しました。Amazon Bedrock で利用可能な Anthropic Claude などのモデルを用いて高度な GraphRAG を構築したい場合はこちらのリポジトリもチェックしてみてください。

Discussion