- Jeff Deanの30億個ベクトルクエリ問題をきっかけに、最適な map-reduce ソリューションを自ら実装してみる過程を扱った技術実験の記録
- 768次元の float32 ベクトル30億個と1,000個のクエリベクトルの dot product を計算する naive 実装から始め、numpy のベクトル化と float32 変換で段階的な性能改善を達成
- 3,000ベクトル基準で naive 方式の約2秒から、ベクトル化後は0.01秒、float32 適用時は 0.0045秒 まで短縮
- 30億個規模に拡張すると結果行列に約 8.6TB のメモリが必要となり、OOM 問題が発生。バッチ処理・メモリマッピング・Rust/C への書き換え・SimSIMD のような追加最適化が必要
- 技術的なソリューションよりも、要件定義(返却形式、マシンスペック、圧縮許容可否など)のほうが難しい中核課題であることを強調
問題の出発点
- Jeff Deanと30億個ベクトルクエリ問題に関するやり取りを見て、最適な map-reduce ソリューションを自分で実装してみようとした試みから始まる
- ベクトルは
n 次元の浮動小数点数配列であり、ベクトル検索は意味的に近い単語や項目を見つけるために使われる
- 検索、推薦、Cursor のような 生成系検索アプリケーションで埋め込みを活用する一般的なパターン
Naive 実装
- 初期の想定は、検索対象となる30億個の document ベクトルと約1,000個のクエリベクトルがあり、いずれもディスクに
.npy 形式で保存されているというもの
- ベクトル次元は 768 で、多くの類似度ベースの埋め込みクエリモデルで使われる一般的なサイズ
- 各クエリベクトルを全文書ベクトルと比較して dot product を計算し、その結果を返す二重ループ方式
- まず3,000個のベクトルでテストしたところ、M2 MacBook で 約2秒 を要した(30億個より100万分の1の規模)
ベクトル化(Vectorized)最適化
- numpy の行列積演算(
vectors_file @ query_vectors.T)で二重ループを置き換える ベクトル化方式を適用
- 3,000ベクトル基準で 0.0107秒 へ大幅改善
- 300万ベクトルに拡張すると 12.85秒 を要した
float32 変換
- データを
np.float32 に変換して追加の最適化を実施
- 3,000ベクトル基準で 0.0045秒 まで短縮
- 300万ベクトルで約13秒なので、30億個ではその1,000倍の約 3,216分 と推定
性能比較の要約
- Naive 方式(3,000ベクトル): 1.9877秒
- ベクトル化方式(3,000ベクトル): 0.0107秒
- ベクトル化方式(300万ベクトル): 12.8491秒
- float32 方式(3,000ベクトル): 0.0045秒
OOM 問題と追加最適化の方向性
- 30億個のオブジェクトを float32(4バイト)で処理する際に必要なメモリは約 8.6TB で、実行時に OOM が発生
- Jeff Dean が示した方向に沿った追加最適化が必要:
- ジェネレーターとバッチ比較演算を使うようコードを変更
- 一定の計算ごとにディスクへ書き出す、または メモリマッピングを使用
- Rust または C でシステムレベル最適化コードへ書き換え
- SimSIMD のような大規模ベクトル類似度比較専用ライブラリを活用
要件定義の重要性
- 追加最適化の前に、問題そのものに不明確な点が多いことが見えてきた:
- 単一クエリベクトルで30億件全体を1回検索して全結果を返すのか、それとも top-k ANN 検索なのか
- 元のベクトルも一緒に返す必要があるのか、類似度にもとづく 再ランキングが必要なのか
- 1,000個のクエリベクトルで全体を1回走査するのか
- ベクトルはすでにメモリ上にあるのか、ディスクから1つずつ読むのか、ストリーミング方式なのか
- ローカル実行なのか、マシンスペックは何か、GPU 利用可能かどうか
- 埋め込みサイズが結果の表現や入出力ベクトルサイズに与える影響
- float64 から float32 への 圧縮が精度面で許容されるのか
- 単発プロジェクトなのか、生成にどれくらいの時間が与えられているのか
- 技術的なソリューションそのものより、要件を正確に定義することが最も難しい課題
まだコメントはありません。