Wrestling with ANN Indexes
ベクトル検索を初めて実装しようと思ったとき、素朴に全件比較を書いた。10万件くらいまでは動いた。100万件を超えたあたりで、レスポンスが秒単位になった。当然だ。100万件の768次元ベクトルに対してコサイン類似度を総当たりで計算すれば、そうなる。
コサイン類似度とは
2つのベクトルがどれくらい同じ方向を向いているかを測る指標。値は-1(真逆)〜1(完全一致)の範囲を取る。テキストや画像を数百次元のベクトルに変換したあと、「似ているか」を判定するのに最もよく使われる。
ここで登場するのがANN——Approximate Nearest Neighbor、近似最近傍探索だ。「近似」という言葉がすべてを物語っている。正確な最近傍を保証する代わりに、「だいたい近い」結果を桁違いの速度で返す。
ワイン1000万本の倉庫
たとえ話をしよう。1000万本のワインが並ぶ巨大な倉庫を想像してほしい。手元にある1本と「最も味が近い」ワインを探したい。愚直な方法は、1000万本すべてをテイスティングすること。味の違いがわかる舌を持っていたとしても、人生が終わる。
ANNは、倉庫の整理の仕方を工夫することで、全部を試さなくても「たぶんこのあたり」にたどり着く技術だ。そして「整理の仕方」にはいくつかの流派がある。
大きく分けると二つ。
IVF——倉庫をエリアに区切る
IVF(Inverted File Index)は、倉庫をエリアに区切る発想だ。赤のボルドー、白のブルゴーニュ、ロゼのプロヴァンス——といった具合にクラスタリングする。
検索時はまず「この1本はどのエリアに近いか」を判定して、そのエリアとその周辺だけを探す。1000万本を1000エリアに分ければ、見る量は1万本で済む。探索範囲が1000分の1になるので速い。
ただし、事前に「どうエリア分けするか」を決めるトレーニングが必要になる。全体のデータをある程度見渡して、クラスタの重心を計算しておかなければならない。この処理は重い。データが変わるたびにやり直すわけにもいかない。FAISSが代表的な実装だ。
k-meansとは
データをK個のグループ(クラスタ)に自動で分ける手法。各クラスタの「重心」を繰り返し計算しながら、近いデータ同士をまとめていく。IVFでは、このクラスタの重心が「エリアの代表地点」になる。
HNSW——ワイン同士を繋いだ地図
HNSW(Hierarchical Navigable Small World)は、ワインとワインを「似ている」という関係で繋いだネットワーク(グラフ構造)を作るアプローチだ。
上の階層は粗い地図、下の階層は詳細な地図。最初は粗い地図で大まかな位置を特定し、下の階層に降りて精密に探す。Google Mapsで国レベルから番地レベルへズームしていく感覚に近い。
HNSWの利点は、新しいワインが入荷したらグラフにノードを追加して近傍と繋げばいいという点だ。IVFのような事前トレーニングがいらない。ただし、グラフ全体をメモリに載せる必要があるので、ノード間のエッジ情報だけでも相当なメモリを消費する。
データが止まらない問題
ここまでがアルゴリズムの話。問題は、現実のシステムではデータが止まらないことだ。
ベクトルは毎秒追加され、その間も検索リクエストは飛んでくる。1000万本の倉庫に毎日新しいワインが届き、棚卸しをしながら客の注文にも応えなければならない。
この「書き込みと検索の両立」が、ベクトルDBの設計思想が分かれるところだ。
FAISSはライブラリなので、この問題を開発者に委ねている。train()とadd()が明確に分離されていて、trainは一度だけ実行し、あとはadd()で逐次追加。「棚卸しのタイミングは倉庫番が決めろ」というスタンスだ。
Milvusはこの問題を、倉庫で言えば「荷受け場」と「整理棚」の二層構造で解決した。新しいワインはまず入口の荷受け場(Growing Segment)に積まれる。棚番号もラベルもない。客から「似た味のワインを探して」と言われたら、スタッフは整理棚を案内マップで探しつつ、荷受け場は端から順に片っ端から味見する。荷受け場がいっぱいになったら閉鎖して、裏方が整理棚に移す作業を始める。その間も新しいワインは別の荷受け場に届くし、客の注文にも応え続ける。
Qdrantも似た二層構造を持っている。WAL(Write-Ahead Log)にまず書き込み、セグメントに反映する。セグメントは最初plain index——棚番号だけ振った未整理の状態——で、データが溜まるとoptimizer threadがバックグラウンドでHNSWインデックスを構築する。構築中も旧セグメントで検索は続く。Milvusとの大きな違いはHNSWベースであること。IVFのような事前トレーニングが不要なぶん、セグメントの切り替えが滑らかだ。
WeaviateはHNSWをデフォルトに据えている。グラフベースなので、新しいベクトルが到着した瞬間にノードとして地図に組み込める。荷受け場すら必要ない。全部が整理棚だ。ただし、グラフ全体をメモリに保持するため、同じデータ量でもIVFベースのアプローチよりメモリを食う。最近ではflatインデックスや、少量のうちはflatで始めて閾値を超えたらHNSWに切り替えるdynamicモードも追加されている。
ElasticsearchはLuceneのセグメントアーキテクチャの上に乗っている。各セグメントが独立したHNSWグラフを持ち、新しいドキュメントはメモリにバッファされたあと新しいセグメントとしてフラッシュされる。バックグラウンドのmerge policyが定期的にセグメントを統合し、そのとき最大セグメントのHNSWグラフを活かしつつ、他のセグメントのベクトルを追加する形で統合する。Elasticsearchを使い慣れた人なら、おなじみのセグメントマージがベクトルの世界でも動いているだけだとわかるだろう。
こうして並べてみると、各DBの設計思想は「トレーニングとサービングの境界をどこに引くか」という問いへの答え方の違いだとわかる。
| DB | 方針 |
|---|---|
| FAISS | 「そこは開発者が考えろ」 |
| Milvus / Qdrant | 「二層構造で吸収する」 |
| Weaviate | 「トレーニング自体をなくす」 |
| Elasticsearch | 「既存のセグメントマージに載せる」 |
どれが正解というわけではない。データの規模、更新頻度、メモリ予算、運用体制。選択の軸はいくつもある。
nvecdで実際にやってみた
nvecdという小さなベクトル検索エンジンを書いている。イベントの共起データとベクトル類似度を融合して、リアルタイムにレコメンデーションを返す。TikTokのフィードのようなユースケースを想定したものだ。
全件比較、意外と速い
nvecdが扱うのは32次元のベクトルだ。冒頭で触れた768次元(Sentence-BERTなどの言語モデル出力)とは桁が違う。次元が低いぶん計算は圧倒的に軽い。SIMD最適化を入れたら、1000万件でも1クエリ28.5ミリ秒で返ってきた。同時接続500人のシナリオでも成功率100%。
| データ量 | 1クエリの時間 | 500人同接 |
|---|---|---|
| 100万件 | 3.4 ms | 問題なし |
| 1000万件 | 28.5 ms | ギリギリ |
正直、しばらくこれでいいんじゃないかと思った。
ただ、28.5ミリ秒は「速い」のではなく「ギリギリ許容」だ。500人が同時にクエリを投げると、スレッドプールで処理待ちが発生する。負荷が1.5倍になった瞬間に破綻する。1000万件が見えているなら、ANNは必要だった。
FAISSを使わずに書いてみる
FAISSに依存する選択肢もあった。でもやめた。ビルドが重い、ライセンスの扱い、なにより「中身を理解しないまま使うライブラリ」を増やしたくなかった。
IVFの原理は単純だ。
- k-meansでベクトルをクラスタに分ける
- 各クラスタにどのベクトルが入っているか記録する(転置リスト——倉庫で言えば「エリアAにはこのワインが入っている」というリスト)
- 検索時は、クエリに近いクラスタだけ見る
自力ならいざ知らず、いまはAI全盛時代だ。書けないはずがない。
書けた。テストも通った。そして、サーバーに組み込んだ瞬間に地獄が始まった。
地獄その1:デッドロック
ベクトルを登録するVECSETコマンドは、VectorStore(ベクトルの保管庫)に書き込んだあと、IVFインデックスに「新しいベクトルが来たよ」と通知する。この通知処理が、VectorStoreの中身を読もうとする。
C++のstd::shared_mutexは、同じスレッドから2回read lockを取る(再帰ロック)ことを保証しない。規格上は未定義動作だ。小さなテストでは動く。16スレッドの負荷テストで黙って止まる。
修正は単純だった。lockを持ったまま通知するのではなく、ベクトルデータをコピーしてlockを解放してから通知するようにした。ただ、原因の特定に時間がかかった。
地獄その2:トレーニングが重すぎる
k-meansのトレーニングは重い処理だ。1000万件に対して1024クラスタを計算すると数十秒かかる。これがリクエスト処理スレッドの中で走ると、その間サーバーが止まる。
FAISSのドキュメントにはさらりと「trainは別で呼べ」と書いてある。あの一文の裏に、どれだけの設計判断が詰まっていたか、自分で実装して初めてわかった。
バックグラウンドスレッドで非同期にトレーニングを走らせ、完了するまではbrute-forceで検索する方式に書き換えた。
地獄その3:クラスタ数の計算ミス
IVFのクラスタ数(nlist)の目安はsqrt(n)——データ件数の平方根だ。100万件なら1000クラスタ。
ところが、トレーニングには全件ではなくサンプル(5万件)を使っていた。nを「サンプル数」で計算してしまい、sqrt(50000) = 223クラスタになった。
100万件を223クラスタに入れると、1クラスタに約4500件。検索時にnprobe=10クラスタを見ると、4500 × 10 = 45000件をスキャンすることになる。全件比較(100万件)よりは速いが、期待した「全体の1%だけ見る」には程遠い。
nlistは全データの件数から計算しなければならない。サンプリングを入れると見失いやすいポイントだ。
地獄その4:1件ずつ追加は遅すぎる
トレーニングが終わったあと、新しいベクトルを1件ずつIVFに追加していた。追加のたびにwrite lock(書き込みロック)を取得・解放する。100万回のlock/unlock。ロード速度は激減した。
さらに、トレーニング時にベクトルデータを全件コピーしていた。1000万件 × 32次元 × 4バイト = 1.2GB。メモリが逼迫し、macOSがプロセスをサイレントに殺した。
PQ(Product Quantization)——メモリ問題のもう一つの解
IVFと組み合わせてよく使われる手法にPQ(直積量子化)がある。ベクトルを短いコードに圧縮し、メモリ使用量を10〜100分の1に削減する。FAISSのIVF-PQは最も実用的な組み合わせとして広く使われている。nvecdでは32次元と低次元なため採用しなかったが、768次元以上のベクトルを扱うなら検討すべき選択肢だ。
二層構造を採用する
各プロダクトがどう解決しているかを調べ直した。共通していたのは、「書き込みと検索を別のデータ構造で扱う」という思想だった。
nvecdでもMilvus方式の二層構造を採用した。
ポイントは3つ。
- 書き込みは常にバッファへ。専用のlockだけで完結するので、IVFの検索やトレーニングと干渉しない
- 検索はバッファとIVFの両方を見て、結果をマージする。バッファが小さければbrute-forceでも十分速い
- sealは非同期。バッファの中身をswapで取り出し(lockは即解放)、バックグラウンドでIVFに流し込む。その間も書き込みと検索は止まらない
結果、100万件のロードが28秒で完了し、500人同接で成功率100%を達成した。
| 指標 | brute-force全件 | IVF二層構造 |
|---|---|---|
| ロード速度 | 128K ops/sec | 35K ops/sec |
| 500人同接 成功率 | 100% | 100% |
| 検索 p50 | 0.18 ms | 19 ms(バッファ込み) |
バッファに残った数万件のbrute-forceが検索レイテンシに効いてp50=19ミリ秒と、全件brute-force(3.4ミリ秒)より遅い。ただし、sealが進んでバッファが空になればIVF検索のみになる。10個のクラスタだけ見る世界は速い。
何がわかったか
自分で実装してみて思ったのは、ANNアルゴリズムそのものは教科書通りに書ける、ということだ。k-meansもIVFも、コアのロジックは100行で収まる。
難しいのは、並行システムの中にそれを組み込むことだ。
- lockの粒度をどうするか
- トレーニングをいつ、どのスレッドで走らせるか
- バッファとインデックスの整合性をどう保つか
- メモリの上限にどう対処するか
アルゴリズムの論文には書いていない。しかしプロダクションでは避けて通れない問題群だ。MilvusやQdrantのコードベースが大きい理由は、アルゴリズムが複雑だからではない。この「組み込み」の部分が膨大なのだ。
Close Enough
ひとつだけ確実に言えることがある。「近似」で十分な場面は、思ったより多い。
100%正確な最近傍と、99.5%の精度の近似最近傍。ユーザーが体感で区別できることは、まずない。完璧を求めて全件比較に固執するより、close enoughを高速に返すほうが、ほぼすべてのケースで正しい選択だ。
ワインの味も、ベクトル空間の距離も、「だいたい近い」で十分なことが多い。