DiskANN、10億スケールのデータセットで高回収率と高QPSを達成したディスクベースのANNSソリューション
ZillizのR&DエンジニアであるChengming Liは、SouthEast Universityでコンピュータサイエンスの修士号を取得。現在、グラフベースと量子化ベースのソリューションを含む、高次元データのANNS問題に注力している。
"DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node "は、2019年にNeurIPSで発表された論文である。この論文では、わずか64GBのRAMと十分な容量のSSDを搭載したシングルマシンを使って、10億点規模のデータセットに対してインデックス構築と検索を実行する最先端の手法が紹介されている。さらに、大規模データセットにおけるANNS(近似最近傍探索)の3つの要件である、高リコール、低レイテンシ、高密度(単一マシンのノード数)を満たす。本手法は、64GBのRAMと16コアのCPUを搭載したシングルマシンを用いて、10億スケールのデータセットSIFT-1Bにグラフベースのインデックスを構築し、5000QPS(queries per second)、95%以上の再現率(recall@1)、3ms以下の平均待ち時間を達成した。
著者
Suhas Jayaram Subramanya:元マイクロソフト・インディア研究所社員、CMU博士課程在籍。主な研究テーマは、ハイパフォーマンス・コンピューティングと大規模データの機械学習アルゴリズム。
デヴヴリットテキサス大学オースティン校大学院研究助手。研究テーマは理論計算機科学、機械学習、ディープラーニング。
ローハン・カデコディテキサス大学博士課程在籍。主な研究分野は、永続ストレージ、ファイルシステム、KVストレージなどのシステムとストレージ。
ラヴィシャンカール・クリシャスワミ:マイクロソフトインド研究所主任研究員。CMU博士。研究テーマは、グラフとクラスタリングに基づく近似アルゴリズム。
Harsha Vardhan Simhadri:マイクロソフト インド研究所 主任研究員。CMU博士。過去に並列アルゴリズムとランタイムシステムを研究。現在の主な仕事は、新しいアルゴリズムの開発とプログラミングモデルの作成。
動機
主流のANNSアルゴリズムのほとんどは、インデックス構築性能、検索性能、想起率の間でトレードオフの関係にある。HNSWやNSGのようなグラフベースのアルゴリズムは、検索性能と想起率の点で、現在のところ最先端の手法である。メモリに常駐するグラフベースのインデックス作成手法は、メモリ占有量が多すぎるため、限られたメモリ資源しか持たない1台のマシンで大規模なデータセットのインデックス作成と検索を行うことは比較的困難である。
ユークリッド距離ベースのANNSを10億個規模のデータセットに適用する場合、多くのアプリケーションで迅速な応答が要求される。以下に2つの主要な解決策を示す:
- 転置インデックス+量子化:データセットをM個のパーティションにクラスタリングし、PQ(積量子化)などの量子化スキームを使ってデータセットを圧縮する。このソリューションでは、データ圧縮によって精度が低下するため、リコールが低くなる。topkを増加させることで、リコールを向上させることができるが、QPSはそれに応じて低下する。
- 分割してインデックスを作成する:データセットをいくつかの不連続なシャードに分割し、それぞれのシャードに対してインメモリーインデックスを作成する。クエリー要求が来ると、各シャードのインデックスに対して検索を行い、マージ後に結果を返す。このソリューションでは、データセットの規模が拡大しすぎるため、1台のマシンのメモリリソースが制限され、より多くのマシンが必要になり、QPSが低くなる。
上記のいずれの解決策も、1台のマシンのメモリ制限によって制限されている。本稿では、この問題を解決するためにSSD常駐型のインデックス作成メカニズムの設計を提案する。SSD常駐型インデクシングの課題は、ランダムディスクアクセス数とディスクアクセス要求数を削減することである。
寄稿
本論文では、大規模データセットの検索を効果的にサポートするSSD常駐ANNS方式DiskANNを提案する。この方式は、本稿で紹介するグラフベースのアルゴリズムに基づいている:Vamana。本論文の貢献は以下の通りである:
- DiskANNは64GB RAMを搭載した1台のマシンで、100次元を超える10億スケールのデータセットにインデックスを付け、検索することができる。
- Vamanaと呼ばれる新しいグラフベースのアルゴリズムは、NSGやHNSWよりも小さな探索半径を持ち、ディスクアクセス回数を最小化するために提案された。
- Vamanaはメモリ上で動作し、その性能はNSGやHNSWよりも低くはない。
- 大規模データセットの重複するパーティションに構築された小さなVamanaインデックスは、接続性を失うことなく1つのグラフにマージすることができる。
- VamanaはPQのような量子化スキームと組み合わせることができる。グラフ構造と元のデータはディスクに保存され、圧縮されたデータはメモリに保存される。
Vamana
このアルゴリズムはNSG[2][4]の考え方に似ている(NSGを理解していない人は参考文献[2]を、論文を読みたくない人は参考文献[4]を参照されたい)。これらの主な違いはトリミング戦略にある。正確には、NSGのトリミング戦略にスイッチ・アルファが追加されている。NSGのトリミング戦略の主な考え方は、ターゲット点の近傍をできるだけ多様に選択することである。もし新しい近傍点が目標点よりも目標点の近傍点に近ければ、この点を近傍点セットに加える必要はない。言い換えると、ターゲット点の各近傍点に対して、周囲半径dist(ターゲット点、近傍点)内に他の近傍点が存在してはならない。このトリミング戦略は、グラフのアウトディグリーを効果的に制御し、比較的過激である。これはインデックスのメモリフットプリントを減らし、探索速度を向上させるが、探索精度も低下させる。Vamanaのトリミング戦略は、パラメータalphaによってトリミングの規模を自由に制御するものである。動作原理は、トリミング条件におけるdist(近傍点、候補点)にパラメータα(1以上)を乗算する。dist(目標点、ある候補点)が拡大された基準距離よりも大きい場合にのみ、トリミング戦略が採用され、目標点の近傍点間の相互排除の許容範囲が拡大される。
Vamanaのインデックス作成プロセスは比較的単純である:
- ランダムなグラフを初期化する;
- これはNSGのナビゲーション・ポイントと同様である。まず、グローバル重心を見つけ、次にグローバル重心に最も近い点をナビゲーション・ポイントとして見つける。VamanaとNSGの違いは、NSGの入力はすでに最近傍グラフであるため、ユーザーは単純に初期近傍グラフ上で直接セントロイド点の近似最近傍探索を行うことができる。しかし、Vamanaはランダムな最近傍グラフを初期化するため、ユーザーはランダムなグラフ上で直接近似探索を行うことができません。そのため、ユーザは直接ランダムグラフ上で近似探索を行うことができない。そのため、ユーザはグローバル比較を行い、後続の反復探索の開始点となるナビゲーション点を得る必要がある。この点の目的は平均探索半径を最小化することである;
- 初期化されたランダム近傍グラフとステップ2で決定された探索開始点に基づいて、各点に対して近似最近傍探索を実行し、探索経路上の全ての点を近傍候補集合とし、α=1を用いたエッジトリミング戦略を実行する。NSGと同様に、ナビゲーション点から始まる探索経路上の点集合を候補近傍集合として選択すると、長いエッジが増え、探索半径を効果的に小さくすることができる。
- α>1(論文では1.2を推奨)に調整し、ステップ3を繰り返す。ステップ3がランダムな最近傍グラフに基づいているのに対し、最初の反復の後、グラフは低品質になります。したがって、想起率に非常に重要なグラフ品質を改善するために、もう1回反復する必要がある。
本稿では、Vamana、NSG、HNSWの3つのグラフインデックスを比較する。インデックス作成とクエリー性能の点では、VamanaとNSGは比較的近い。データについては以下の実験セクションを参照。
2.png
Vamanaインデックスの構築プロセスを視覚化するために、この論文では200の2次元点を用いて2回の反復をシミュレートしたグラフを提供している。最初の行では、alpha = 1を用いてエッジをトリミングしている。このトリミング戦略は比較的急進的であり、多くのエッジがトリミングされていることがわかる。alphaの値を大きくし、トリミングの条件を緩くすると、明らかに多くのエッジが追加される。最終的なグラフでは、かなり長い辺が追加されている。これは効果的に探索半径を小さくすることができる。
DiskANN
64GBのメモリしか持たないパソコンでは、インデックスを構築するどころか、10億個の生データを保持することすらできない。そこには2つの課題がある:1.1.限られたメモリー資源で、このような大規模なデータセットにインデックスを作成するには?2.元データがメモリにロードできない場合、検索時にどのように距離を計算するか?
この論文では、以下の解決策を提案している:
- 最初の課題:まず、k-meansを用いてデータをk個のクラスターに分割し、各点を最も近いi個のクラスターに割り当てる。各クラスタに対してメモリベースのVamanaインデックスを構築し、最後にk個のVamanaインデックスを1つに統合する。
- 2つ目の課題:元のベクトルにインデックスを作成し、圧縮されたベクトルをクエリする。元のベクトルにインデックスを作成することでグラフの品質を確保し、圧縮ベクトルはメモリにロードして粗視化検索を行うことができる。圧縮ベクトルで検索すると精度が落ちる可能性があるが、グラフの品質が十分高い限り、大まかな方向は正しい。最終的な距離結果は元のベクトルを使って計算される。
DiskANNのインデックスレイアウトは、一般的なグラフインデックスと似ています。各点の近傍集合と元のベクトルデータは一緒に保存される。これはデータの局所性をより有効に利用するためである。
前述したように、インデックスデータがSSDに格納されている場合、検索遅延を少なくするために、ディスクアクセスの回数とディスクの読み書きの要求をできるだけ減らす必要があります。そこで DiskANN は2つの最適化戦略を提案します:
- キャッシュ ホットスポット: メモリ内の開始点から C ジャンプ以内のすべての点をキャッシュします。Cの値は3~4以内に設定するのがよい。
- ビーム探索:簡単に言えば、近傍情報を事前に読み込むことである。点pを検索するとき、pの近傍点がメモリにない場合はディスクからロードする必要がある。少量のSSDランダムアクセス動作はSSDシングルセクタアクセス動作と同程度の時間を要するので、一度にW個の非アクセス点の近傍情報をロードすることができる。Wは大きすぎても小さすぎてもいけない。Wが大きいと計算資源とSSD帯域幅を浪費し、小さいと探索遅延が増大する。
実験
実験は3つのグループで構成される:
メモリベースのインデックス間の比較:Vamana VS.NSG VS.HNSW
データセットデータセット:SIFT1M(128次元)、GIST1M(960次元)、DEEP1M(96次元)、DEEP1Bからランダムにサンプリングした1Mのデータセット。
インデックス・パラメータ(すべてのデータセットで同じパラメータ・セットを使用):
HNSW:M = 128、efc = 512。
Vamana:R = 70, L = 75, α = 1.2.
NSG:R = 60、L = 70、C = 500。
検索パラメータは論文には記載されていないが、これはインデックス作成パラメータと一致する可能性がある。パラメータ選択については、論文で言及されているNSGのパラメータは、NSGのGitHubリポジトリに記載されているパラメータを基に、よりパフォーマンスの良いグループを選択している。VamanaとNSGは比較的近いので、パラメータも近い設定になっている。しかし、HNSWのパラメータ選択の理由は示されていない。我々は、HNSWのパラメータMが比較的大きく設定されていると考えている。グラフベースのインデックスのアウト度が同じレベルに設定されていないと、グラフベースのインデックス間の比較に説得力がなくなる可能性がある。
上記のインデックス作成パラメータの下で、Vamana、HNSW、NSGのインデックス作成時間はそれぞれ129秒、219秒、480秒である。NSGのインデックス作成時間には、EFANN[3]による初期近傍グラフの作成時間が含まれる。
Recall-QPS曲線:
3.png
図3から、Vamanaは3つのデータセットで優れた性能を発揮し、NSGと同等、HNSWよりわずかに優れていることがわかる。
検索半径の比較
図2.cから、VamanaはNSGとHNSWと比較して、同じ想起率の下で、平均検索経路が最も短いことがわかる。
一度だけ構築されたインデックスと大規模な統合インデックスの比較
データセットSIFT1B
1回限りのインデックスパラメータL = 50, R = 128, alpha = 1.2。1800GのDDR3マシンで2日間実行した結果、ピークメモリは約1100G、平均アウト度は113.9となった。
マージに基づくインデックス作成手順:
- データセットに対してkmeansを用いて40クラスタを訓練する;
- 各点は最も近い2つのクラスタに分配される;
- 各クラスタについてL = 50, R = 64, alpha = 1.2のVamanaインデックスを構築する;
- 各クラスタのインデックスをマージする。
このインデックスは384GBのインデックスを生成し、平均アウトオブディグリーは92.1であった。このインデックスは64GB DDR4マシンで5日間実行された。
比較結果は以下の通りである(図2a):
4.png
結論として
- 一度だけ構築されたインデックスは、マージ・ベースのインデックスよりも有意に優れている;
- マージベースのインデックスも優れている;
- マージベースのインデックス作成スキームは、DEEP1Bデータセットにも適用可能である(図2b)。
ディスクベースのインデックス:DiskANN VS.FAISS VS.IVF-OADC+G+P
IVFOADC+G+Pは文献[5]で提案されたアルゴリズムである。
IVFOADC+G+PはFAISSよりも優れていることが参考文献[5]で証明されているため、本稿ではDiskANNとIVFOADC+G+Pのみを比較する。また、FAISSはGPUリソースを必要としますが、これはすべてのプラットフォームでサポートされているわけではありません。
IVF-OADC+G+PはHNSWとIVF-PQの組み合わせのようです。HNSWを用いてクラスタを決定し、ターゲットクラスタにいくつかの刈り込み戦略を追加して探索を行う。
その結果が図2aである。図中の16と32はコードブックサイズである。データセットはOPQで定量化されたSIFT1Bである。
コード実装の詳細
DiskANNのソースコードは https://github.com/microsoft/DiskANN でオープンソース化されている。
2021年1月、ディスクソリューションのソースコードがオープンソース化された。
以下では主にインデックス作成プロセスと検索プロセスについて紹介する。
インデックス構築
インデックス構築には8つのパラメータがある:
data_type: float/int8/uint8を含む。
data_file.bin:元データのバイナリファイル。ファイルの最初の2つの整数は,それぞれデータセット・ベクトルの総数nとベクトルの次元dimを表す.最後のndimsizeof(data_type)バイトは,連続ベクトルデータである.
index_prefix_path:出力ファイルのパスプレフィックス。インデックスが構築されると、複数のインデックス関連ファイルが生成される。このパラメータは、それらが格納されるディレクトリの共通接頭辞である。
R: グローバルインデックスの最大アウト次数。
L: VamanaインデックスのパラメータLで、候補集合サイズの上限。
B: 問い合わせ時のメモリ閾値。PQコードブックのサイズをGB単位で制御する。
M: インデックスを構築する際のメモリ閾値。フラグメントのサイズをGB単位で指定する。
T: スレッド数。
インデックス作成処理 (エントリ関数: aux_utils.cpp::build_disk_index):
- index_prefix_pathに従って様々な出力ファイル名を生成する。
- パラメータチェック。
- data_file.binのmetaを読み込んでnとdimを得る。Bとnに従ってPQのコードブック部分空間数mを決定。
- generate_pq_pivots:PQ学習セットの中心点をp = 1500000/nのサンプリングレートで一様にサンプリングし、PQを大域的に学習させる。
- generate_pq_data_from_pivots:グローバルなPQコードブックを生成し、中心点とコードブックを別々に保存する。
- build_merged_vamana_index: 元のデータセットをスライスし、Vamanaインデックスを分割して作成し、最後にインデックスを1つに統合する。
- partition_with_ram_budget:kmeansを用いてデータセットをサンプリングし、各ポイントを2つの最も近いクラスターに振り分ける。データセットを断片化し、各断片はデータファイルとIDファイルの2つのファイルを生成する。IDファイルとデータ・ファイルは互いに対応し、IDファイルの各IDはデータ・ファイルのベクトルに対応する。IDは元データの各ベクトルに0からn-1までの番号を振って得られる。IDは比較的重要であり、マージに関係する。
- 1500000/nのサンプリング・レートでトレーニング・セットを一様にサンプリングする;
- num_parts = 3 を初期化する:
- ステップiのトレーニングセットに対してnum_parts-means++を行う;
- サンプリング率0.01を使用してテスト集合を一様にサンプリングし,テスト集合を最も近い2つのクラスタに分割する.
- 各クラスタ内のポイント数を数え、それをサンプリング・レートで割って各クラスタ内のポイント数を推定する;
- ステップ3で最大のクラスタが必要とするメモリをVamanaインデックスサイズに従って推定し、それがパラメータMを超えない場合はステップiiiに進み、そうでない場合はnum_parts ++ ステップ2に戻る;
- 元のデータセットをnum_parts個のグループファイルに分割する。各グループファイルには、断片化されたデータファイルと、断片化されたデータに対応するIDファイルが含まれる。
- ステップaですべてのスライスに対して別々にVamanaインデックスを作成し、ディスクに保存する;
- merge_shards: num_partsシャードのVamanaをグローバルインデックスにマージする:
- num_partsフラグメントのIDファイルをidmapに読み込む。このidmapはfragment->idのフォワードマッピングを確立するのと同じである;
- idmapに従ってid->フラグメントの逆マッピングを確立し、各ベクトルがどの2つのフラグメントにあるかを知る;
- 1GBのキャッシュを持つリーダーを使ってnum_partsスライスのVamanaインデックスを開き、1GBのキャッシュを持つライターを使って出力ファイルを開く;
- Vamanaインデックスのnum_partsナビゲーションポイントを、検索時に使用するセンターポイントファイルに配置する;
- 小から大までIDに従ってマージを開始し、逆マッピングに従って各フラグメントの各オリジナルベクトルの近傍点セットを順番に読み込み、重複排除、シャッフル、切り捨て、出力ファイルへの書き込みを行う。スライスは元々大域的に順序付けされており、今回のマージも順序付けされているため、最終的にフラッシュされるインデックスのIDと元データのIDは一対一に対応する。
- フラグメントファイル、フラグメントインデックス、フラグメントIDファイルを含む一時ファイルを削除する。
create_disk_layoutを実行する:create_disk_layout:ステップ6で生成されたグローバルインデックスには、コンパクトな隣接表しかない。このステップはインデックスの整列を行う。隣接表と元のデータは一緒に格納される。検索時には、正確な距離計算のために、隣接表をロードし、元のベクトルを一緒に読み込む。SECTORという概念もあり、デフォルトのサイズは4096です。各SECTORは4096 / node_sizeのベクトル情報しか持ちません。 node_size = 単一ベクトルサイズ + 単一ノード隣接テーブルサイズです。
最後に、150000 / nのグローバル一様サンプリングを行って保存し、検索時のウォームアップに使用します。
検索
10個の検索パラメータがある:
- index_type:オプションにはFloat/int8/uint8があり、インデックスを構築する際の最初のパラメータdata_typeと同様である。
- index_prefix_path:インデックス・パラメータ index_prefix_path を参照。
- num_nodes_to_cache:キャッシュ・ホットスポットの数。
- num_threads:検索スレッド数。
- beamwidth: プリロードポイント数の上限。0に設定するとシステムが判断する。
- query_file.bin:クエリーセットファイル。
- truthset.bin:結果セットファイル。"null "は、結果セットが提供されていないことを意味し、プログラムが自分で計算する;
- K: topk;
- result_output_prefix:検索結果を保存するパス;
- L*:検索パラメーターリスト。複数の値を追加できる。各Lについて、異なるLで検索した場合の統計情報が与えられる。
検索プロセス
- 関連データの読み込み:クエリーセット、PQ中心点データ、コードブックデータ、検索開始点、その他のデータを読み込み、インデックスメタを読み込む。
- インデックス作成時にサンプリングされたデータセットを用いてcached_beam_searchを行い、各点のアクセス時間をカウントし、アクセス頻度の高いnum_nodes_to_cache点をキャッシュにロードする。
- デフォルトでWARMUP操作がある。ステップ2と同様に、このサンプルデータセットもcached_beam_searchに使用されます。
- 与えられたパラメータLの数に応じて、各Lはクエリセットを用いて再度cached_beam_searchが実行され、想起率やQPSなどの統計情報が出力されます。ウォームアップやホットスポットデータの統計処理はクエリ時間にカウントされません。
cached_beam_searchについて:
- 候補開始点からクエリ点に最も近い候補を探す。ここではPQ距離が使用され、始点は検索キューに追加される。
- 検索を開始する:
- 検索キューから、beam_width + 2以上の未訪問点がない。これらの点がキャッシュにあれば、キャッシュヒットキューに追加する。ヒットしなかった場合は、ミスキューに追加する。ミスキューのサイズがbeam_widthを超えないようにする。
- ミスキューのポイントに非同期ディスクアクセス要求を送る。
- キャッシュにヒットした点については、元データとクエリデータを使って正確な距離を計算し、結果キューに追加する。検索キューの長さはパラメータによって制限される。
- ステップ c と同様に、ステップ a でキャッシュされたミス点を処理する。
- 探索待ち行列が空になったら探索を終了し、結果待ち行列 topk を返す。
要約
これは比較的長い作品だが、全体的に優れている。論文とコードのアイデアは明快である。重なり合うバケットをk-meansによって分割し、バケットを分割してマップ・インデックスを作成し、最後にインデックスをマージするという、比較的新しいアイデアである。メモリベースのグラフインデックスVamanaに関しては、基本的にNSGのランダム初期化バージョンであり、トリミングの粒度を制御することができる。クエリー時には、キャッシュ+パイプラインをフル活用し、io時間の一部をカバーし、QPSを向上させる。しかし、この論文によると、マシンのコンディションが特別でなくても、トレーニングに5日もかかり、使い勝手は比較的悪い。今後、トレーニングの最適化が必要であることは間違いない。コードの観点からは、品質は比較的高く、本番環境でそのまま使用できる。
参考文献
- Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, Rohan Kadekodi.DiskANN: シングルノード上での高速高精度10億点最近傍探索.NeurIPS 2019.
- [コン・フー、チャオ・シアン、チャンシュー・ワン、デン・カイ。Navigating Spreading-outグラフによる高速近似最近傍探索.PVLDB, 12(5):461 - 474, 2019. doi: 10.14778/3303753.3303754.].(http://www.vldb.org/pvldb/vol12/p461-fu.pdf)
- コン・フー、デン・カイGitHub - ZJULearning/efanna: ANN検索とKNNグラフ構築のための高速ライブラリ.
- AI用検索エンジン:高维数据检索工业级解决方案
Like the article? Spread the word