基于k-means 特征的適應性近似最近鄰搜索算法

胡文潔， 楊凱祥， 譚宗元

（東華大學計算機科學與技術學院， 上海 201620）

0 引 言

隨著互聯網技術的快速發展，文字、圖片和視頻等信息呈指數式增長，給信息檢索帶來了巨大的挑戰。 最近鄰檢索問題最早采用暴力搜索的方式，而暴力搜索通常采用線性掃描的方式，即計算目標點與每個數據點的相似性，從而返回與目標點最相似的結果。 針對十億規模的高維數據，線性掃描效率過于低下的問題，研究者們開始研究海量數據中近似最近鄰搜索問題，力求在數據集龐大的條件下，返回與目標最相似的K個鄰居。 近似最近鄰搜索在信息檢索、人臉識別、文件檢索、推薦系統等方面發揮著重要作用。

面對大規模的高維數據，近似最近鄰搜索算法主要解決兩個問題：降低內存占用和加快搜索時間。在降低內存占用方面，近似最近鄰采用量化技術，將原始數據壓縮為二進制編碼進行存儲，很大程度上降低了內存開銷，其代表性算法包括乘積量化算法［1］、優化的乘積量化算法［2］和加和量化算法［3］。在加快搜索時間方面，近似最近鄰搜索領域最常采用的解決措施是對數據構建索引，使用不同的方法對數據進行劃分，在查找目標數據的最近鄰時，僅僅搜索與目標數據最接近的一類數據，減少了數據的比對個數，降低了查詢時間。 其中，基于向量量化的近似最近鄰搜索算法［1-4］，根據數據構建倒排索引；而基于近鄰圖的近似最近鄰搜索算法，是對數據構建圖索引［5-6］。

目前，在根據索引搜索目標數據的最近鄰算法中，通常只搜索固定的粗聚類或者固定的數據點，這會導致搜索訪問不必要的數據點。 因此，本文在IVF-HNSW 算法的基礎上，結合數據的k-means 特征，動態預測每個查詢點需要搜索的粗聚類中心的個數，降低了查找次數以及查詢點平均的搜索時間。

1 相關工作

1.1 相關定義

最近鄰搜索將若干個數據點與給定目標進行匹配，并返回距離目標最近的數據點。 當計算目標向量和數據集向量間的距離時，常用的指標包括余弦相似度、漢明距離和歐式距離。 其中，歐式距離使用向量間的絕對距離表示兩者的遠近程度。 本文研究中，采用歐式距離計算向量間的距離。 歐式距離定義為：在n維空間中，給定數據集中任意一條向量x＝（x1，x2，…，xn） 和目標向量y＝（y1，y2，…，yn），向量x和向量y間的距離計算如式（1）：

歐式距離的值越小，表示兩向量越近；反之，說明兩向量距離越遠。

最近鄰搜索針對目標數據返回精確的搜索結果。 但是當數據量增加到百萬甚至十億規模并且維度上升到成百上千，精確的最近鄰查找的代價非常高，于是在搜索時，犧牲可接受搜索精度范圍內提高搜索的準確率，返回與目標數據最相似的K個最近鄰，即近似最近鄰搜索。 近似最近鄰搜索可定義為：

在D維空間中，給定包含N條向量的數據集合B＝｛b1，b2，…，bn｝ 和目標數據t， 并確定返回最近鄰的個數（K值），則近似最近鄰搜索根據距離函數d（t，b），返回目標數據t在數據集合B中的K個最近鄰，即：

適應性搜索指給定一個查詢點，利用數據點的k-means 特征和回歸模型，預測其終止條件為τ，則此查詢點僅搜索τ個粗聚中心或者τ個候選節點。

1.2 近似最近鄰搜索

目前主流的近似最近鄰搜索方法包括基于哈希的近似最近鄰快速查找技術、基于樹的最近鄰搜索方法、基于向量量化的最近鄰搜索和基于圖的最近鄰搜索等。

基于樹的搜索方法以KD-樹［7］為代表。 KD-樹首先會選取數據各維度中方差最大的維度p； 其次以維度p為基準對數據排序，選取位于中位數的數據點b；最后將數據點b作為閾值，則全部數據點被分為兩部分。 將數據點b作為根節點，兩部分數據分別作為其左右子樹，并遞歸將其左右子樹展開，最終將全部數據構建為一顆索引樹。 在查詢目標數據的最近鄰時，根據二分搜索，查找左子樹或右子樹，直到搜索到葉子節點，返回目標數據的最近鄰。

基于哈希的近似最近鄰查找算法中，局部敏感哈希（LSH）［8］占據重要地位，其主要思想是利用哈希函數簇，將原始數據映射到低維空間，經過哈希函數簇之后，原始空間中相似的兩條數據以更大的概率被分配在同一個桶中；反之，不相似的兩條數據被分配在不同的桶中，并建立桶編號和數據點的索引關系。 在查找目標數據的最近鄰時，利用哈希函數簇對目標數據進行哈希，并計算目標數據和其對應的桶中的全部數據的距離，返回最近鄰。

基于向量量化的方法，以乘積量化（PQ）［1］為代表。 乘積量化的思想是對高維數據進行正交分解，將高維特征空間劃分為M個低維子空間，在每一個低維子空間中運用k-means 方法對所有的數據進行聚類，并對聚類中心進行編碼，每個數據用其對應聚類中心的編碼表示。 在查詢過程中，使用ADC 方法計算目標向量與數據庫向量間的距離，并利用lookups 表存儲查詢向量的每一段子向量到各編碼的距離，通過lookups 表查詢M段距離并進行加和，最終將距離排序，返回目標向量的最近鄰。

在十億規模的數據集上，原有的ADC 方式雖然加快了查詢時間，但仍然是線性掃描，耗時較多，因此Jegou 等人提出結合倒排索引系統［4］進行搜索。倒排索引首先使用k-means 方法對所有的數據進行粗聚類，并計算每個數據的殘差，即原始數據和對應粗聚類中心的差值，之后再使用乘積量化算法對所有殘差進行量化。 最終每一類建立一個倒排列表，倒排列表中存儲該類中所有數據的id 和PQ 量化的結果。 在查詢向量搜索最近鄰的過程中與所有的粗聚類中心計算距離，并選取w個距離最近的粗聚類中心，從w個類中包含的所有數據點中選取最近鄰。

基于圖的方法，代表性的方法為分層可導航小世界（HNSW）［7］，其做法如下：在構建圖索引的過程中，當插入一個新數據點時，計算新數據點插入的層數l，將其插入到0-l層，并在每一層選取新數據點的m個最近鄰并連接，直到數據集中的數據點全部插入圖索引，則多層索引構建完成。 查詢向量搜索最近鄰時，從最上層的入口點開始，計算其在當前層最近的鄰居，并將其作為下層搜索的起始點。 當搜索到最底層時，構建長度為w的優先隊列，用于存儲候選節點。 查詢向量在優先隊列中進行精確的查找，返回其最近鄰。

1.3 IVF-HNSW 算法

IVF-HNSW 算法針對十億規模數據集，結合倒排索引和圖索引，實現了更細粒度的劃分。 一方面，IVF-HNSW 算法將聚類得到的區域劃分為更小的區域；另一方面，在查找最近鄰時，利用剪枝策略過濾部分數據點，加快了檢索速度。 具體實現過程如下：

在構建倒排索引的過程中，首先將整個空間劃分為若干個區域，并計算每個區域的粗聚類中心（質心）。 在每個區域中再次進行劃分，與普通的劃分方式不同的是，IVF-HNSW 算法采用凸組合的形式表示子質心，即每個區域的子質心由質心和此質心相鄰的L個最近質心的凸組合表示。 凸組合形式的優勢在于不需要存儲碼本，節省了內存空間的占用。 其次每條數據庫向量被量化到最近的子質心后，計算其殘差，即原始向量和其對應子質心的差值，并利用乘積量化算法對殘差進行編碼。 最后建立倒排索引，鍵為質心，值為質心對應區域中包含的數據點，并且同一個子區域中的數據排列在一起。

在查詢向量搜索最近鄰的過程中，數據集的規模導致查詢向量搜索質心的效率低下，因此IVFHNSW 算法選擇對質心構建HNSW 索引，用于快速查找最近的質心。 搜索到最近的質心S后，查詢向量和質心S中包含的子質心計算距離，并計算距離的均值E，當查詢向量和子質心的距離大于距離E時，則跳過對應的子區域，過濾不可能成為最近鄰的數據點；當查詢向量和子質心的距離小于或等于距離E時，則在對應的子區域中進行搜索，即計算查詢向量的殘差，并利用ADC 的方式計算查詢向量和已壓縮向量的距離，將距離排序，最終返回查詢向量的最近鄰。

2 自適應的提前終止查詢算法

本文通過分析固定的終止條件對近似最近鄰搜索算法產生的影響，在查詢階段將原始的IVFHNSW 算法和基于k-means 特征的提前終止算法相結合，使得每個查詢點進行適應性搜索，優化了查詢過程，并在sift 和deep 數據集上通過實驗進行證明。

2.1 問題分析

基于倒排索引結構的近似最近鄰搜索算法（如IVFADC 算法等），通常使用固定的搜索終止條件（如Nprobe），Nprobe 指每個查詢點所需要搜索的粗聚類中心的數量。 基于圖索引的搜索算法（如HNSW 算法），通常也使用固定的搜索終止條件（如Efsearch）。 Efsearch 參數指圖索引的0 層使用長度為efsearch 的優先隊列，用于存儲候選節點，每個查詢點在優先隊列中進行精確的搜索。 因此終止條件（即nprobe 和efsearch）的設置，是影響搜索準確度的關鍵指標。 隨著終止條件的增大，查詢點需要搜索數據點的個數增多，則搜索準確度越高，花費的搜索時間也越長；反之，雖然搜索時間較少，但搜索準確度會逐漸降低。

在實際搜索過程中，存在每個查詢點終止條件不同的情況，即有些查詢點只需要很小的終止條件α1就可以檢索到最近鄰，而一些查詢點需要更大的終止條件α2才能搜索到最近鄰，其中α2＞α1。 對于這種情況，為了滿足搜索精度的要求，終止條件的設置需要保證絕大部分查詢點搜索到最近鄰，即終止條件大于α2。 由于所有的查詢點使用相同的終止條件，則導致“容易”搜索到最近鄰的查詢點花費時間搜索不必要的數據點，造成了平均搜索時間升高的問題。

由于IVF-HNSW 算法在查詢階段使用固定的終止條件進行搜索，因此本文采用k-means 特征，動態預測每個查詢點的終止條件，并使得每個查詢點在對應的終止條件下停止搜索，實現了適應性搜索，降低了平均查詢時間。

2.2 算法實現

本文在IVF-HNSW 算法上的優化包括以下步驟：收集數據集的特征、訓練模型、預測查詢點的停止條件、建立IVF-HNSW 索引并整合模型。

2.2.1 收集數據集的特征

數據集特征包括訓練集的輸入特征和輸出特征兩部分。 輸入特征表示訓練集中每個數據點的kmeans 特征，輸出特征指訓練集中每個數據點在IVF-HNSW 算法中真實的停止條件。 收集k-means特征首先利用k-means 方法，將訓練集聚到一萬個類聚類中心，計算每個數據點最近的50 個聚類中心。 選取特征見表1。

其中，t表示訓練集中任一數據點，dist（t，ith nearest coarse cluster centroid）表示數據點和其第i個聚類中心的距離，因此選取的特征共10 個，包括數據點到其第5、10、15 等聚類中心與數據點到最近的聚類中心的距離的比值。

收集輸出特征是在查找過程中，判斷每個數據點是否搜索到其真實最近鄰，如果搜索到真實最近鄰，則記錄搜索的倒排列表的個數。 由于數據庫向量中存在重復的向量，因此每個數據點可能會有多個距離相同的真實最近鄰，于是在收集輸出特征的過程中，只要數據點搜索到其真實最近鄰中的一個，則認為此數據點搜索到真實最近鄰并記錄下最小訪問點的個數。

2.2.2 訓練模型

將每個數據點的k-means 特征作為輸入，對應的真實停止條件作為輸出，選擇神經網絡模型對數據進行回歸擬合。 神經網絡回歸模型的輸入層為數據的10 個k-means 特征；隱藏層的個數為2，其神經元個數設置為100；輸出層個數為1，即神經網絡預測的數據點的終止條件。 當訓練次數為50 次時，模型訓練完成。

2.2.3 預測查詢點的停止條件

首先收集查詢集中每個查詢點的k-means 特征，計算每個查詢點的110 個聚類中心，選取查詢點到其最近的聚類中心的距離和查詢點到其第5、10、15 等聚類中心的距離的比值，作為10 個k-means特征。 查詢點的k-means 特征作為神經網絡回歸模型的輸入，預測其終止條件，即需要搜索的倒排列表的個數和候選節點的個數。

2.2.4 建立IVF-HNSW 索引并進行適應性搜索

將訓練集聚類到一百萬個粗聚類中心后，建立相應的倒排索引并根據粗聚類中心建立HNSW 索引。 在IVF-HNSW 索引的基礎上，加載每個查詢點的停止條件，并將nprobe 和efsearch 參數設置為對應的終止條件，保證每個查詢點進行適應性搜索。

3 實驗

本文采用sift1B、deep1B 以及spacev1B 3 個公開的數據集。 每個數據集包括3 部分：數據集、訓練集和查詢集。 訓練集用于訓練回歸模型，數據集和查詢集用于評估最近鄰搜索的準確度。 spacev1B 數據集包含的訓練集是從數據庫集中切分的一千萬條向量，所以數據庫集僅包含九億九千萬條向量。 有關數據集的詳情見表2。

表2 數據集Tab. 2 Datasets

SIFT 數據集是從圖像中提取的局部特征描述符，每張圖片的描述符組成128 維的特征向量，其中每個坐標為0 ～128 的整數；DEEP 數據集是由自然圖像經過CNN 后生成的96 維特征向量，其中每個坐標的取值范圍為-1 ～1 的浮點數；SPACEV 數據集是微軟必應最先發布的數據集，由Microsoft SpaceV 模型生成的數據庫向量和查詢向量，其中每條向量的維度為100 維，每個坐標的取值范圍為-128～128 的整數。

本文采用召回率recall1＠100 衡量搜索最近鄰的準確度，即搜索結果的第1 最近鄰在其前100 個真實最近鄰中的查詢點個數占所有查詢點個數的比例。 召回率越高，表示搜索的準確度越高。 在比較原始的IVF-HNSW 算法和本文優化的算法時，采用控制變量法，即達到相同的召回率時，對比兩種算法的平均查詢時間。 在相同的召回率下，平均查詢時間越短，證明算法的性能越好。

原始的IVF-HNSW 算法和本文優化的自適應搜索算法在sift1B、deep1B 和spacev1B 3 個數據集上得到的實驗結果如圖1 ～圖3 所示。 其中縱軸表示算法達到的召回率recall1＠100，橫軸表示在對應的召回率下，算法對于所有查詢點的平均查詢時間（以ms 為單位）。 藍色的實線表示基準的IVFHNSW 算法，紅色的實線表示基于IVF-HNSW 算法的自適應搜索算法。

圖1 平均查詢時間-召回率（SIFT1B）Fig. 1 Average query time-Recall （SIFT1B）

圖2 平均查詢時間-召回率（DEEP1B）Fig. 2 Average query time-Recall （DEEP1B）

圖3 平均查詢時間-召回率（SPACEV1B）Fig. 3 Average query time-Recall （SPACEV1B）

從中可以看出，自適應搜索算法在低召回率和高召回率下的平均查詢時間均低于IVF-HNSW 算法，并且隨著召回率的升高，減少的查詢時間越多。在SIFT1B 數據集上，當兩種算法的搜索準確度都達到最高召回率0.958 6 時，IVF-HNSW 算法的平均查詢時間為6.03 ms，而自適應搜索算法的平均查詢時間為5.15 ms，在原始算法的基礎上降低了14.6%；當兩種算法的召回率為0.95 時，基準算法的平均查詢時間為3.78 ms，而自適應搜索算法的平均查詢時間為2.52 ms，下降了33%。 結果表明，在SIFT 數據集上，高召回率下的時間下降更多。

然而在DEEP1B 數據集上，兩種算法達到最高召回率0.948 9 時，IVF-HNSW 算法的平均查詢時間為5.13 ms，而自適應優化算法的平均查詢時間為3.75 ms，降低了27%；在召回率0.90 下，IVF-HNSW算法的平均查詢時間為1.94 ms；而自適應優化算法的平均查詢時間為1.52 ms，降低了21%。 說明隨著召回率的升高，平均查詢時間在原始IVF-HNSW 算法的基礎上下降的幅度更大。

對于SPACEV1B 數據集，在低召回率下兩者沒有區別，但是在最高召回率0.956 6的情況下，平均查詢時間從17.42 ms 降低到15.25 ms，降低了12.5%，說明自適應算法在高召回率下具有更高的優勢。

從實驗結果分析可以看出，本文優化后的自適應搜索算法的性能全面高于原始的IVF-HNSW 算法，在相同的召回率下，進一步降低了平均查詢時間。 由于不同數據集生成特征向量的方式不同，所以同一種算法在不同的數據集上表現不同。 在sift數據集上，自適應搜索算法在低召回率上更有優勢；在deep 數據集的高召回率下，查詢時間下降的百分比更多；在spacev 數據集上，自適應算法在低召回率下的表現幾乎和原始的IVF-HNSW 算法重合，但是在高召回率下的查詢時間有很大程度的降低。

4 結束語

本文在IVF-HNSW 算法的基礎上，建立并訓練神經網絡回歸模型，利用數據的k-means 特征預測每個查詢點的終止條件，進而實現適應性搜索。 實驗結果表明本文優化后的算法進一步降低了近似最近鄰搜索算法在十億規模上的搜索時間，使得近似最近鄰搜索算法以更低的平均查詢時間達到了相同的準確度。