この記事は、 NTT Communications Advent Calendar 2024 22日目の記事です。
こんにちは、イノベーションセンターの鈴ヶ嶺です。普段はAI/MLシステムに関する業務に従事しています。
本記事では、LLMのRetrieval Augmented Generation(RAG)などで用いられる近似最近傍探索(ANN)手法の1つとして最近注目されているDiskANNについて紹介します。 DiskANNはbillionクラスの大規模なデータセットに対してスケーラブルに機能し、SSDなどのDiskを利用可能なことからコスト効率の良い点が特徴です。 本記事ではまず、グラフベースの近似最近傍探索と従来手法の課題点について説明します。次にDiskANNについて、中核となるアルゴリズムであるVamana AlgorithmやDiskへのデータ保存の詳細、検索方法、利用方法について説明します。さらに、DiskANNベースの手法を採用しているクラウド・ベクトルデータベースベンダーの対応状況を紹介します。
グラフベースの近似最近傍探索について
近似最近傍探索は与えられた検索対象のクエリから、類似される(距離が近い)データを効率的に見つけることが目的です。 その中で、グラフベース手法はデータセットをノードとエッジにより構成されたグラフで表現して探索するものになります。 例えばグラフベース手法で有名なHierarchical Navigable Small World(HNSW)1は、次のように階層的なグラフ構造となります。 上層では疎なグラフを探索して、徐々に階層を重ねることでより密なグラフの探索をします。
このような階層性や探索におけるランダムアクセス性の要件からHNSWのグラフはin-memoryに保存することが望まれます。 一方で大規模なデータセットをメモリに展開するのはコストが高く、利用が難しい状況にあります。
また、その他グラフベース手法の詳細については説明を割愛しますが、以下の松井先生の資料などがわかりやすいため補足としてご覧いただければと思います。
DiskANN
HNSWなどのグラフベース手法の課題を背景に、Microsoft Researchの研究チームはbillionクラスの大規模なデータに対してスケーラブルに機能し、さらに低コストなSSDなどのDiskを用いることが可能な手法であるDiskANNを考案しました。 この手法は2019年のNeural Information Processing Systems(NeurIPS)で採択されています。2
実装は、オープンソースとして公開されています。
https://github.com/microsoft/DiskANN
コアはC++で実装されており、Pythonのdiskannpyからも利用可能です。
一部Rustによる実装も進められていますが、コミッターの回答からは明確なロードマップがあるわけではないようです。
We do not have a strong roadmap yet.
Vamana Algorithm
このセクションではDiskANNのグラフ作成アルゴリズムであるVamana Algorithmについて説明します。 Vamana Algorithmはグラフの走査を高速化するためのもので、そこには2つのアルゴリズムが寄与しています。
Vamanaのグラフはフラットな単一のレイヤーからなるグラフで構成されます。 最初にランダム接続された初期グラフを改善することによってクエリを探索可能なグラフとして構成します。
Algorithm 1のGreedySearchでは、入力されたクエリにより距離が近いノードを探索します。 始点ノードであるsから、クエリx_qにより近い隣接ノードを探索して候補に加えて、さらにその候補から隣接ノードを探索するというような形で、よりクエリx_qとの距離を縮める方向で貪欲的に探索を進めます。 探索中に候補数がL以上になった場合、上位Lのノードだけ残して探索を進めます。
この処理は、後に説明するAlgorithm 3のグラフ構築時にも、クエリの検索時にも利用されます。
Algorithm 2のRobustPruneではノードのエッジの刈り込みを行いグラフ探索の効率を向上させます。 まず、特定のノードpを起点として、候補ノードの中から最も距離が近いノードp＊を選択し、pとp*の間にエッジを新規に追加します。 次に候補から任意のノードp'において、追加されたノードp＊からp'の距離をパラメータαの乗数で制御して、特定のノードpからp'の距離よりも小さい場合、p'を削除します。 この処理で、ある程度疎なグラフとなり、長距離のエッジも形成されます。 最終的にノードあたりのエッジ数はR個になるように調整されます。
上図に、RobustPruneにおけるパラメータαによる違いを例示します。 α=1の場合、ノードp'は削除されます。一方、α=2の場合、削除されず長距離のエッジ追加に貢献します。 ちなみに、HNSWやNSG3では、この調整可能なパラメータαを持っておらず、暗黙的にα=1を利用しているため長距離のエッジが作成されにくい傾向にあります。
グラフ構築のインデックス処理はAlgorithm 1, 2を組みわせて行われます。 まず、ランダムに設定されたノードをデータセットのmedoid（全てのノードからの距離の合計が最小）からGreedySearchを行い探索で訪れたノードを記録します。 それらのノードは長距離のエッジ構築を促進する目的で、次のRobustPruneの候補ノードとして活用され、エッジの刈り込みがされます。 これらの処理を行いエッジが刈り込まれて、探索の効率が向上します。
次の画像は、実際にエッジが刈り込まれる様子を図示されたものになります。 1回目にα=1として処理され、次の2回目に長距離エッジを導入するためにユーザ定義のα>1で処理されます。
Diskへのインデックス
billionスケール規模のデータセットにおいて、例えばfloat32の100次元データとして考えるとインデックス構築時には4Byte x 100 x 1000000000 = 約373GBほど必要となります。 先述の通り、これをメモリに展開することはコストが高く、小規模な計算環境では現実的ではありません。 DiskANNではクラスタリングによる分割によってこの問題の解決を図ります。 k-meansを利用してk個(例 k=40)のクラスタに分割して各データセット点はl個(例 l=2)のクラスタに属するように設定します。 次に、各クラスタに割り当てられたデータセットでVamamaインデックスを構築します(データ数はNl/kとなり、例 4Byte x 100 x 1000000000 x 2 / 40 = 約 18.6GB)。 最終的に各クラスタで構築したグラフのエッジを単純に結合させて1つのグラフにマージします。
最終的に作成したグラフを全てメモリに載せることは難しいので、DiskANNではProduct Quantization4を利用した圧縮ベクトルをメモリに格納して、フル精度のグラフはSSDなどのDiskに保存します。Diskへの保存のレイアウトは各ノードiについてフル精度のベクトルとR個の近傍ノードのIDを記録します。その際にノードの次数がRより小さい場合は0埋めしてオフセットの計算を容易にします。
クエリ検索時には圧縮ベクトルを用いてまず大まかな計算をして、実際の距離を計算してランク付けする処理はDiskにあるフル精度のものを利用します。検索の様子は次のMicrosoftが公開している動画が非常にわかりやすく表現されているのでおすすめです。
派生手法
DiskANNを提案したMicrosoft Researchの研究チームはDiskANNを改良した派生を活発に研究しています。
FreshDiskANN
DiskANNはランダムな初期グラフから刈り込みをすることでグラフを生成するため、新しいデータを追加して再インデックス化するコストがかかっていました。 その課題に対してFreshDiskANN 5では、リアルタイムにデータを更新してインデックスに反映可能な手法を提案しています。
Filtered-DiskANN
Filtered-DiskANN 6は日付、価格帯、言語などでラベルによってフィルタリングされたインデックス内の探索を効率化した手法を提案しています。
利用方法(diskannpy)
DiskANNのPythonでの実装、diskannpyの利用方法は次のようになります。
sample_diskann.py
import numpy as np import diskannpy as dap from pathlib import Path rng = np.random.default_rng(1234) rows = 512 dim = 100 data = rng.random((rows, dim), dtype=np.float32) index_directory = "index" Path(index_directory).mkdir(exist_ok=True) # Build Disk Index dap.build_disk_index( data=data, distance_metric="l2", index_directory=index_directory, graph_degree=16, complexity=32, vector_dtype=np.float32, search_memory_maximum=0.00003, build_memory_maximum=1, num_threads=0, pq_disk_bytes=0 ) index = dap.StaticDiskIndex( distance_metric="l2", vector_dtype=np.float32, index_directory=index_directory, num_threads=16, num_nodes_to_cache=10, ) # Search queries = rng.random((5, dim), dtype=np.float32) identifiers, distances = index.batch_search(queries, k_neighbors=5, complexity=5, beam_width=2, num_threads=1)
実行方法は次のようになり、indexディレクトリにメタデータやインデックスが保存されます。
python sample_diskann.py tree index/ # Diskに保存されたメタデータやインデックス # index/ # ├── ann_disk.index # ├── ann_metadata.bin # ├── ann_pq_compressed.bin # ├── ann_pq_pivots.bin # ├── ann_sample_data.bin # ├── ann_sample_ids.bin # └── ann_vectors.bin
クラウド・ベクトルデータベースベンダーの対応状況
このセクションではDiskANNベースのアプローチをとり実装している各種クラウド事業者やいくつかのベクトルデータベースベンダーの対応状況を記載します。
Azure
Microsoftで考案された手法であるため、クラウドにおいてはAzureでの利用が先行している印象です。今年2024年のMicrosoft Build 2024やMicrosoft Ignite 2024でAzure Database for PostgreSQL - Flexible ServerとAzure Cosmos DBの対応が発表されました。
DiskANN Vector Index in Azure Database for PostgreSQL
Azure Database for PostgreSQL(Flexible Server)において pg_diskann 拡張を有効化することで利用可能となりました。
次のような例でDiskANNを使用可能です。
-- pg_diskannのextension有効化 CREATE EXTENSION IF NOT EXISTS pg_diskann CASCADE; -- インデックスをdiskannを利用して作成 CREATE INDEX my_table_embedding_diskann_idx ON my_table USING diskann (embedding vector_cosine_ops);
DiskANN for Azure Cosmos DB
引用: https://devblogs.microsoft.com/cosmosdb/diskann-for-azure-cosmos-db-now-in-open-public-preview/
Azure Cosmos DBでもDiskANNが利用可能です。Cosmos DBの水平スケーリング機能と相乗効果で多くのクエリをコスト効率を最適化しながら処理することが可能となります。
次のように EnableNoSQLVectorSearch を設定することで利用可能です。
az cosmosdb update \
--resource-group <resource-group-name> \
--name <account-name> \
--capabilities EnableNoSQLVectorSearch
Timescale pgvectorscale
Timescaleはpostgresql拡張として pgvectorscale をOSSとして公開しています。 内部ではStreamingDiskANNというDiskANNにインスパイアされたアルゴリズムを使用しています。
余談ですがpgvectorscaleはRustでPostgresql拡張を実装するPGRX frameworkを用いてRustで実装されています。
TimescaleはいくつかDiskANNに関する技術ブログを公開しているため参考におすすめです。
- https://www.timescale.com/learn/hnsw-vs-diskann
- https://www.timescale.com/blog/understanding-diskann
Zilliz Milvus
Zillizのベクトルデータベース Milvus はDiskANNに対応しています。インデックスタイプをDISKANNと設定することで利用可能です。
https://milvus.io/docs/ja/disk_index.md
ZillizはDiskANNに関する技術ブログを多数公開しています。
- https://milvus.io/blog/2021-09-24-diskann.md
- https://zilliz.com/learn/DiskANN-and-the-Vamana-Algorithm
- https://zilliz.com/blog/diskann-a-disk-based-anns-solution-with-high-recall-and-high-qps-on-billion-scale-dataset
Weaviate
ベクトルデータベース Weaviate はVersion 1.18からDiskANNベースの処理を組み合わせています。
More specifically, in 1.18 we implemented the solution proposed in DiskANN whereby we combined our HNSW indexing with vectors that have been compressed using product quantization (PQ).
WeaviateもDiskANNに関する技術ブログを公開しています。HNSWとのRecallやLatencyなども比較検証結果を公開しています。
まとめ
本記事ではDisk-friendlyで低コストにスケール可能な近似最近傍探索手法 DiskANNのグラフ構築アルゴリズムVamanaやその利用方法を紹介しました。また、各クラウド・ベクトルデータベースベンダーの対応状況を紹介し低コストで社会実装を視野に入れた取り組みが活発化している印象を感じています。大規模な検索システムを検討している方にこちらの内容を参考にしていただければ幸いです。
明日のアドベントカレンダーもお楽しみに。
- Malkov, Yu A., and Dmitry A. Yashunin. "Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs." IEEE transactions on pattern analysis and machine intelligence 42.4 (2018): 824-836.↩
- Jayaram Subramanya, Suhas, et al. "Diskann: Fast accurate billion-point nearest neighbor search on a single node." Advances in Neural Information Processing Systems 32 (2019).↩
- Fu, Cong, et al. "Fast approximate nearest neighbor search with the navigating spreading-out graph." arXiv preprint arXiv:1707.00143 (2017).↩
- Jegou, Herve, Matthijs Douze, and Cordelia Schmid. "Product quantization for nearest neighbor search." IEEE transactions on pattern analysis and machine intelligence 33.1 (2010): 117-128.↩
- Singh, Aditi, et al. "Freshdiskann: A fast and accurate graph-based ann index for streaming similarity search." arXiv preprint arXiv:2105.09613 (2021).↩
- Gollapudi, Siddharth, et al. "Filtered-diskann: Graph algorithms for approximate nearest neighbor search with filters." Proceedings of the ACM Web Conference 2023. 2023.↩
