DiskANN： 十億規模數據集上高召回高 QPS 的 ANNS 單機方案 | Paper Reading

摘要

「DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node」 [1] 是 2019 年發表在 NeurIPS 上的論文。該文提出了一種基於磁盤的 ANN 方案，該方案能夠在單個 64 G 內存和足夠 SSD 的機器上對十億級別的數據進行索引、存儲和查詢， 而且可以知足大規模數據 ANNS 的三個需求： 高召回、低查詢時延和高密度（單節點能索引的點的數量）。該文提出的方法作到了在 16 核 64G 內存的機器上對十億級別的數據集 SIFT1B 建基於磁盤的圖索引，而且 recall@1 > 95% 的狀況下 qps 達到了 5000， 平均時延不到 3ms。

論文做者

Suhas Jayaram Subramanya： 前微軟印度研究院員工，CMU 在讀博士。主要研究方向有高性能計算和麪向大規模數據的機器學習算法。

Devvrit：Graduate Research Assistant at The University of Texas at Austin。研究方向是理論計算機科學、機器學習和深度學習。

Rohan Kadekodi：德克薩斯大學的博士研究生。研究方向是系統和存儲，主要包括持久化存儲、文件系統和 kv 存儲領域。

Ravishankar Krishaswamy：微軟印度研究院 Principal Researcher。 CMU 博士學位。 研究方向是基於圖和聚類的近似算法。

Harsha Vardhan Simhadri：微軟印度研究員 Principal Researcher。CMU 博士學位。之前研究並行算法和運行時系統，如今主要工做是開發新算法，編寫編程模型。

論文動機

目前有不少向量檢索的 ANN 算法，這些算法在建索引性能、查詢性能和查詢召回率方面各有取捨。當前在查詢時間和召回率上表現較好的是基於圖的索引如 HNSW 和 NSG。因爲圖索引內存佔用比較大，在單機內存受限的狀況下，常駐內存的方案能處理點集的規模就十分有限。

許多應用須要快速在十億級別數據規模上作基於歐幾里得距離的近似查詢，目前有兩種主流的方法：

1. 基於 倒排表+量化 的方法。缺點是召回率不高，由於量化會產生偏差。雖然能夠提升 topk 以改善召回率，可是相應的會下降 qps。
2. 基於 分治。將數據集分紅若干個不相交的子集，每一個子集建基於內存的索引，最後對結果作歸併。這種方法比較耗內存和機器。例如對 100M, 128 維的 float 數據集上建 NSG 索引，假設出度上限爲 50， 大約須要 75G 內存。所以處理十億級別的數據就須要多臺機器。在阿里巴巴淘寶的實際應用中，將 20 億 128 維浮點數據分到 32 個分片中分別建 NSG 索引，recall@100 爲 98% 的時延大約在 5ms。當數據規模增長到千億級別時須要數千臺機器。^[2]

以上兩種方法的侷限性在於太依賴內存。因此這篇論文考慮設計一種索引常駐 SSD 的方案。索引常駐 SSD 的方案主要面臨的挑戰是如何減小隨機訪問 SSD 的次數和減小發起 SSD 訪問請求的數量。將傳統的基於內存的 ANNS 算法放到 SSD 上的話平均單條查詢會產生數百個讀磁盤操做，這會致使極高的時延。

論文貢獻

這篇論文提出了可以有效支持大規模數據的常駐 SSD 的 ANNS 方案： DiskANN。該方案基於這篇論文中提出的另外一個基於圖的索引： Vamana，後面會詳細介紹。這篇論文的主要貢獻包括但不限於：

1. DiskANN 能夠在一臺 64G 內存的機器上對十億級別的維度大於 100 的數據集進行索引構建和提供查詢服務，並在單條查詢 recall@1 > 95% 的狀況下平均時延不超過 5ms。
2. 提出了基於圖的新索引 Vamana，該索引相比目前最早進的 NSG 和 HNSW 具備更小的搜索半徑，這個性質能夠最小化 DiskANN 的磁盤訪問次數。
3. Vamana 搜索性能不慢於目前最好的圖索引 NSG 和 HNSW。
4. DiskANN 方案經過將大數據集分紅若干個相交的分片，而後對每一個分片建基於內存的圖索引 Vamana，最後將全部分片的索引合併成一個大索引，解決了內存受限的狀況下對大數據集創建索引的問題。
5. Vamana 能夠和現有的量化方法如 PQ 結合，量化數據能夠緩存在內存中，索引數據和向量數據能夠放在 SSD 上。

Vamana

這個算法和 NSG^[2] ^[4]思路比較像（不瞭解 NSG 的能夠看參考文獻 2，不想讀 paper 的話能夠看參考文獻 4），主要區別在於裁邊策略。準確的說是給 NSG 的裁邊策略上加了一個開關 alpha。NSG 的裁邊策略主要思路是：對於目標點鄰居的選擇儘量多樣化，若是新鄰居相比目標點，更靠近目標點的某個鄰居，咱們能夠沒必要將這個點加入鄰居點集中。也就是說，對於目標點的每一個鄰居節點，周圍方圓 dist(目標點，鄰居點）範圍內不能有其餘鄰居點。這個裁邊策略有效控制了圖的出度，而且比較激進，因此減小了索引的內存佔用，提升了搜索速度，但同時也下降了搜索精度。Vamana 的裁邊策略其實就是經過參數 alpha 自由控制裁邊的尺度。具體做用原理是給裁邊條件中的 dist(某個鄰居點，候選點） 乘上一個不小於 1 的參數 alpha，當 dist(目標點，某個候選點）大於這個被放大了的參考距離後才選擇裁邊，增長了目標點的鄰居點之間的互斥容忍度。

Vamana 的建索引過程比較簡單：

1. 初始化一張隨機圖；
2. 計算起點，和 NSG 的導航點相似，先求全局質心，而後求全局離質心最近的點做爲導航點。和 NSG 的區別在於：NSG 的輸入已是一張近鄰圖了，因此直接在初始近鄰圖上對質心點作一次近似最近鄰搜索就能夠了。可是 Vamana 初始化是一張隨機近鄰圖，因此不能在隨機圖上直接作近似搜索，須要全局比對，獲得一個導航點，這個點做爲後續迭代的起始點，目的是儘可能減小平均搜索半徑；
3. 基於初始化的隨機近鄰圖和步驟 2 中肯定的搜索起點對每一個點作 ANN，將搜索路徑上全部的點做爲候選鄰居集，執行 alpha = 1 的裁邊策略。這裏和 NSG 同樣，選擇從導航點出發的搜索路徑上的點集做爲候選鄰居集會增長一些長邊，有效減小搜索半徑。
4. 調整 alpha > 1(論文推薦 1.2）重複步驟 3。由於 3 是基於隨機近鄰圖作的，第一次迭代後圖的質量不高，因此須要再迭代一次來提高圖的質量，這個對召回率很重要。

論文比較了 Vamana、NSG、HNSW 3種圖索引，不管是建索引性能仍是查詢性能， Vamana 和 NSG 都比較接近，而且都稍好於 HNSW。具體數據能夠看下文實驗部分。

爲了直觀地表現建 Vamana 索引過程，論文中給出了這麼一張圖，用 200 個二維點模擬了兩輪迭代過程。第一行是用 alpha = 1 來裁邊，能夠發現改裁邊策略比較激進，大量的邊被裁剪。通過放大 alpha，裁邊條件放鬆後，明顯加回來了很多邊，而且第二行最右這張圖，即最終的圖中，明顯加了很多長邊。這樣能夠有效減小搜索半徑。

DiskAnn

一臺只有 64G 內存的我的電腦連十億原始數據都放不下，更別說建索引了。擺在咱們面前的有兩個問題：1. 如何用這麼小的內存對這麼大規模的數據集建索引？2. 若是原始數據內存放不下如何在搜索時計算距離？

本文提出的方法：

1. 對於第一個問題，先作全局 kmeans，將數據分紅 k 個簇，而後將每一個點分到距離最近的 I 個簇中，通常 I 取 2 就夠了。對每一個簇建基於內存的 Vamana 索引，最後將 k 個 Vamana 索引合併成一個索引。
2. 對於第二個問題，可使用量化的方法，建索引時用原始向量，查詢的時候用壓縮向量。由於建索引使用原始向量保證圖的質量，搜索的時候使用內存能夠 hold 住的壓縮向量進行粗粒度搜索，這時的壓縮向量雖然有精度損失，可是隻要圖的質量足夠高，大方向上是對的就能夠了，最後的距離結果仍是用原始向量作計算的。

DiskANN 的索引布局和通常的圖索引相似，每一個點的鄰居集和原始向量數據存在一塊兒。這樣作的好處是能夠利用數據的局部性。

前面提到了，若是索引數據放在 SSD 上，爲了保證搜索時延，儘量減小磁盤訪問次數和減小磁盤讀寫請求。所以 DiskANN 提出兩種優化策略：

1. 緩存熱點：將起點開始 C 跳內的點常駐內存，C 取 3~4 就比較好。
2. beam search： 簡單的說就是預加載，搜索 p 點時，若是 p 的鄰居點不在緩存中，須要從磁盤加載 p 點的鄰居點。因爲一次少許的 SSD 隨機訪問操做和一次 SSD 單扇區訪問操做耗時差很少，因此咱們能夠一次加載 W 個未訪問點的鄰居信息，W 不能過大也不能太小，過大會浪費計算和 SSD 帶寬，過小了也不行，會增長搜索時延。

實驗

實驗分三組：

1、 基於內存的索引比較： 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 repository ^[3]裏列出的參數中選擇出性能比較好的那組，Vamana 和 NSG 比較接近，所以參數也比較接近，可是沒有給出 HNSW 的參數選擇理由。在筆者看來，HNSW 的參數 M 選擇偏大，同爲圖索引，出度應該也要在同一水平才能更好作對比。

在上述建索引參數下，Vamana、HNSW 和 NSG 建索引時間分別爲 129s、219s 和 480s。NSG 建索引時間包括了用 EFANN [3] 構建初始化近鄰圖的時間。

召回率-qps 曲線：

從 Figure 3 能夠看出，Vamana 在三個數據集上都有着優秀的表現，和 NSG 比較接近，比 HNSW 稍好。

比較搜索半徑：

這個結果能夠看 Figure 2.c，從圖中能夠看出 Vamana 相比 NSG 和 HNSW，在相同召回率下平均搜索路徑最短。

2、 比較一次性建成的索引和多個小索引合併成一個大索引的區別

數據集： SIFT1B

一次建成的索引參數：L = 50， R = 128， alpha = 1.2. 在 1800G DDR3 的機器上跑了 2 天， 內存峯值大約 1100 G，平均出度 113.9。

基於合併的索引步驟：

1. 將數據集用 kmeans 訓練出 40 個簇；
2. 每一個點分到最近的 2 個簇；
3. 對每一個簇建 L = 50， R = 64， alpha = 1.2 的 Vamana 索引；
4. 合併每一個簇的索引。

這個索引生成了一個 384G 的索引，平均出度 92.1。這個索引在 64G DDR4 的機器上跑了 5 天。

比較結果以下（Figure 2a）：

結論：

1. 一次建成的索引顯著優於基於合併的索引；
2. 基於合併的索引也很優秀；
3. 基於合併的索引方案也適用於 DEEP1B 數據集（Figure 2b）。

3、 基於磁盤的索引： DiskANN VS. FAISS VS. IVF-OADC+G+P

IVFOADC+G+P 是參考文獻 [5] 提出的一種算法。

這篇論文只和 IVFOADC+G+P 比較，由於參考文獻 [5] 中已經證實比 FAISS 要更優秀，另外 FAISS 要用 GPU，並非全部平臺都支持。

IVF-OADC+G+P 好像是一個 hnsw + ivfpq。用 hnsw 肯定 cluster，而後在目標 cluster 上加上一些剪枝策略進行搜索。

結果在 Figure 2a 裏。圖裏的 16 和 32 是碼本大小。數據集是 SIFT1B，用的 OPQ 量化。

代碼實現細節

DiskANN 的源碼已經開源：[GitHub 地址][https://github.com/microsoft/DiskANN]

2021 年 1 月開源了磁盤方案的源碼，稍微看了下，大概介紹一下磁盤方案的實現細節。

這個代碼質量比較高，可讀性比較強，建議有能力的能夠讀一下。

下面主要介紹建索引過程和搜索過程。

建索引

建索引參數有 8 個：

data_type: float/int8/uint8 三選一

data_file.bin: 原始數據二進制文件，文件前 2 個 int 分別表示數據集向量總數 n 和向量維度 dim。後 n * dim * sizeof(data_type) 字節就是連續的向量數據。

index_prefix_path: 輸出文件的路徑前綴，索引建完後會生成若干個索引相關文件，這個參數是公共前綴。

R: 全局索引的最大出度。

L: Vamana 索引的參數 L，候選集大小上界。

B: 查詢時的內存閾值，單位 GB，控制 pq 碼本大小。

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: 用 p = 1500000/n 的採樣率全局均勻採樣出 pq 訓練集訓練 pq 中心點。

5. generate_pq_data_from_pivots: 生成全局 pq 碼本，中心點和碼本分別保存。

6. build_merged_vamana_index: 對原始數據集切片，分段建 Vamana 索引，最後合併。

   • 6.1： partition_with_ram_budget: 根據參數 M 肯定分片數 k。對數據集採樣作 kmeans，每一個點分到最近的兩個簇中，對數據集進行分片，每一個分片產生兩個文件：數據文件和 id 文件。id 文件和數據文件一一對應，id 文件中每一個 id 對應數據文件中每條向量一一對應。這裏的 id 能夠認爲是對原始數據的每條向量按 0 ～ n-1 編號。這個 id 比較重要，跟後面的合併相關。

     • 6.1.1： 用 1500000 / n 的採樣率全局均勻抽樣出訓練集
     • 6.1.2： 初始化 num_parts = 3。從 3 開始迭代。
       • 6.1.2.1：對步驟 6.1.1 中的訓練集作 num_parts-means++；
       • 6.1.2.2：用 0.01 的採樣率全局均勻採樣出一個測試集，將測試集中的分到最近的 2 個簇中；
       • 6.1.2.3：統計每一個簇中的點數，除以採樣率估算每一個簇的點數；
       • 6.1.2.4：按 Vamana 索引大小估算步驟 6.1.2.3 中最大的簇須要的內存，若是不超過參數 M，轉步驟 6.1.3，不然 num_parts ++ 轉步驟 6.1.2；
     • 6.1.3： 將原始數據集分紅 num_parts 組文件，每組文件包括分片數據文件和分片數據對應的 id 文件。

   • 6.2： 對步驟 6.1 中的全部分片單獨創建 Vamana 索引並存盤；

   • 6.3： merge_shards: 將 num_parts 個分片 Vamana 合併成一個全局索引。

     • 6.3.1： 讀取 num_parts 個分片的 id 文件到 idmap 中，這個 idmap 至關於創建了分片->id 的正向映射；

     • 6.3.2：根據 idmap 創建 id-> 分片的反向映射，知道每條向量在哪兩個分片中。

     • 6.3.3：用帶 1G 緩存的 reader 打開 num_parts 個分片 Vamana 索引，用帶 1G 緩存的 writer 打開輸出文件，準備合併；

     • 6.3.4：將 num_parts 個 Vamana 索引的導航點落盤到中心點文件中，搜索的時候會用到；

     • 6.3.5：按 id 從小到大開始合併，根據反向映射依次讀取每條原始向量在各個分片的鄰居點集，去重，shuffle，截斷，寫入輸出文件。由於當初切片時是全局有序的，如今合併也按順序來，因此最終的落盤索引中的 id 和原始數據的 id 是一一對應的。

     • 6.3.6：刪除臨時文件，包括分片文件、分片索引、分片 id 文件；

7. create_disk_layout: 步驟 6 中生成的全局索引只有鄰接表，並且仍是緊湊的鄰接表，這一步就是把索引對齊，鄰接表和原始數據存在一塊兒，搜索的時候加載鄰接表順便把原始向量一塊兒讀上去作精確距離計算。這裏還有一個 SECTOR 的概念，默認大小是 4096，每一個 SECTOR 只放 4096 / node_size 條向量信息，node_size = 單條向量大小+單節點鄰接表大小。

8. 最後再作一個 150000 / n 的全局均勻採樣，存盤，搜索時用來 warmup。

搜索

搜索參數主要有 10 個：

index_type: float/int8/uint8 三選 1，和建索引第一個參數 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 下的統計信息。

搜索流程：

1. 加載相關數據：查詢集、pq 中心點數據、碼本數據、搜索起點等數據，讀取索引 meta；
2. 用建索引時採樣的數據集作 cached_beam_search，統計每一個點的訪問次數，將 num_nodes_to_cache 個訪問頻次最高的點加載到緩存；
3. 默認有一個 WARMUP 操做，和步驟 2 同樣，也是用這個採樣數據集作一次 cached_beam_search，不太明白這一步有什麼用，由於步驟 2 其實已經起到了 warm up 的做用了，難道是再刷一遍系統緩存？
4. 根據給的參數 L 個數，每一個參數 L 都會用查詢集作一遍 cached_beam_search，輸出召回率、qps 等統計信息。warm up 和 統計熱點數據的過程不計入查詢時間。

關於 cached_beam_search:

1. 從候選起點中找離查詢點最近的，這裏用的 pq 距離，起點加入搜索隊列；

2. 開始搜索：

   • 2.1：從搜索隊列中看不超過 beam_width + 2 個未訪問過的點，若是這些點在緩存中，加入緩存命中的隊列，若是不命中，加入未命中的隊列，保證未命中的隊列大小不超過 beam_width；

   • 2.2：對於未命中隊列中的點，發送異步磁盤訪問請求；

   • 2.3：對於緩存命中的點，用原始數據和查詢數據算精確距離加入結果隊列，而後對這些點未訪問過的鄰居點，用 pq 算距離後加入搜索隊列，搜索隊列長度受參數限制；

   • 2.4：處理步驟 a 中緩存未命中的點，過程和步驟 c 同樣；

   • 2.5：搜索隊列爲空時結束搜索，返回結果隊列 topk。

總結

這篇論文加代碼花了一點時間，整體來講仍是很優秀的。論文和代碼思路都比較清晰，經過 kmeans 分若干 overlap 的桶，而後分桶建圖索引，最後合併的思路仍是比較新穎的。至於基於內存的圖索引 Vamana，本質上是一個隨機初始化版本的能夠控制裁邊粒度的 NSG。查詢的時候充分利用了緩存 + pipline，掩蓋了部分 io 時間，提升了 qps。不過按論文中說的，就算機器配置不高，訓練時間長達 5 天，可用性也比較低，後續能夠考慮對訓練部分作一些優化。從代碼來看，質量比較高，能夠直接上生產環境那種。感受建索引那段代碼在算法和實現方面仍是有加速空間的。

參考文獻

1. Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, Rohan Kadekodi. DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node. NeurIPS 2019.
2. Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai. Fast approximate nearest neighbor search with the navigating spreading-out graphs. PVLDB, 12(5):461 – 474, 2019. doi: 10.14778/3303753.3303754. URL http://www.vldb.org/pvldb/vol12/p461- fu.pdf.
3. [NSG GitHub][https://github.com/ZJULearning/efanna]
4. Search Engine For AI：高維數據檢索工業級解決方案
5. Dmitry Baranchuk, Artem Babenko, and Yury Malkov. Revisiting the inverted indices for billion-scale approximate nearest neighbors.

筆者簡介

李成明，Zilliz 研發工程師，東南大學計算機碩士。主要關注大規模高維向量數據的類似最近鄰檢索問題，包括但不限於基於圖和基於量化等向量索引方案，目前專一於 Milvus 向量搜索引擎 knowhere 的研發。喜歡研究高效算法，享受實現純粹的代碼，熱衷壓榨機器的性能。

[Github][https://github.com/op-hunter]