DiskANN：10億スケールのデータセットで高い回収率とQPSを達成したディスクベースの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を搭載した1台のマシンを用いて、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つの主要な解決策を示す：

1.転置インデックス＋量子化：データセットをM個のパーティションにクラスタリングし、PQ（積量子化）などの量子化スキームを用いてデータセットを圧縮する。このソリューションでは、データ圧縮によって精度が低下するため、リコールが低くなる。topkを増加させることはリコールの改善に役立つが、QPSはそれに応じて低下する。

2.分割してインデックスを作成する：データセットをいくつかの不連続なシャードに分割し、それぞれのシャードに対してインメモリーインデックスを作成する。クエリー要求が来ると、各シャードのインデックスに対して検索が実行され、マージ後に結果が返される。このソリューションでは、データセットの規模が拡大しすぎるため、1台のマシンのメモリリソースが制限され、より多くのマシンが必要となり、QPSが低くなる。

上記のいずれの解決策も、1台のマシンのメモリ制限によって制限されている。本稿では、この問題を解決するためにSSD常駐型のインデックス作成メカニズムの設計を提案する。SSD常駐型インデクシングの課題は、ランダムディスクアクセス数とディスクアクセス要求数を削減することである。

投稿論文

本論文では，大規模データセットの検索を効果的にサポートするSSD常駐ANNS方式DiskANNを提案する．この方式は本稿で紹介するグラフベースのアルゴリズムに基づいている．Vamana。本論文の貢献は以下の通りである．

1.DiskANNは64GBのRAMを搭載した1台のマシンで、100次元を超える10億スケールのデータセットをインデックス化し、検索することができる。

2.Vamanaと呼ばれる新しいグラフベースのアルゴリズムは、NSGやHNSWよりも小さな探索半径を持ち、ディスクアクセス回数を最小化する。

3.Vamanaはメモリ上で動作可能であり、その性能はNSGやHNSWよりも低くはない。

4.大規模データセットの重複するパーティションに構築された小さなVamanaインデックスは、接続性を失うことなく1つのグラフにマージすることができる。

5.VamanaはPQのような量子化スキームと組み合わせることができる。グラフ構造と元のデータはディスクに保存され、圧縮されたデータはメモリに保存される。

Vamana

このアルゴリズムはNSG[2][4]の考え方に似ている（NSGを理解できない人は参考文献[2]を、論文を読みたくない人は参考文献[4]を参照されたい）。これらの主な違いはトリミング戦略にある。正確には、NSGのトリミング戦略にスイッチ・アルファが追加されている。NSGのトリミング戦略の主な考え方は、ターゲット点の近傍をできるだけ多様に選択することである。もし新しい近傍点が目標点よりも目標点の近傍点に近ければ、この点を近傍点セットに加える必要はない。言い換えると、ターゲット点の各近傍点に対して、周囲半径dist（ターゲット点、近傍点）内に他の近傍点が存在してはならない。このトリミング戦略は、グラフのアウトディグリーを効果的に制御し、比較的過激である。これはインデックスのメモリフットプリントを減らし、探索速度を向上させるが、探索精度も低下させる。Vamanaのトリミング戦略は、パラメータalphaによってトリミングの規模を自由に制御するものである。動作原理は、トリミング条件におけるdist（近傍点、候補点）にパラメータα（1以上）を乗算する。dist（目標点、ある候補点）が拡大された基準距離よりも大きい場合にのみ、トリミング戦略が採用され、目標点の近傍点間の相互排除の許容範囲が拡大される。

Vamanaのインデックス作成プロセスは比較的単純である：

1.ランダムなグラフを初期化する；

2.2.NSGのナビゲーションポイントと同じように、始点を計算する。まずグローバル重心を求め、次にグローバル重心に最も近い点をナビゲーション点として求める。VamanaとNSGの違いは、NSGの入力はすでに最近傍グラフであるため、ユーザーは単純に初期近傍グラフ上で直接セントロイド点の近似最近傍探索を行うことができる。しかし、Vamanaはランダムな最近傍グラフを初期化するため、ユーザーはランダムなグラフ上で直接近似探索を行うことができません。そのため、ユーザは直接ランダムグラフ上で近似探索を行うことができない。そのため、ユーザはグローバル比較を行い、後続の反復探索の開始点となるナビゲーション点を得る必要がある。この点の目的は平均探索半径を最小化することである；

3.3.初期化されたランダム近傍グラフとステップ2で決定された探索開始点に基づいて、各点に対して近似最近傍探索を実行し、探索経路上の全ての点を近傍候補セットとし、α=1を用いたエッジトリミング戦略を実行する。NSGと同様に、ナビゲーション点を始点とする探索経路上の点集合を候補近傍集合として選択することで、長いエッジが増え、探索半径を効果的に小さくすることができる。

4.α＞1に調整し（論文では1.2を推奨）、ステップ3を繰り返す。ステップ3がランダムな最近傍グラフに基づいているのに対し、最初の反復の後、グラフは低品質になります。従って、想起率に非常に重要なグラフ品質を向上させるために、もう1回反復する必要がある。

本稿では、Vamana、NSG、HNSWの3つのグラフインデックスを比較する。インデックス作成とクエリー性能の点では、VamanaとNSGは比較的近い。データについては以下の実験セクションを参照。

図1](https://assets.zilliz.com/2_906f6a4def.png)

Vamanaインデックスの構築プロセスを視覚化するために、この論文では200の2次元点を用いて2ラウンドの反復をシミュレートしたグラフを提供している。最初の行はα=1を使ってエッジをトリミングしている。このトリミング戦略は比較的急進的であり、多くのエッジがトリミングされていることがわかる。alphaの値を大きくし、トリミングの条件を緩くすると、明らかに多くのエッジが追加される。最終的なグラフでは、かなり長い辺が追加されている。これは効果的に探索半径を小さくすることができる。

DiskANN

64GBのメモリしかないパソコンでは、インデックスを構築するどころか、10億個の生データを保持することすらできない。そこには2つの課題がある：1.1.限られたメモリー資源で、このような大規模なデータセットにインデックスを付けるには？2.元データがメモリにロードできない場合、検索時の距離をどのように計算するか？

本論文では以下の解決策を提案した：

1.最初の課題：まず、k-meansを使ってデータをk個のクラスターに分割し、各点を最も近いi個のクラスターに割り当てる。各クラスタに対してメモリベースのVamanaインデックスを構築し、最後にk個のVamanaインデックスを1つに統合する。

2.2つ目の課題：元のベクトル上にインデックスを構築し、圧縮されたベクトルをクエリする。元のベクトルにインデックスを構築することで、グラフの品質を確保し、圧縮ベクトルはメモリにロードして粗視化検索を行うことができる。圧縮ベクトルで検索すると精度が落ちる可能性があるが、グラフの品質が十分高い限り、大まかな方向は正しい。最終的な距離結果は元のベクトルを使って計算される。

DiskANNのインデックスレイアウトは、一般的なグラフインデックスと似ています。各点の近傍集合と元のベクトルデータは一緒に保存される。これはデータの局所性をより良く利用するためである。

前述したように、インデックスデータがSSDに格納されている場合、検索遅延を少なくするために、ディスクアクセスの回数とディスクの読み書きの要求をできるだけ減らす必要があります。そこで DiskANN は 2 つの最適化戦略を提案します：

1.1. ホットスポットをキャッシュする：メモリ上の開始点から C ジャンプ以内のすべての点をキャッシュする。Cの値は3～4以内に設定するのがよい。

2.ビームサーチ：簡単に言えば、近傍情報をあらかじめ読み込んでおくことである。点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, alpha = 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曲線：

図2](https://assets.zilliz.com/3_dcdb9452ca.png)

図3からわかるように、Vamanaは3つのデータセットで優れた性能を発揮し、NSGと同程度、HNSWよりわずかに優れている。

検索半径の比較

図2.cから、VamanaはNSGとHNSWと比較して、同じ想起率の下で、平均探索経路が最も短いことがわかる。

一度だけ構築されたインデックスと大規模な統合インデックスの比較

データセットSIFT1B

一度だけ構築されたインデックスのパラメータ：L = 50, R = 128, alpha = 1.2。1800GのDDR3マシンで2日間実行した結果、ピークメモリは約1100G、平均アウト度は113.9。

マージに基づくインデックス作成手順：

1.kmeansを用いてデータセットに40クラスタを学習させる；

2.各ポイントは最も近い2つのクラスターに分配される；

3.3. 各クラスタについて、L = 50, R = 64, α = 1.2のVamanaインデックスを構築する；

4.各クラスターのインデックスをマージする。

このインデックスは384GBのインデックスを生成し、平均アウトオブディグリーは92.1であった。このインデックスは64GB DDR4マシンで5日間実行された。

比較結果は以下の通りである（図2a）： 図3.](https://assets.zilliz.com/4_ea421b98c3.png)

結論として

1.一度だけ構築されたインデックスは、マージ・ベースのインデックスよりも有意に優れている；

2.マージベースのインデックスも優れている；

3.マージに基づくインデックス作成スキームは、DEEP1Bデータセットにも適用可能である（図2b）。

ディスクベースのインデックス：DiskANN VS.FAISS VS.IVF-OADC+G+P

IVFOADC+G+Pは文献[5]で提案されたアルゴリズムである。

参考文献[5]ではIVFOADC+G+PがFAISSより優れていることが証明されているため、本稿では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を表す．最後のn * dim * sizeof(data_type)バイトは，連続ベクトルデータである．

index_prefix_path：出力ファイルのパスプレフィックス。インデックスが構築されると、複数のインデックス関連ファイルが生成される。このパラメータは、それらが格納されるディレクトリの共通接頭辞である。

R: グローバルインデックスの最大アウト次数。

L: Vamana インデックスのパラメータ L で、候補集合サイズの上限。

B: 問い合わせ時のメモリ閾値。PQコードブックのサイズをGB単位で制御する。

M: インデックスを構築する際のメモリ閾値。GB単位でフラグメントのサイズを決定する。

T: スレッド数。

インデックス作成処理(エントリ関数: aux_utils.cpp::build_disk_index)：

1.index_prefix_pathに従って様々な出力ファイル名を生成する。

2.パラメータチェック。

3.data_file.binのmetaを読み込んでnとdimを得る。Bとnに従ってPQのコードブック部分空間数mを決定する。

4. generate_pq_pivots：PQ学習セットの中心点をp = 1500000/nのサンプリングレートで一様にサンプリングし、PQをグローバルに学習する。

5. generate_pq_data_from_pivots：グローバルなPQコードブックを生成し、中心点とコードブックを別々に保存する。

6. 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 を超えない場合はステップ 3 に進み、そうでない場合は 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ファイルを含む一時ファイルを削除する。

7.create_disk_layoutを作成する：ステップ6で生成されたグローバルインデックスは、コンパクトな隣接表しか持っていない。このステップでは、インデックスの位置合わせを行う。隣接表と元のデータは一緒に格納される。検索時には、正確な距離計算のために、隣接表をロードし、元のベクトルを一緒に読み込む。SECTORという概念もあり、デフォルトのサイズは4096です。各SECTORは4096 / node_sizeのベクトル情報しか持ちません。node_size = 単一ベクトル・サイズ + 単一ノード隣接テーブル・サイズです。

8.最後に、150000 / nのグローバル一様サンプリングを行って保存し、検索時のウォームアップに使用する。

検索

10個の検索パラメータがある：

• index_type：インデックスを作成する際の最初のパラメータ 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について、異なるLで検索した場合の統計情報が与えられる。

検索プロセス

1.関連データの読み込み：クエリーセット、PQ中心点データ、コードブックデータ、検索開始点などのデータを読み込み、インデックスメタを読み込む。

2.2.インデックス作成中にサンプリングされたデータセットを使用して、cached_beam_searchを行い、各点のアクセス時間をカウントし、アクセス頻度の高いnum_nodes_to_cache点をキャッシュにロードする。

3.デフォルトでWARMUP操作がある。ステップ2と同様に、このサンプル・データセットもcached_beam_searchに使用されます。

4.4.与えられたパラメータLの数に応じて、各Lはクエリーセットを用いて再度cached_beam_searchが実行され、再現率やQPSなどの統計が出力される。ウォームアップやホットスポットデータの統計処理は、クエリ時間にはカウントされない。

cached_beam_searchについて：

1.始点候補からクエリーポイントに最も近い候補を探す。ここではPQ距離が使用され、始点が検索キューに追加される。

2.検索を開始する：

• 検索キューから、beam_width + 2 未訪問点以下。これらの点がキャッシュにあれば、キャッシュヒットキューに追加する。ヒットしなかった場合は、ミスキューに追加する。ミスキューのサイズがbeam_widthを超えないようにする。

• ミスキューのポイントに非同期ディスクアクセス要求を送る。

• キャッシュにヒットした点については、元のデータとクエリーデータを使用して正確な距離を計算し、結果キューに追加する。検索キューの長さはパラメータによって制限される。

• ステップ c と同様に、ステップ a でキャッシュされたミスポイントを処理する。

• 探索待ち行列が空になったら探索を終了し、結果待ち行列 topk を返す。

まとめ

これは比較的長い作品であるが、全体的に優れている。論文とコードのアイデアは明快である。k-meansを通して重なり合う多数のバケットを分割し、バケットを分割してマップ・インデックスを構築し、最後にインデックスをマージするという、比較的新しいアイデアである。メモリベースのグラフインデックスVamanaに関しては、基本的にNSGのランダム初期化バージョンであり、トリミングの粒度を制御することができる。クエリー時には、キャッシュ＋パイプラインをフル活用し、io時間の一部をカバーし、QPSを向上させる。しかし、この論文によると、マシンのコンディションが特別でなくても、トレーニングに5日もかかり、使い勝手は比較的悪い。今後、トレーニングの最適化が必要であることは間違いない。コードの観点からは、品質は比較的高く、本番環境でそのまま使用できる。

参考文献

1.Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, Rohan Kadekodi.DiskANN: シングルノード上での高速高精度10億点最近傍探索.NeurIPS 2019.

2.[Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai.Navigating spreading-out graphs による高速近似最近傍探索．PVLDB, 12(5):461 - 474, 2019. doi: 10.14778/3303753.3303754.].(http://www.vldb.org/pvldb/vol12/p461-fu.pdf)

3.Cong Fu と Deng Cai.GitHub - ZJULearning/efanna: ANN検索とKNNグラフ構築のための高速ライブラリ

4.AI用検索エンジン：高维数据检索工业级解决方案

5. Dmitry Baranchuk, Artem Babenko, and Yury Malkov.億スケールの近似最近傍探索のための転置インデックスの再検討