Alternative Architecture DOJO

オルターブースのクラウドネイティブ特化型ブログです。

RAG & AI ワークフロー設計の悟りを求めて🤔

はじめに

こんにちは。寒すぎて手袋をつけてキーボードが叩けるように成りたいと考えている釘宮です。 今回は、RAG を含めた AI ワークフローを実現する際の設計で 四苦八苦 したので、その時に考えていた事と試した事についてご紹介します。

本記事はオルターブース Advent Calendar 2024の14日目の記事です。 (14日の土曜日です)

adventar.org

RAG とは

人工知能(AI)の進化に伴い、Retrieval Augmented Generation (RAG) は、情報検索と自然言語生成を融合させた革新的な手法として注目を集めています。RAGは、大規模言語モデル(LLM)と外部の知識ベースを組み合わせることで、より正確で関連性の高い応答を生成します。しかし、その実装には多くの技術的課題が伴います。

本記事では、RAGシステムの構築において、当初はKernel Memory (KM) を採用し、その後の課題を経てAzure Cosmos DB を活用する新たなアーキテクチャへの転換を模索する過程について、詳細に解説します。技術者としての視点から、具体的な実装のポイント、直面した問題、得られた洞察を共有します。

RAGシステム構築への挑戦:Kernel Memoryの採用

RAGシステムの重要性と課題

Retrieval Augmented Generation (RAG) とは、LLMと外部知識ベースを統合し、ユーザーのクエリに対して高度な応答を生成する手法です。このアプローチにより、モデルが事前学習データに含まれない最新情報や専門的な知識を提供できるため、多くの利点があります。

メリット

  • 最新情報の提供:動的な知識ベースを活用し、最新の情報を反映。
  • ドメイン特化:特定業界や企業の内部知識を組み込むことで、専門的な応答が可能。
  • 高精度な応答:関連性の高い情報を組み合わせ、より正確な回答を提供。
  • LLMの「幻覚」の軽減:事実に沿った文脈(データ)を参照させる事で、正確性を向上。

課題

  • 知識ベースの構築と管理:大量のデータを効率的に取り込み、インデックス化する必要。
  • 検索性能の最適化:高速かつ高精度な検索アルゴリズムの実装。
  • スケーラビリティ:データ量やアクセス頻度の増加に対応するアーキテクチャ設計。

Kernel Memoryの採用と期待

RAGシステムの基盤として、当初はKernel Memory (KM) を採用しました。KMは、Semantic Kernel (SK) のプラグインとしても利用可能なパッケージであり、以下の特徴を持っています。

  • 多様なデータ形式のサポート:PDF、Word、JSON、画像など、様々な形式のデータを取り込み可能。
  • セマンティック検索の実装:ベクトル検索やハイブリッド検索をサポートし、意味的な関連性を考慮した情報検索が可能。
  • 柔軟な統合オプション:Webサービス、Dockerコンテナ、ChatGPT/Copilot/Semantic Kernelのプラグインとして利用可能。これにより、既存のシステムやAIワークフローへのシームレスな統合が実現。
  • カスタム連続データハイブリッドパイプライン:大規模データセットをインデックス化し、リアルタイムでの情報更新と検索を可能にする。

期待した効果

  • 短時間での開発プロセスの加速。
  • 豊富な機能セットによる多様なニーズへの対応。
  • Azure AI Searchとの連携によるスムーズなデプロイ。

Kernel Memoryの活用と直面した課題

インデックス作成と検索性能の最適化

KMを用いた初期の実装では、データの取り込みから検索まで一連の流れを迅速に構築できました。しかし、データ量が増加するにつれ、インデックス作成や検索性能に課題が生じました。

具体的な課題

  • インデックスサイズの肥大化:大量データの取り込みにより、インデックスが大きくなり、検索速度が低下。
  • クエリパフォーマンスの低下:複雑なクエリや高頻度のアクセス時に応答速度が遅延。

Azure AI Searchのレートリミットの壁

さらに深刻だったのは、Azure AI Search のレートリミットに直面したことです。

  • レートリミットの詳細
    • Freeプラン:3 QPS(Queries Per Second)
    • Basicプラン:15 QPS
    • Standard S1プラン:25 QPS

この制限は、高トラフィックな環境や大量のデータを処理するシステムにとって、大きなボトルネックとなりました。

影響

  • スループットの制限:同時に処理できるクエリ数が制限され、ユーザー体験に影響。
  • スケーラビリティの制約:システムの拡張性が制限され、将来的なビジネス拡大に対応困難。

対応策の検討

  • サービスプランのアップグレード:より高いプランへの移行を検討しましたが、コスト増加の問題が浮上。
  • クエリの最適化:キャッシュの導入やクエリ頻度の抑制などを試みましたが、根本的な解決には至らず。

新たなソリューションの探索:Azure Cosmos DBの可能性

Azure Cosmos DBへの着目

これらの課題を受け、代替ソリューションとしてAzure Cosmos DB に着目しました。Cosmos DBは、グローバルに分散されたNoSQLデータベースで、高いスケーラビリティと低レイテンシーを提供します。

Azure Cosmos DBの特徴

  • グローバル分散と高可用性:世界中のリージョンでのデータ複製と99.999%の可用性。
  • スケーラブルなスループット:リクエストユニット(RU/s)の設定により、必要なスループットを柔軟に調整可能。
  • マルチモデルデータのサポート:ドキュメント、キー・バリュー、グラフ、列ファミリーなど様々なデータモデルに対応。

ベクトル検索機能の可能性

Azure Cosmos DBでは、ベクトル類似度検索をサポートする機能が追加されつつあります。これにより、高次元ベクトルを用いたセマンティック検索が可能となり、RAGシステムの要求を満たすことが期待されます。

ベクトル検索のメリット

  • 高速な検索性能:大規模なデータセットでも高速な類似検索が可能。
  • 柔軟なクエリ:コサイン類似度やユークリッド距離など、様々な類似度指標を使用可能。

注意点

  • 機能の成熟度:ベクトル検索機能はまだ開発中の側面があり、最新情報の収集と検証が必要。
  • 実装の複雑さ:カスタムのストアドプロシージャやUDFの作成が必要な場合がある。

Azure Cosmos DBを用いたRAGシステム再設計の試み

新たなワークフローの設計

現在、Azure Cosmos DBを活用した新しいRAGシステムのワークフローの設計と実装を進めています。以下は、その概略です。

  1. データの取り込みと前処理

    • ドキュメントをテキストデータに変換。
    • NLP処理を施し、埋め込みベクトルを生成。
  2. データの格納

    • 埋め込みベクトルと関連メタデータをCosmos DBに格納。
    • 適切なパーティションキーとインデックスを設定。
  3. クエリ処理

    • ユーザーのクエリを同様に埋め込みベクトルに変換。
    • Cosmos DB上で類似度検索を実行。
  4. 応答生成

    • 検索結果をLLMに入力し、最終的な応答を生成。

実装のポイント

  • パーティショニング戦略:データの特性に応じて最適なパーティションキーを選択し、クエリ性能を向上。
  • ベクトル類似度計算:カスタムUDFを用いて効率的な類似度計算を実装。
  • セキュリティとアクセス制御:Azure Active Directoryやロールベースのアクセス制御を活用。

初期結果

実装はまだ進行中ですが、スモールスケールでのテストでは以下のような成果が見られています。

  • パフォーマンスの向上:従来のKMベースのシステムに比べ、クエリ応答時間が改善。
  • スループットの拡大:レートリミットの制約が緩和され、同時接続数の増加に対応。

しかし、同時に新たな課題も浮上しています。

新たな課題

  • コスト管理:スループット設定に応じてコストが変動するため、最適な設定の見極めが必要。
  • 機能の制約:最新のベクトル検索機能における制限事項の把握と対応。

技術的な展望

データモデルとパーティショニング

Cosmos DBでのデータ格納において、データモデルの設計はクエリ性能に直結します。その為、以下の点に注力していく必要を感じています。

  • パーティションキーの選択:クエリパターンを分析し、データの均一な分散と効率的なアクセスを両立するキーを選択。
  • インデックス設定:必要なフィールドのみをインデックス化し、不要なオーバーヘッドを削減。

ベクトル類似度検索の実装詳細

Cosmos DBでのベクトル検索は、以下のアプローチを考察しています。

  • カスタムUDF(ユーザー定義関数)の使用

    • コサイン類似度や内積計算を行うUDFを作成。
    • クエリ内でUDFを呼び出し、類似度計算を実行。
  • クエリの最適化

    • フィルタリング条件を活用し、検索範囲を限定。
    • 必要に応じて、マテリアライズドビューの活用も検討。
  • DiskANN の活用

    • より先進的で高精度な実装を求めて検討(非常に強力な機能な為、後日に別記事を執筆します)

セキュリティとコンプライアンス

機密情報を扱う場合、以下のセキュリティ対策を実施する事が可能です。

  • データの暗号化:保存時(at rest)と転送時(in transit)の両方でデータを暗号化。
  • アクセス制御:ロールベースのアクセス制御(RBAC)を実装し、最小権限の原則を適用。
  • ネットワークセキュリティ:仮想ネットワークやファイアウォール設定を適切に構成。

結論と今後の展望

Kernel Memoryを用いたRAGシステムの構築では、多くの学びを得られましたが、ユースケースの拡大や、精度の向上を目指す過程で、スケーラビリティと性能面での課題に直面しました。また、Azure Cosmos DBの活用は、その課題を解決する可能性を秘めていると考えています。

今後の取り組み

  • 継続的な実装と検証:現行の実装を完了し、大規模なデータセットでの性能評価を実施。
  • コミュニティとの情報共有:得られた知見やベストプラクティスを技術コミュニティと共有。
  • 新技術の追求:Cosmos DBの新機能や他のソリューションも探索し、最適なアーキテクチャを追求。

おわりに

今回のような機会と学びを通して、RAGシステムの構築は、技術的な挑戦が多く変化の絶えない領域ですが、その先にある可能性は非常に大きいと感じました。技術者として、直面する課題に真摯に向き合い、最適なソリューションを模索し続けることが重要だと考えており、本記事が、同様の課題に取り組む皆様の参考となれば幸いです。

本記事では技術者としての探求と成長の過程を記述しました。課題に直面し、それを乗り越えるための試行錯誤、そして新たな可能性の発見。このプロセス自体が、技術革新の原動力であり、エンジニアリングの醍醐味であると感じています。

最後までご覧いただきありがとうございました🍩

参考リンク

:本記事の内容は、著者が現在進行中のプロジェクトとその試行錯誤を基に作成されています。技術は日々進化しており、最新の情報を常に確認することをお勧めします。


サービス一覧 www.alterbooth.com cloudpointer.tech www.alterbooth.com