Ml-tips

faissのIndexIVFFlatにおける速度と精度のトレードオフ

はじめに

埋め込みベクトルで「似ているものを探す」とき、素直にやるなら全ベクトルとの距離を総当たりで計算します。正確ですが、ベクトルが数百万件になると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が必要、速度差はバッチサイズにも依存、という点に注意

参考リンク