本記事ではFaissやScaNNといったライブラリに実装されているIVF-PQ系の近似近傍探索手法のパラメータチューニングの際に気をつける点を紹介します。pythonのプログラム上で動かすことを想定していて、vertex AI vector searchのようなAPIで行うものは対象外です。ただ、OpenSearchではfaissを近似近傍探索として選ぶことができるため、チューニングの参考になるかもしれません。
はじめに: ANN-Benchmarksの罠
ANNの性能とパフォーマンスの参考になるサイトとして、ANN-Benchmarksというサイトがあります。このサイトでは各近似近傍探索のパフォーマンスが様々なベンチマークにより比較されており、近年ではFaissに実装されているFastScanやTensorFlow recommendersから使えるScaNNといった、高速化されたIVF-PQ系の手法が人気です。しかし、データセットの特徴やパラメータチューニングによってベストな手法は変わるので、この表のみを信じるのは危険です。まずはIVF-PQの仕組みについて解説します。
IVF-PQの概要
IVF
IVFは事前にデータをクラスタリングし、各クラスタ中心から各データに対して転置インデックスを貼り、検索時は検索データに近いクラスタの転置インデックスのアイテムを検索する手法です。faissではindex作成時にIVFXXでXX個のクラスタを作り、検索時のパラメータn_probeで検索するクラスタの個数を指定します。ScaNNではインスタンス作成時にnum_leavesでクラスタの個数を指定し、num_leaves_to_searchで検索するクラスタの個数を指定します。
PQ
PQはProduct Quantization(直積量子化)で、ざっくり説明するとD次元のベクトルをN次元ごと(通常は2次元)に分割し、各次元でクラスタリングを行った後クラスタ中心間の距離をコードブックとして持ち、推論時はコードブックから距離を計算します。一時はHNSWといったグラフ系の手法が優位に立ったのですが、FastScanやScaNNではこれをレジスタ内に収まるようなコードブックの大きさにすることで大幅な高速化をしました。
IVF-PQのチューニングの注意点
IVF-PQによって高速にベクトル検索が可能ですが、以下の3つの点に注意する必要があります。
- リランキングを設定する
- 小規模なデータに対してはIVFの分割数を減らす
- スループットが高速化しているか計測する
リランキングを設定する
PQのコードブックにおける計算はざっくりとしているため、精度が低いです。そのためfaissにおいてもScaNNにおいても、PQでスコアを計算したあとに正確な距離で計算し直すrerankingのパラメータが設定されています。
faissの場合
index factoryでfastscanを使う場合はRefineやRFlatを指定します。例えばIVFで1024個に分割し、28個のコードブックを作ってPQを行い、最後にtop100について正確な距離で計算するには
IVF1024,PQ28x4fs,RFlatを指定しindexを作成することでリランキングが設定されます。
リランキングのパフォーマンスと効果は、faissの公式wikiに結果が載っています(Preliminary experiment: IVF re-ranking)
ScaNNの場合
tfrs.layers.factorized_top_k.ScaNNにおいてはnum_reordering_candidatesが該当します。デフォルトではNone(リランキングしない)になっているので注意が必要です。
小規模なデータに対してはIVFの分割数を減らす
IVF-PQでは検索時に指定したIVFのクラスタ数のみ探索するため、探索したクラスタのアイテムの合計数が返して欲しいアイテム数を下回る場合、アイテム数が減った状態で返却されます。100のデータに対してIVFを100に設定し、10クラスタ探索する設定で50個アイテムを返してもらうことはできません。余裕を持ってIVFのクラスタ数は設定しましょう。
faissの場合
index作成時にIVFの分割数を決定するほか、index.nprobeに適当な数を入力することで探索するクラスタの数を増やすことができます。ただし、nprobeの数だけ検索時間がおよそ線形に増えていくので、パフォーマンスと比較しながらnprobeを決定しましょう
ScaNNの場合
tfrs.layers.factorized_top_k.ScaNNにおいてはnum_leavesで事前分割のクラスタ数を決定し、num_leaves_to_searchで探索数を決めます。デフォルトはそれぞれ100, 10と大きめなので、これもパフォーマンスと相談しながらチューニングしましょう。
スループットが高速化しているか計測する
IVF-PQを使うと理論上の計算量は減りますが、実際はIVFのみで直接検索したい距離を計算したほうがスループットが早い場合があります。近似近傍探索を用いてスループットが重要なケースにおいては、複数の手法を試しパフォーマンスを測定することが大切です。
実例:RVCにおけるチューニング
RVCは生の波形からHuBERTで特徴量をベクトルで抽出し、それを変換します。ただし学習データと予測データにおいては特徴量の分布が異なるので、学習データのHuBERT特徴量を保存し、予測データのHuBERT特徴量をクエリに学習データのHuBERT特徴量に近いものを検索し使うことで精度を上げています。ボイスチェンジャーAIのRVCでは部分的にfaissが使われているのでそのチューニングを行いドキュメントを作成したのですが、indexの量によってはFastScanよりもIVFの方がスループットが早い事がわかりました。また、学習データのHuBERTの特徴量をすべて保存し検索するのは重くて非効率です。そこで学習データのHuBERTの特徴量をkmeansでクラスタリングし、そのクラスタ中心からIVFで検索することで精度は落とさず大幅な検索パフォーマンスの向上を達成できました。(ddPn08さんのフォークに作った当時のPR)
まとめ
IVF-PQ系の近似近傍探索を用いる際のチューニングで気をつけることをまとめました。近似近傍探索の具体的なパラメータ設定はfaissの公式wikiを読むと詳しく書いてあります。例えばGuidelines to choose an indexはおすすめです。またFastScanの技術的詳細は東大松井研の4-bit PQの解説もおすすめです。近似近傍探索の仕組みを知ったうえで、ベクトル検索を活用していきましょう。
Discussion