論文賞析：十億級別單機向量檢索方案DiskAnn

摘要

“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 帶寬，太小了也不行，會增加搜索時延。

實驗

實驗分三組：

一、 基於內存的索引比較： 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，在相同召回率下平均搜索路徑最短。

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

數據集： 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）。

三、 基於磁盤的索引： 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]