Faissインデックスの使い方

背景

CLIPモデルを使った画像検索を実装したところ、検索対象の画像の枚数が10万枚になったくらいから検索速度が遅くなってきました。
「検索＝クエリと検索対象の特徴ベクトルを総当たりで類似度計算してランキングを返す」という処理なので検索時間は検索対象の増加に応じて線形に増加します。
私の場合、検索対象の数は億くらいまであり得るのでこのまま対策を考えないと詰むと思いつつも解が思い当たらず諦めていました。
そんな中、同じようにCLIPモデルを使った画像検索をしているブログを見つけて読んでいたところ、最後のほうで「もっと画像が増えたらベクトル計算用のDBとかFaissとか使って高速化しないといけないのかなぁ」としれっと書かれていました。
そこでFaissとかいうものを使えば高速化できるのか？と調べ始めたのがきっかけです。

使い方

まず、私は会社から支給されたノートPC（Win10, Core-i5, メモリ16GB）を使ってPythonで開発しているので下記の通りインストールしました。

pip install faiss-cpu

CUDA搭載のWin機もあったのでGPU版の「faiss-gpu」を導入することも考えたのですが、Linuxでしか使えないようでした。

IndexFlatL2

Faissで調べると近似最近傍探索というキーワードをよく見かけるので文字通り近似で、総当たりの検索と比較すると少なからず精度は失われると思っていたのですが、FlatとつくFaissインデックスは生の特徴量を全て持つので総当たりで検索すれば精度の低下はありません。
ただ、速度は普通にPythonで総当たりするより遥かに高速で、私の場合、768次元*25万ベクトルの検索に2分程度かかっていたものが数秒になり、感動しました。

# faissをインポート
import faiss

# dは特徴量ベクトルの次元数、特徴量ベクトルはnumpy配列のマトリクスでfloat32
d = feature_vectors.shape[1]

# インデックスを構築
index = faiss.IndexFlatL2(d)

# インデックスに特徴量ベクトルを追加
index.add(feature_vectors)     

# インデックスをファイルに書き出す
faiss.write_index(index,"faiss_index.bin")

たったのこれだけでIndexFlatL2は完成です。
私はこのインデックスをファイルに書き出して検索のときはDBアクセスを殆どしないようにしています。検索対象が数十万〜数百万程度で精度が重要ならこれで十分かもしれません。注意点としてはFlatでは特徴量のインデックスは追加した順に0から連番という感じで振られるのでもともとの検索対象の特徴量を独自のIDで管理している場合はFaissのインデックスと独自IDの対応付けを管理しておかなければなりません。
私の場合は、特徴量ベクトルを1から始まるIDで管理していたので1ずれの罠で少しハマりました。

IndexIVFFlat

これもFlatなので生の特徴量を全て格納します。IndexFlatL2との違いは特徴量ベクトルをクラスタリングするところです。この関係でquantizerという別のインデックスが必要になります。どの特徴量がどのクラスタに所属しているのかを特定するために二重でインデックスが必要なんだろうと考えるようにして納得しています。このquantizerには基本IndexFlatL2を使えば良いようです。

# クラスタの数
nlist = 100

# もうひとつのインデックス
quantizer = faiss.IndexFlatL2(d)

# dは検索対象の特徴量ベクトルの次元
index = faiss.IndexIVFFlat(quantizer, d, nlist)

# クラスタリングの処理
index.train(feature_vectors)

# IndexFlatL2のときより時間がかかります
index.add(feature_vectors)

# インデックスをファイルに書き出す
faiss.write_index(index,"faiss_index.bin")

クラスタの数nlistを特徴量ベクトルの数と同じにすると1クラスタに1ベクトルが格納されることになります。検索時にnprobeという変数で何個のクラスタから検索するかを決められるのですが、これをクラスタの数nlistと同じにすると結果的には総当たりと同じになります。この辺りを理解しておくと100万ベクトルあるときにnlist=100でnprobe=10とかにすると1クラスタあたり1万ベクトルくらい含まれているのだろう（100万割る100）から、そのクラスタを10個検索するとなると10万ベクトル分（1万掛ける10）の類似度計算くらいで済むのかな、という感覚を得られるのでパラメータ調整の目安になると思います。

IndexIVFPQ

ここまではロスレス（生特徴量を持っている）なインデックスでしたが検索対象が大きくなるほど対応しきれなくなってきます。そこで特徴量ベクトルをロッシーな方法で圧縮することで情報量を削減するオプションもあり、IVFPQがその一例です。mというsubquantizerのバイト数が新たに変数として加わります。先ほど追加したquantizerに続けてsubquantizerが追加されたわけなのでさらにもう一層処理が増えていることが予想できます。厳密さを激しく無視した自分のイメージを言うとベクトルを量子化して間引きするイメージです。mが小さいほど間引く数が増えるので、もともと持っている情報は失われる代わりに軽くなると理解しています。なお、mはdの約数である必要があります。私の場合はd=768次元だったので控えめにm=384にしました。さらにFlatでは使えなかったindex.add_with_idsが使えます。これは読んだ通りで特徴量を追加するときに自分で管理している特徴量のIDをインデックスとして採用してもらえます。つまりFaissインデックスと自分の管理用IDを対応づけてくれます。

# クラスタの数
nlist = 100

# subquantizerのバイト数
m = 384

# もうひとつのインデックス
quantizer = faiss.IndexFlatL2(d)

# 8はサブベクトルが8bitで符号化されるということ
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)

# クラスタリングの処理
index.train(feature_vectors)

# 独自で特徴量ベクトルを管理しているIDがあれば渡すと紐づけてくれる
index.add_with_ids(feature_vectors, image_ids)

# インデックスをファイルに書き出す
faiss.write_index(index,"faiss_index.bin")

検索

ここまで、いくつかのインデックスとその作り方を見てきました。ここからは作ったインデックスを使って実際に検索をする方法を説明します。
index.searchを使うだけなのですが、いくつかパラメータを設定する必要があります。一つ目は上位何番目のインデックスまで返すかというtop_kと、二つ目は何個のクラスタを検索するのかというnprobeです。特にnprobeは指定しないとデフォルトでは1で検索されてしまうのでクラスタの数などによっては全然精度が出ないということが起こってしまいます。ここでも少しハマりました。

# 上位何位までのインデックスを返すか
top_k = 100

# Faissインデックスの読み込み
index = faiss.read_index("faiss_index.bin")

# 何個のクラスタから検索するか（デフォルトは1）
index.nprobe = 10

# 検索処理 Dに距離、Iにインデックスが返される
D,I = index.search(query_vector, top_k)

結果

私の環境では117万ベクトルに対してnlist=128,m=384のIVFPQを算出し、nprobe=10で検索したところ、0.2秒程度で総当たりとほとんど遜色ない結果になりました。（定量的な評価は一切やっていないので感覚の話です。）計算速度に絶望していた頃から考えたら信じられない境地です。IVFPQの算出時にメモリをかなり喰いますが一度ファイル化してしまえば以後はほぼメモリ消費せずに使えます。実際、メモリ8GBのサーバでも爆速で検索できて驚きました。こんな技術が無料で使えてしまうのだからMetaに感謝せざるを得ません。

最後に

何も知らない状況からLLM任せでFaiss対応しようとしたのですが、全然うまくいきませんでした。そこで様々な日本語のFaiss解説ブログを読んでみたのですが、結局公式の英語のwikiのチュートリアルが一番わかりやすかったのでこの記事もwikiの流れ通りに書いてみました。自身の備忘録として書きましたが、誰かの助けになれば幸いです。