最近、文をembeddingsといった特徴ベクトルに変換するユースケースが増えている。そのベクトルから類似ベクトルを探す時に、数千ベクトルならほぼ何も考えなくともよく、数万ベクトル〜になると検索速度を高速化するためHNSW等のANNの近似最近傍探索アルゴリズムを使い、そして数百万ベクトル〜になってくると現実的なデータサイズ収めるために量子化等々を組み合わせた最適化を行うことが多いだろう。

これら類似ベクトル検索のための最適化(HNSW・IVFといったアルゴリズムや量子化)では、検索速度、データサイズ(メモリに乗るか)、精度、三つのトレードオフが発生する。これらトレードオフを踏まえた最適化戦略を考えるのだが、最適化時の正確さの計測結果として recall@10 や recall@100 が掲載されていることを多く見かける。例えばChoose the k-NN algorithm for your billion-scale use case with OpenSearchでは recall@10 で評価しているし、Byte-quantized vectors in OpenSearch では recall@100 である。

embeddings以外の情報と組み合わせて検索する場合や、reranker によって再ソートする場合は recall@10,100 でも良いかもしれないが、RAG での利用の場合、LLM や Reader でのタスクの In-Context に入れる際に top-10 を入れることはあまりないと感じる。top-3 や top-5 を入れるケースが自分の場合は多い。というわけで、recall@1,2,3,5 を faiss + 代表的なアルゴリズム+量子化で評価し、検索速度・データサイズ・精度(recall)を測定してみた。なおライブラリにfaissを使った。faiss でなくとも主だったアルゴリズムや量子化はベクトル検索エンジンにはたいてい実装されているし、OpenSearch では内部で faiss をベクトル検索エンジンとして利用することも可能だ。

ベクトル検索データベースでの、アルゴリズムとパラメータ

測定の前に、インデクス作成時のアルゴリズムとパラメータについてのおさらい。メジャーどころとして、IVFやHNSWの探索アルゴリズムと、PQ(Product Quantization)による圧縮、各々のパラメータが挙げられる。以下は各々の説明のGPT4出力に多少手を加えたもの。なお、そもそもの近傍探索については、松井先生のグラフを用いた近似最近傍探索の理論と応用にしっかりと書かれているので、一読をお勧めする。

• HNSW (Hierarchical Navigable Small World)
  • アプローチ: HNSWは、グラフベースのアプローチを使用します。各ノード（データ点）は、隣接するノードとのリンクを持っており、これらのリンクを介して効率的に検索を行います。
  • 特徴:
    • 高速な検索: グラフの階層構造により、近似的な最近傍探索が非常に高速に行えます。
    • 動的な追加: 新しいデータ点を動的に追加することが可能です。
    • 高精度: 他の多くの近似アルゴリズムと比較して高い精度を提供します。
    • メモリ使用量: グラフ構造のため、メモリ使用量が比較的大きい可能性があります。
  • パラメータ: M
    • 各データポイントが持つ近傍ノードの最大数
• IVF (Inverted File Index)
  • アプローチ: IVFは、データを複数のクラスタに分割し、それぞれのクラスタに対して個別のインデックスを作成します。検索時には、クエリに最も近いクラスタを識別し、そのクラスタ内でのみ検索を行います。
  • 特徴:
    • 効率的な大規模検索: 大規模なデータセットに対して、検索を効率化できます。
    • スケーラビリティ: データセットが大きい場合にも適用可能で、スケーラブルです。
    • カスタマイズ可能: クラスタの数（nlist）や量子化のレベルなど、多くのパラメータをカスタマイズできます。
    • メモリ使用量: HNSWよりもメモリ効率が良いことが多いですが、クラスタ数と量子化のレベルにより異なります。
  • パラメータ: nlist
    • クラスタの数
• 両者の比較
  • 精度: HNSWはIVFよりも一般的に高い精度を提供しますが、メモリ使用量が増加する傾向にあります。
  • スピード: HNSWは高速な検索が可能ですが、IVFは大規模なデータセットに対してよりスケーラブルで、メモリ効率が良いです。
  • 使用場面: HNSWは精度が重要な場合やリアルタイム検索に適しています。IVFは大規模なデータセットや限られたメモリリソースでの使用に適しています。
• Product Quantization (PQ) の概念
  • Product Quantizationは、高次元ベクトルを効率的に圧縮するためのテクニックです。このアプローチでは、以下のステップが含まれます：
    • ベクトルの分割: 各ベクトルは、より小さい次元の複数のサブベクターに分割されます。
    • サブベクターの量子化: それぞれのサブベクトルは、個別の小さなコードブック（事前定義された値のセット）を用いて量子化されます。これにより、各サブベクターはコードブック内の最も近い値にマッピングされます。
    • 圧縮された表現: 最終的に、元のベクトルはこれらの量子化されたサブベクターの組み合わせとして表現されます。

HNSWの場合の主パラメータはM、IVFはnlist(とここには出てきていないがmbit)となる。またPQを使う場合、サブベクター数となる。サブベクター数は、大元のベクトルの次元を割れる数、例えば大元が384次元なら32,64,96などとなる。また、先ほどのパラメータ指定は学習前に必要だが、グラフノードやクラスタをどれだけ探索するかは efSearch(HNSW), nprobe(IVF) といったパラメータで「検索の実行時」に決めることができる。

なお、faiss + python 環境では、以下のように index_factory() を使うことで、文字列から生成する方法もあり、便利である。

faiss.index_factory(d, "IVF2048,PQ64") # nlist = 2048, PQ = 64
faiss.index_factory(d, "HNSW32,PQ64") # M = 32, PQ = 64

利用するデータセットとコード

今回は faiss のコードでよく使われているANN_SIFT1Mを使った。名前の通り1M(100万件)の128次元のベクトルデータセットである。これらの中から検索用クエリ10000件を使って、recall@N を計測する。faiss は検索にGPUも使えるが、たいていの検索ではCPUで行うであろうため、CPU(Ryzen9 5950)を利用している。

ベンチマークに使ったコードはbench_gpu_sift1m_ivf_hnsw.py。このソースコードを git clone した faiss の benchs ディレクトリに入れて、データセットを適切に配置して実行すれば再現できる、はず。

測定結果

• https://docs.google.com/spreadsheets/d/1ZsMJZf-4tgKgfSa4zvLcmZrBZpoy6Op_EOgY17GPOfY/edit?usp=sharing

結果のシートは二枚あり、全部を載せると大きすぎるので、片方のシートは抜粋のデータ。

IVFやHNSW、PQ(量子化)は冒頭に述べた通りのトレードオフの関係となった。IVFは抜群に省メモリ(小さなデータサイズ)だが速度と正確性がHNSWよりも悪い。ただメモリに乗らなくなった途端に、HNSWも速度が劣化するであろう。そして精度のRecall@1-100の値を見ると面白い。Recall@100, 10 は割とすぐ 1.0 に迫るが、上位3・5件に同じデータが含まれるかの Recall@3, 5 が指標として大事なシステムだとすると、ちょうど良い塩梅のパラメータはよく考えて選択する必要がありそう。また、Recall@1 の top-1 は PQ 時には一致する確率が結構低い。ので、top-1 が大切な状況ではどうすべきかなど、よく考慮する必要がありそうだ。

まとめ

RAG の top-3 や top-5 が大切なシステムにおいて、Reranker 使わないのに大きめな Recall で評価して問題ないと判断してしまうと、実際に欲しい結果と乖離がでそうな結果。というわけで、当たり前の話だけど、やりたいことに対して指標を考え、それに適した最適なパラメータをセットしましょう、というお話でした。この結果もあくまで SIFT1M データセットに対しての各種パラメータの結果なので、また別のデータセットでは違った結果になるはず。

何も考えないベストなパラメータなどは無いが、100万件を超えるような(とりわけ数千万件〜にも到達しそうな)ベクトルに対してはIVF+量子化を使い、nlist は1024(実際の件数のsqrt(トータル件数)付近のあたりが良さそうなので、1000万件なら4096など)で、PQはできる限り大きく、例えば128次元のベクトルなら64など、が良いと感じている。精度と速度重視で〜数百万件ベクトルなら、HNSW + PQも十分現実的なデータサイズで扱えそうなので、こちらも検討すべきであろう。

なお、faiss の benchs ディレクトリには、IVFやHNSW以外の計測や、例えばPCAで次元削減するベンチマークなど、ベクトル検索でこんなことやりたいよう、という計測も結構行われていたりするので、一見の価値がある。

おまけ・kaggle コンペで失敗した話

kaggle のすでに終了したコンペLLM Science Examでは、RAGでtop1~5 を入れてスコアを伸ばす解法が多く、多分にもれず自分も行っていたのだけど、この時 faiss のパラメータを割と適当に決めていたため、データサイズはかなり小さくなったのだが、実は精度が結構下がってしまっていたことに気づかずじまいであった。なので、精度が大切な場面では、ちゃんと計測してパラメータを考えよう、と過去の自分に教えてあげたい。