はじめに
埋め込みベクトルで「似ているものを探す」とき、素直にやるなら全ベクトルとの距離を総当たりで計算します。正確ですが、ベクトルが数百万件になると1クエリごとに全件をなめることになり、検索が重くなります。
faiss(Facebook AI Similarity Search)は、この類似検索を高速化するライブラリです。少しの取りこぼしを許す代わりに、検索を桁違いに速くする近似最近傍(ANN)が使えます。この記事では、全探索と近似検索の速度と精度(recall)のトレードオフを実測します。
pip install faiss-cpu
データ:GloVe-100(ANN-Benchmarks)
ANN-Benchmarksが配布しているglove-100-angularを使います。GloVeの単語ベクトル約118万件・100次元に、クエリ1万件と**その正解の最近傍(ground truth)**が同梱されており、recallの評価にそのまま使えます。
import urllib.request, h5py, numpy as np
# ann-benchmarks.com はそのままだと 403 になるので User-Agent を付ける
url = "https://ann-benchmarks.com/glove-100-angular.hdf5"
req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
with urllib.request.urlopen(req) as r, open("glove-100-angular.hdf5", "wb") as f:
f.write(r.read())
f = h5py.File("glove-100-angular.hdf5", "r")
xb = np.asarray(f["train"], dtype="float32") # (1183514, 100) 検索対象
xq = np.asarray(f["test"], dtype="float32") # (10000, 100) クエリ
gt = np.asarray(f["neighbors"], dtype="int64") # (10000, 100) 正解の近傍index
print(xb.shape, xq.shape)
距離はangular(コサイン)なので、ベクトルをL2正規化して内積を類似度として扱います。
import faiss
faiss.normalize_L2(xb)
faiss.normalize_L2(xq)
全探索(IndexFlatIP)
まずは総当たりの全探索です。IndexFlatIPは内積で全件と比較する、近似なしの正確な検索です。各クエリの上位10件を取り、正解と突き合わせてrecall@10を測ります。ここでは1件ずつ検索(オンライン検索を想定)し、1スレッドでレイテンシを測ります。
検索時間はクエリごとにばらつくので、1件ずつ個別に計測して中央値(p50)と p10–p90 の幅で見ます。recallの方はインデックスが固定なら決定的なので、1回で確定します。
import time, numpy as np
faiss.omp_set_num_threads(1)
K = 10
flat = faiss.IndexFlatIP(xb.shape[1])
flat.add(xb)
def recall_at_10(index):
_, I = index.search(xq, K)
hit = sum(len(set(I[i]) & set(gt[i, :K])) for i in range(len(xq)))
return hit / (len(xq) * K)
def latency_percentiles(index, n):
for i in range(30): index.search(xq[i:i+1], K) # warmup
ts = np.empty(n)
for i in range(n):
t = time.perf_counter_ns()
index.search(xq[i:i+1], K) # 1件ずつ計測
ts[i] = (time.perf_counter_ns() - t) / 1e6 # ms
return np.percentile(ts, [10, 50, 90])
p10, p50, p90 = latency_percentiles(flat, 1500)
print(f"FlatIP: recall@10={recall_at_10(flat):.4f} median={p50:.2f} ms (p10-p90: {p10:.2f}-{p90:.2f})")
FlatIP: recall@10=1.0000 median=9.84 ms (p10-p90: 9.68-10.30)
全探索なのでrecallは1.0(これが正解の基準)ですが、118万件をなめるため1クエリ約9.8msかかります。これを速くするのが近似最近傍です。
近似最近傍(IndexIVFFlat)
IndexIVFFlatは、あらかじめベクトルをnlist個のクラスタに分けておき(trainが必要)、検索時はクエリに近い一部のクラスタだけを調べます。調べるクラスタ数がnprobeで、大きくするほど正確(recall↑)だが遅くなります。
nlist = 4096
ivf = faiss.IndexIVFFlat(faiss.IndexFlatIP(xb.shape[1]), xb.shape[1],
nlist, faiss.METRIC_INNER_PRODUCT)
ivf.train(xb) # クラスタ中心を学習
ivf.add(xb)
ivf.nprobe = 32
p10, p50, p90 = latency_percentiles(ivf, 3000)
print(f"IVF(nprobe=32): recall@10={recall_at_10(ivf):.4f} median={p50:.3f} ms")
IVF(nprobe=32): recall@10=0.8167 median=0.159 ms
同じデータで、nprobe=32なら全探索の約62倍速で、recallは0.82まで出ています。
速度と精度のトレードオフ
nprobeを振ると、recallと検索時間のトレードオフが描けます。
for nprobe in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
ivf.nprobe = nprobe
p10, p50, p90 = latency_percentiles(ivf, 3000)
print(f"nprobe={nprobe:4d}: recall@10={recall_at_10(ivf):.4f} "
f"median={p50:.4f} ms (p10-p90: {p10:.4f}-{p90:.4f})")
nprobe= 1: recall@10=0.328 median=0.0394 ms (p10-p90: 0.0376-0.0425)
nprobe= 2: recall@10=0.447 median=0.0452 ms (p10-p90: 0.0423-0.0497)
nprobe= 4: recall@10=0.560 median=0.0530 ms (p10-p90: 0.0487-0.0596)
nprobe= 8: recall@10=0.661 median=0.0688 ms (p10-p90: 0.0615-0.0795)
nprobe= 16: recall@10=0.746 median=0.0988 ms (p10-p90: 0.0871-0.1170)
nprobe= 32: recall@10=0.817 median=0.1591 ms (p10-p90: 0.1382-0.1919)
nprobe= 64: recall@10=0.875 median=0.2770 ms (p10-p90: 0.2392-0.3273)
nprobe= 128: recall@10=0.922 median=0.5238 ms (p10-p90: 0.4619-0.6081)
nprobe= 256: recall@10=0.957 median=0.9798 ms (p10-p90: 0.8802-1.0988)

左の図は横軸にnprobeを取ったものです。nprobeを上げるほど、調べるクラスタが増えてrecallは上がりますが、レイテンシ(中央値の線、帯は p10–p90)も伸びていきます。右の図はこれを「recall vs レイテンシ」で描き直したもので、横棒は p10–p90 の幅、右上の★が全探索(recall 1.0・約9.8ms)です。IVFの曲線は、nprobeを上げるほど★(正確な全探索)に近づきながら右(低速)へ動きます。
読み取れるのは、少しのrecallを譲るだけで、検索が大きく速くなるという点です。例えばrecall 0.92(nprobe=128)なら全探索の約19倍速、recall 0.96(nprobe=256)でも約10倍速です。どこまでrecallを求めるかは用途次第で、nprobeはその操作点を調整するつまみになります。
注意点
- コサイン類似度で検索したいときは、ベクトルを
normalize_L2してからMETRIC_INNER_PRODUCTを使います(正規化した内積=コサイン)。 IndexIVFFlatはtrain(クラスタ中心の学習)が必要です。IndexFlatIPのような全探索indexには不要です。nlist(クラスタ数)は√N前後がよく使われます。増やすと1クラスタが小さくなり、同じrecallをより少ないnprobeで得やすくなります。- 速度差はバッチサイズにも依存します。今回のように1件ずつ検索すると全探索は不利ですが、多数のクエリをまとめて検索するとBLASが効いて全探索も速くなり、差は縮まります。オンライン検索(1件ずつ)ほど近似検索の効果が出ます。
IndexIVFFlatは元ベクトルをそのまま保持するためメモリを食います。メモリを抑えたい場合はベクトルを圧縮するIndexIVFPQなどがあります。
まとめ
- faissは大量ベクトルの類似検索を高速化するライブラリ
IndexFlatIPは正確な全探索(recall 1.0)だが、件数に比例して遅いIndexIVFFlatはクラスタに分けて一部だけ探索する近似最近傍で、nprobeで速度と精度を調整する- GloVe-100(118万件)では、recall 0.92を保ったまま全探索の約19倍速で検索できた
- コサインは正規化+内積、IVFは
trainが必要、速度差はバッチサイズにも依存、という点に注意