電通総研 テックブログ

電通総研が運営する技術ブログ

RAGに捧げるベクトル検索パフォーマンスチューニング

こんにちは。電通総研 金融ソリューション事業部の若本です。
今回は、RAG内部で使用するベクトル検索の設定のチューニングについての記事になります。

RAGとは

RAG(Retrieval-Augmented Generation)は、ユーザーのクエリに対して外部の知識を検索し、それをもとにChatGPTのような大規模言語モデル(LLM: Large Language Model)に回答を生成させる手法になります。
具体的な方法は山下さんの記事でも紹介されていますので、興味のある方はぜひご一読ください。

ベクトル検索になぜチューニングが必要なのか?

RAGの検索において、よく使われる検索方法の1つがベクトル検索です。
ベクトル検索では文章を「ベクトル」という表現に変換して検索を行います。意味合いが似ている文章を、使われる単語が異なっていても検索できるため、非常に重宝される技術です。

しかし、「ベクトル」という複雑な表現同士で検索するには工夫が必要です。
例えば、langchainのfaissのVector storeではデフォルト設定がL2 normを用いた総当たりのベクトル検索(IndexFlatL2)になっています。
総当たりのベクトル検索では、ユーザーのクエリのベクトルとDBに保持しているベクトルを1件ずつ比較する必要があるため、DBに保持しているベクトル数が多いほど検索に時間がかかってしまい(O(n))、プロダクション向きではありません。

そこで、より高速なベクトル検索として近似近傍探索(ANN: Approximate nearest neighbor search)という技術があります。

近似近傍探索

近似近傍探索では、様々なアプローチで高速なベクトル検索の実装が試みられています。近似近傍探索の種類については、ANNベンチマーク)などを参照ください。
しかし、近似近傍探索も万能ではなく、速度が総当たりのベクトル検索より早いぶん総当たりと同じ検索結果を取得できない場合があります。検索結果をもとに回答を行うRAGの仕組みでは、本来取得できたはずの検索結果が抜け落ちてしまうことは性能の悪化に繋がります。

近似近傍探索の有名な手法の1つであるHNSWについて実装し、RAGを想定したパフォーマンスのチューニングを試みます。例えば、Azure AI Searchでも近似近傍探索としてHNSWが使用されています

HNSW(Hierarchical Navigable Small Worlds)は、ベクトル同士が連結されたグラフを順番に辿ることで、最も似ているベクトルを探策する手法です。詳細な理論については、Pineconeの記事などにわかりやすくまとめられています。
HNSWには3つのパラメータがあり、それぞれ増やすメリット/デメリットが存在します。

M:

  • 接続できるノードの数を表す
  • 増やすと検索精度が上がるが、検索時間と消費メモリも増加する

efConstruction:

  • インデックスの構築時に探索されるエントリポイント数
  • 増やすと検索精度が上がるが、検索時間※とインデックス構築時間も増加する

efSearch:

  • 検索中にレイヤー間で探索されるエントリポイント数
  • 増やすと検索精度が上がるが、検索時間が増加する

※Mが高い場合のみ

参考までに、Azure AI SearchにおけるHNSWのデフォルトのパラメータはmが低く、その他のパラメータが高く設定されています(引用元)。

今回は、上記のパラメータをベンチマークとして、HNSWを用いたベクトル検索でRAGを行うためのパラメータ探索を行います。

実験

検索対象とクエリについてデータを用意し、ベクトル検索の実行と評価を行いました。

検索対象

検索対象として、様々なトピックのPowerpointファイルを120件用意しました。
上記のファイルをlang chainのUnstructuredPowerPointLoaderで読み込み、Text splitterで分割しています。

なお、今回のText splitterの設定は下記としました。

chunk size(分割した文章単位の文字数): 1000
overlap size(文章のオーバーラップ): 500

作成後のindexのサイズは1069となりました。
また、ベクトル検索に使用するembeddingはOpen AI APIのtext-embedding-ada-002を使用しました。ada-002の次元数(ベクトルのサイズ)は1536です。

検索クエリ

今回はRAGで使用することを想定したチューニングのため、実際のRAGでユーザーが入力しそうなクエリをGPT-4で大量に生成しました。生成した件数は200件です。

評価方法

今回評価する内容は2つです。

  • 一致率(%)
    • クエリでベクトル検索を行い、ヒットした結果のセットが総当たり検索の結果セットと完全一致するかを確認します(単一の正解を定めてRecallを取っているわけではない)。今回は検索数を10としました。
    • このとき、結果セットの順序変動はRAGへの影響が(実際にはありますが)少ないとみて無視します。
  • 1件当たりの検索時間(s)
    • バッチ的に処理した検索時間を件数で割ったもの。
      faissに実装されているHNSWを用いて、M、efSearch、efConstructionのパラメータを探索しながら上記の評価を行います。

その他、本来であればインデックス作成時間なども計測すべきところですが、今回は対象外としました。

結果

M、efSearch、efConstructionを変えたときのHNSWの評価結果は下記のようになりました。

まずは一致率です。Y軸が一致率(%)を、X軸がefSearchを表しています。Mを変えた場合はラベルで、efConstructionを変えた場合はサブプロットで表現しています。

同様に、Y軸を検索時間(s)に変えて評価結果を集計しました。

ここで、Azure AI Searchで使用されているデフォルトのHNSWパラメータでは、今回の実験設定において5.5%のクエリで検索結果が異なる結果となっていました。検索速度は1件当たり約0.25ミリ秒を計測しています。
※今回はローカル上のfaiss実装(CPU)で実験を行っています。そのため実際のAzure AI Searchでの検索結果/検索速度ではないことにご注意ください。

実験の結果より、M、efSearch、efConstructionを抑えつつ、一致率を100%にする設定は下記となりました。
検索速度についても1件当たり約0.1ミリ秒となっており、上記の設定よりも高速に検索できています。

精度と検索速度が上がった一方、メモリの消費量は増加しました。HNSWにおける1件当たりのメモリ消費量は以下のとおりです(引用)。

4 × d + x × M × 2 × 4 ※d=ベクトルの次元数、x=インデックスの数

インデックスのメモリサイズを試算すると、デフォルトのHNSWパラメータでは43MB(1件あたり約40KB)、一致率100%の設定では79MB(1件あたり約75KB)となりました。精度と検索速度が改善した代わりに、インデックスのメモリサイズが約2倍に増加しています。

そこで、メモリの増加を防ぐためにMについてより詳細に探索します。Mを4から8の間に設定し、再度検証を行いました。

上記より、M=5のとき一致率が100%となる設定を発見できました。検索速度もデフォルトパラメータよりよい結果が得られています。

インデックスのメモリサイズについても52MB(1件あたり49KB)となり、ある程度抑えることができているため、今回の設定では上記のパラメータで検索インデックスを作成するのがよさそうです。

また余談ですが、今回は8-bit PQ(Product Quantization)を組み合わせたチューニングも実施しましたが、今回の実験設定では効果が出ない結果となりました。

おわりに

次元数の多いembeddingでのHNSWパラメータの探索や、RAG向けのベクトル検索についての技術記事があまり見当たらなかったため、今回は独自で検証を行いました。
結果として、今回はより一致率や検索時間に優れたパラメータを探索することができました。ベクトル検索のチューニングにおいては精度・検索速度・メモリのトレードオフとなるため、提供するRAGシステムの問題設定(ファイル数、文章ドメイン、RAGの仕様)やサービス設計などに合わせて同様の検証が求められます。 今回はHNSWに閉じた検証となりましたが、用途に応じて他の近似近傍探索手法を採用することも考えると非常に奥が深いですね。 この記事がどなたかの役に立てれば幸いです。ここまでお読みいただきありがとうございました。

執筆:@wakamoto.ryosuke、レビュー:@yamada.y
Shodoで執筆されました