RAGを超えて:AWS上のエンタープライズAI向けタスク認識知識圧縮
従来のRAGは、数百のドキュメントにまたがる分析タスクにおいて限界に達します。この記事では、AWS上でタスク認識知識圧縮(TAKC)を使用して、知識ベース全体をタスク固有の表現に事前圧縮し、複数の忠実度階層でキャッシュし、各クエリを適切な階層にルーティングする方法を、デプロイ可能なオープンソース実装とともに紹介します。
検索拡張生成(RAG)を、財務デューデリジェンスや規制コンプライアンスレビューなど、数百のドキュメントにまたがる複雑な分析タスクに使用している場合、おそらくその限界に直面したことがあるでしょう。類似性検索は関連する断片を見つけ出しますが、ドキュメント間のつながりを見落とすことがよくあります。この記事では、タスク認識知識圧縮(TAKC)と呼ばれる手法を使用してそのギャップを埋める方法を紹介します。TAKCは、知識ベース全体をタスク固有の表現に事前圧縮し、AWS上にデプロイします。完全なオープンソース実装を自分のアカウントにデプロイできます。
タスク認識知識圧縮
製造会社を5億ドルで買収しようとしているプライベートエクイティファームを考えてみましょう。デューデリジェンスチームは、12の子会社と5年間の財務諸表を分析する必要があります。さらに、200以上のサプライヤー契約、8つの施設からの環境コンプライアンス報告書、50以上の訴訟案件にも直面しています。アナリストが「現在のサプライヤー条件と係争中の訴訟を考慮した総合的な財務リスクは?」と質問した場合、RAGの類似性検索では回答を導き出せません。数百のドキュメントに関連情報が含まれていますが、それらの間のつながりには語彙的な類似性がありません。
TAKCは、LLMを使用してドキュメントのより短くタスクに焦点を当てた要約を生成することで、この種の問題に対処します。異なるタスクには異なる要約が生成されます。同じドキュメントでも、財務分析用に圧縮されたものは収益、利益率、キャッシュフローのデータを必要としますが、コンプライアンスレビュー用に圧縮されたものは規制の引用や違反履歴を必要とします。汎用的な要約はすべてをカバーしようとするため、特定のユースケースに対する情報密度が薄まります。TAKCは、特定のタスクのレンズを通してドキュメントを圧縮し、重要な情報を保持して残りを破棄します。インジェストパイプラインのセクションでは、圧縮プロンプトが保持する情報を具体的に指定する方法を示しています。本番環境では、タスクタイプのプロンプトをバージョン管理された設定(AWS Systems Manager Parameter Storeや専用のAmazon Simple Storage Service(Amazon S3)プレフィックスなど)に保存し、プロンプトの変更を監査可能にし、プロンプトが更新されたときに再圧縮をトリガーできるようにします。
システムは、タスクタイプごとにドキュメントをオフラインで1回圧縮します。クエリ時には、元のドキュメントではなく、事前圧縮された表現を取得します。システムは圧縮バージョンを使用して質問に答えます。圧縮表現に十分な詳細がない場合、クエリ複雑性アナライザーは質問をより多くのコンテキストを保持する低い圧縮階層にルーティングします。
TAKCは、類似性検索が返すトップkチャンクだけでなく、圧縮形式で知識ベース全体へのアクセスを提供します。圧縮時にドキュメントが一緒に処理されるため、システムはドキュメント間のつながりを保持します。また、同じソースマテリアルから異なるタスクに対して異なる圧縮出力を生成します。10-Kファイルの財務圧縮は、同じファイルの法的リスク圧縮とはまったく異なります。圧縮により、タスク関連情報の保持を対象としながら、トークン数を8倍から64倍に削減します。
マルチレート圧縮
異なるクエリには異なるレベルの忠実度が必要です。「第3四半期の収益は?」という質問には、サプライヤーの支払条件と子会社間の四半期キャッシュフローの関係を分析するリクエストよりもはるかに少ないコンテキストで十分です。
TAKCは、タスクタイプごとに4つの圧縮階層を維持することでこれに対処します。最も軽い圧縮(8倍)では、システムは約87.5%のコンテキストを保持します。多段階の推論とクロスドキュメントの統合に十分な詳細を保持します。中程度の階層(16倍)では、コンテキスト削減率は約93.8%です。このレベルは、中程度の複雑さの分析クエリに適しています。高階層(32倍)ではコンテキストが約96.9%削減され、事実確認や明確に定義された質問に適しています。超階層(64倍)では約98.4%削減され、分類タスクやキーワード検索に適しています。
クエリ複雑性アナライザーは、クエリ長、質問タイプ、分析言語の存在などのシグナルに基づいて、受信クエリを適切な階層にルーティングします。単純な事実に関する質問は超圧縮キャッシュにヒットし、複雑な分析質問は軽度圧縮キャッシュを使用します。これはユーザーには透過的に行われます。
ほとんどのエンタープライズクエリは、より高い圧縮階層から最小限のコストで処理できるルックアップです。まれな複雑なクエリは、必要な場合にのみ大きなコンテキストバジェットを消費します。この階層ベースのルーティングは、メタデータフィルタリングやクエリ再フォーマットなどの既存のRAG最適化を補完し、圧縮を適用する前にドキュメントセットを絞り込むことができます。圧縮品質を検証するには、各階層のLLM応答を、完全な非圧縮ドキュメントから生成された応答と比較します。参考実装には、特定のタスクタイプとドキュメントに対してこの比較を実行するテストスクリプトが含まれています。
AWS上のアーキテクチャ
この実装は、AWS上で2つの分離されたサーバーレスパイプラインとして実行されます。1つはインジェスト用、もう1つはクエリ用です。図1に両方のパイプラインを示します。
TAKCアーキテクチャ:パイプラインフローはデータの取り込みと圧縮を処理します。ユーザーフローは認証されたクエリを処理します。
AWS Lambdaをコンピューティングに選択した理由は、各関数呼び出しが短命でイベント駆動型であるためです。ワークロードはインジェスト中にバーストでデータを処理し、バースト間に可変のクエリ負荷を処理するため、サーバーレスは自然な選択です。
Amazon API Gatewayを選択して、クエリインターフェースをRESTエンドポイントとして公開しました。キャッシュには、シャード管理なしで複合キー(takc:{task}:{rate})での読み取りにAmazon ElastiCache Serverlessを選択しました。Amazon Cognitoは、カスタム認証コードなしでJWT発行とトークンリフレッシュを処理し、実装面積を削減します。
インジェストパイプライン
ドキュメントがタスクタイプのプレフィックス(例:raw-data/financial/)でAmazon S3に配置されると、S3イベント通知がAWS Lambda関数をトリガーします。この関数は、ドキュメントを256トークンのチャンクに分割し、境界での情報損失を防ぐために50トークンのオーバーラップを持たせます。次に、各チャンクに対して非同期的に圧縮Lambda関数を呼び出し、並列処理を可能にします。大規模なインジェストの場合は、圧縮関数で予約された同時実行数を構成し、チャンク化ステップと圧縮ステップの間にAmazon Simple Queue Service(Amazon SQS)キューを配置して、スロットリングを適切に処理します。
2番目の関数は、Amazon Bedrockを呼び出して、4つの圧縮階層すべてでチャンクを圧縮します。各圧縮呼び出しには、モデルに保持する情報を指示するタスク認識プロンプトが含まれます:
タスク:財務分析。収益指標、利益率、キャッシュフロー、債務、財務リスク指標を保持する。 圧縮目標:元の長さの約1/16に削減。 指示:
- タスクに関連する事実と関係に焦点を当てる
- 数値データと指標を保持する
- エンティティとその属性を維持する
- 因果関係と依存関係を保持する
- 冗長または無関係な情報を削除する
モデルは、プロンプトがタスクにとって重要なものを指定するため、何を保持すべきかを認識します。この違いが、これを汎用圧縮ではなくタスク認識型にしている理由です。システムは圧縮出力をAmazon ElastiCache Serverlessに、キーをtakc:financial:mediumのように格納し、耐久性のためにS3にバックアップします。Redis OSSデータモデルは、マルチレートキャッシュルックアップに必要な階層キー構造をサポートします。キャッシュエントリは24時間のTTLを使用し、S3にバックアップされます。キャッシュエントリが退避または期限切れになった場合、クエリ関数はS3バックアップにフォールバックし、読み取り時にキャッシュを再入力します。
クエリパイプライン
ユーザーはAmazon Cognitoで認証され、JWTを受け取り、Amazon API Gatewayを介してクエリを送信します。AWS WAFは、レート制限と脅威保護のためにAPIの前面に配置されます。Lambda関数はヒューリスティック(キーワードシグナルとクエリ長)を使用してクエリの複雑さを分析し、Amazon ElastiCache Serverlessから適切な圧縮キャッシュを取得し、圧縮コンテキストとクエリをAmazon Bedrockに送信して推論を実行します。ルーティングの信頼度が低い場合、システムはセーフフォールバックとして中程度の圧縮階層をデフォルトで使用します。
コストのかかるBedrock圧縮呼び出しは、インジェスト中に1回だけ行われます。クエリパスは、キャッシュルックアップと圧縮コンテキストでの推論です。
このスタックは、圧縮と推論にAmazon Bedrock(Anthropic Claude 3 Haiku、Claude 3 Sonnet、Amazon Titan Text)を使用します。モデルの選択は、コード変更なしでCDKコンテキスト値を介して構成可能です。AWS Lambda(Python 3.12+)はデータ処理とクエリロジックを処理します。Amazon ElastiCache Serverlessは圧縮キャッシュを格納し、Amazon S3は生データ、チャンク、キャッシュバックアップ(KMS暗号化)を保持します。Amazon API GatewayはRESTエンドポイントを公開し、Amazon CognitoはJWTベースの認証を提供します。AWS WAFはレート制限とマネージドセキュリティルールでAPIを保護します。Amazon CloudWatchは監視とメトリクスを提供し、AWS Key Management Service(AWS KMS)は暗号化キーを管理します。
インフラストラクチャを単一のAWS Cloud Development Kit(AWS CDK)スタックとして定義し、単一のコマンドでデプロイします。AWS CDKは反復可能なデプロイを提供し、圧縮チャンクサイズ、Lambdaメモリ割り当て、ElastiCacheストレージ制限などのパラメータをコンテキスト値を介してカスタマイズできます。本番環境では、ステートフルリソース(S3、ElastiCache、Cognito)を独自のスタックに分離して、爆発半径を減らし、コンピューティング層とストレージ層の独立したライフサイクル管理を可能にすることを検討してください。
コスト比較
100,000トークンの知識ベースを1日1,000回クエリした場合の入力トークン使用量(出力トークンは応答長がコンテキストサイズに依存しないため除外):
アプローチ クエリあたりの入力トークン 1日あたりの入力トークン 相対入力コスト 完全コンテキスト 100,000 100,000,000 100% RAG(上位10チャンク) ~10,000 10,000,000 10% TAKC軽度(8倍) ~12,500 12,500,000 12.5% TAKC中程度(16倍) ~6,250 6,250,000 6.25% TAKC高(32倍) ~3,125 3,125,000 3.1% TAKC超(64倍) ~1,563 1,563,000 1.6%
圧縮によるトークン節約は、圧縮率に直接由来します。実際の節約額は、特定のBedrockモデル価格とクエリパターンによって異なります。
TAKCでは、インジェスト中に1回限りのBedrock呼び出しによる圧縮コストが事前に必要です。このコストは、変更がめったになく繰り返しクエリされる知識ベースの場合に償却されます。1時間ごとに変更される知識ベースの場合は、RAGのクエリごとの検索モデルの方が実用的な場合があります。
TAKC、RAG、または両方を選択する場合
次の表は、ワークロードの特性に基づいて、各アプローチが適している場合をまとめたものです:
要素 TAKCが有利 RAGが有利 クエリタイプ クロスドキュメント推論、統合 直接的な事実確認 知識ベースの安定性 1日以下で変更 1時間以上で変更 タスクの予測可能性 明確に定義されたタスクタイプ 予測不可能なクエリパターン カバレッジ要件 コーパス全体を考慮する必要がある 関連するドキュメントが少数のみ ソース帰属 不要 必要(ユーザーがソースを確認する必要がある) トークンバジェット 厳しい 柔軟
実際には、本番システムは多くの場合、両方の恩恵を受けます。RAGは迅速なルックアップを効率的に処理します。TAKCは、検索ベースのアプローチでは見逃されるつながりを必要とする分析クエリを処理します。クエリ複雑性アナライザーはそれらの間でルーティングできます。ユーザーが応答を特定のソースドキュメントにトレースする必要がある場合はRAGを使用します。クロスドキュメント推論と監査可能性の両方が必要な規制対象ワークロードの場合は、TAKCとRAGを組み合わせます。TAKCを分析応答に使用し、RAGを使用して監査証跡のためのサポートソースドキュメントを取得します。
はじめに
前提条件
この参考実装をデプロイするには、次のものがあることを確認してください:
Amazon Bedrockモデルアクセス権を持つAWSアカウント。
AWS CDK CLI。
Python 3.12以降。
Node.js 18以降。
デプロイ
[コスト削減のため原文は省略]