K近鄰算法哪家強？KDTree、Annoy、HNSW原理和使用方法介紹
原創 · 作者 | Giant
學校 | 浙江大學
研究方向 | 對話系統、text2sql
知乎專欄 | 大熊貓遊樂園
1、什麼是K近鄰算法
K近鄰算法（KNN）是一種常用的分類和迴歸方法，它的基本思想是從訓練集中尋找和輸入樣本最相似的k個樣本，如果這k個樣本中的大多數屬於某一個類別，則輸入的樣本也屬於這個類別。
關於KNN算法，一個核心問題是： 如何快速從數據集中找到和目標樣本最接近的K個樣本？
本文將從這個角度切入，介紹常用的K近鄰算法的實現方法。具體將從原理、使用方法、時間開銷和準確率對比等方面進行分析和實驗。
2、距離度量
在介紹具體算法之前，我們先簡單回顧一下KNN算法的三要素： 距離度量、k值的選擇和分類決策規則 。
其中機器學習領域常用的距離度量方法，有歐式距離、餘弦距離、曼哈頓距離、dot內積等
主流的近鄰算法都支持上述不同的距離度量。其中n維特徵空間的a、b向量的 歐式距離 體現數值上的絕對差異，而餘弦距離基於餘弦相似度（兩個向量間夾角的餘弦值），體現方向上的相對差異。 如果對向量做歸一化處理，二者的結果基本是等價的。
實際應用中，需要根據業務目標來選擇合適的度量方法。
3、K近鄰算法的實現方法
K近鄰的實現方式多達數十種，筆者從中挑選了幾種常用、經典的方法作爲分析案例。
首先最直觀的想法（暴力法），是線性掃描法。將待預測樣本和候選樣本逐一比對，最終挑選出距離最接近的k個樣本即可，時間複雜度O(n)。對於樣本數量較少的情況，這種方法簡單穩定，已經能有不錯的效果。但是數據規模較大時，時間開銷嚴重無法接受。
所以實際應用中，往往會尋找其他類型的數據結構來保存特徵，以降低搜索的時間複雜度。
常用的存儲結構可以分爲樹和圖兩大類。樹結構的代表是KDTree，以及改進版BallTree和Annoy等；基於圖結構的搜索算法有HNSW等。
4、KDTree和BallTree
KDTree
kd 樹是一種對k維特徵空間中的實例點進行存儲以便對其快速檢索的樹形數據結構。
kd樹是二叉樹，核心思想是對 k 維特徵空間不斷切分（假設特徵維度是768，對於(0,1,2,...,767)中的每一個維度，以中值遞歸切分）構造的樹，每一個節點是一個超矩形 ，小於結點的樣本劃分到左子樹，大於結點的樣本劃分到右子樹。
樹構造完畢後，最終檢索時（1） 從根結點出發，遞歸地向下訪問kd樹 。若目標點 當前維 的座標小於切分點的座標，移動到左子樹，否則移動到右子樹，直至到達葉結點；（2）以此葉結點爲“最近點”， 遞歸地向上回退，查找該結點的兄弟結點中是否存在更近的點 ，若存在則更新“最近點”，否則回退；未到達根結點時繼續執行（2）；（3） 回退到根結點時，搜索結束。
kd樹在維數小於20時效率最高，一般適用於訓練實例數遠大於空間維數時的k近鄰搜索；當空間維數接近訓練實例數時，它的效率會迅速下降，幾乎接近線形掃描。
BallTree
爲了解決kd樹在樣本特徵維度很高時效率低下的問題，研究人員提出了“球樹“ BallTree 。KD 樹沿座標軸分割數據， BallTree將在一系列嵌套的超球面上分割數據，即使用超球面而不是超矩形劃分區域。
具體而言，BallTree 將數據遞歸地劃分到由質心 C 和 半徑 r 定義的節點上，以使得節點內的每個點都位於由質心C和半徑 r 定義的超球面內。通過使用三角不等式 
減少近鄰搜索的候選點數。
coding 實驗
以下實驗均在CLUE下的今日頭條短文本分類數據集上進行，訓練集規模：53360條短文本。
實驗環境：Ubuntu 16.04.6，CPU: 126G/20核，python 3.6
requirement：scikit-learn、annoy、hnswlib
實驗中我使用了bert-as-service服務將文本統一編碼爲768維度的特徵向量，作爲近鄰搜索算法的輸入特徵。
工具包sklearn提供了統一的kdtree和balltree使用接口，可以用一行代碼傳入特徵集合、距離度量方式。
爲了減少推理時間，我這裏僅選取驗證集中前 200條 文本作爲演示。
觀察實驗發現，以 歐式距離 爲度量標準，從5w條知識庫中查找和輸入文本最接近的 top3 ，200條驗證集中kd樹和球樹均正確檢索出153條，但是kd樹檢索200條花費了37秒（185ms/條），球樹花費15秒（75ms/條）， balltree的檢索時間比kdtree快了1倍以上。
5、Annoy
annoy全稱“Approximate Nearest Neighbors Oh Yeah”，是一種適合實際應用的快速相似查找算法。Annoy 同樣通過建立一個二叉樹來使得每個點查找時間複雜度是O(log n)，和kd樹不同的是，annoy沒有對k維特徵進行切分。
annoy的每一次空間劃分，可以看作聚類數爲2的KMeans過程 。收斂後在產生的兩個聚類中心連線之間建立一條垂線（圖中的黑線），把數據空間劃分爲兩部分。
在劃分的子空間內不停的遞歸迭代繼續劃分，直到每個子空間最多隻剩下K個數據節點，劃分結束。
最終生成的二叉樹具有如下類似結構，二叉樹底層是葉子節點記錄原始數據節點，其他中間節點記錄的是分割超平面的信息。
查詢過程和kd樹類似，先從根向葉子結點遞歸查找，再向上回溯即可，完整構建、查找過程可以參考快速計算距離Annoy算法。
coding 實驗
annoy包封裝了算法調用的python接口，底層經C++優化實現。繼續使用頭條文本數據集，調用方法如下：
首先構建一個“AnnoyIndex”索引對象，需指定特徵維度和距離度量標準（支持多種距離度量方式），並將所有數據集樣本特徵順序添加到索引對象中。
之後需要在 build(n_trees) 接口中指定棵數。annoy通過構建一個森林（類似隨機森林的思想）來提高查詢的精準度，減少方差。構建完成後，我們可以將annoy索引文件保存到本地，之後使用時可以直接載入。（完整說明文檔參考annoy的github倉庫）
最後，我們對輸入的200條文本依次查找top3近鄰。
我們發現，正確查找的樣本數和之前相差不大（153 -> 149），但是 查找速度從之前的15秒（75ms/條）降到了0.08秒（0.4ms/條），提升了100倍以上 ，達到了實際開發中的延時要求。
最後提一點，annoy接口中一般需要調整的參數有兩個： 查找返回的topk近鄰和樹的個數 。一般樹越多，精準率越高但是對內存的開銷也越大，需要權衡取捨（tradeoff）。
6、HNSW
和前幾種算法不同，HNSW（Hierarchcal Navigable Small World graphs）是 基於圖存儲 的數據結構。
圖查找的樸素思想
假設我們現在有13個2維數據向量，我們把這些向量放在了一個平面直角座標系內，隱去座標系刻度，它們的位置關係如上圖所示。
HNSW算法就是對上述樸素思想的改進和優化。爲了達到快速搜索的目標，hnsw算法在構建圖時還至少要滿足如下要求： 1）圖中每個點都有“友點”；2）相近的點都互爲“友點”；3）圖中所有連線的數量最少；4）配有高速公路機制的構圖法。
HNSW低配版NSW論文中配了這樣一張圖，短黑線是近鄰點連線，長紅線是“高速公路機制”，如此可以大幅減少平均搜索的路徑長度。
在NSW基礎之上，HNSW加入了 跳錶結構 做了進一步優化。最底層是所有數據點，每一個點都有50%概率進入上一層的有序鏈表。這樣可以保證 表層是“高速通道”，底層是精細查找 。通過層狀結構，將邊按特徵半徑進行分層，使每個頂點在所有層中平均度數變爲常數，從而將NSW的計算複雜度由多重對數複雜度降到了對數複雜度。
關於HNSW的詳細內容可以參考原論文Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs和博客HNSW算法理論的來龍去脈。
coding 實驗
通過 hnswlib 庫，可以方便地調用hnsw算法。
同樣，首先將輸入特徵載入索引模型並保存到本地，下一次可以直接載入內容。具體測試實驗：
最終， 預測200條樣本耗時0.05秒（0.25ms/條） ，速度優於annoy。
此外，同樣的53360條特徵向量（768維度），保存爲靜態索引文件後 ann 索引的大小是227MB，hnsw索引是171MB，從這一點看hnsw也略勝一籌，可以節約部分內存。
參數設置中，ef表示最近鄰動態列表的大小（需要大於查找的topk），M表示每個結點的“友點”數，是平衡時間/準確率的超參數。可以根據服務器資源和查找的召回率等，做相應調整。
7、小結
本文介紹了幾種常用的k近鄰查找算法，kdtree是KNN的一種基本實現算法；考慮到併發、延時等要素，annoy、hnsw是可以在實際業務中落地的算法，其中 bert/sentence-bert+hnsw 的組合會有不錯的召回效果。
除此之外，還有衆多近鄰算法。感興趣的同學可以閱讀相關論文做進一步研究。
Reference：
1.Annoy - Github
2.HNSW - Github
3.Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs
4.Five Balltree Construction Algorithms
5.李航 《統計學習方法》P53-P57: K近鄰法的實現: kd樹
6.快速計算距離Annoy算法原理及Python使用
7.HNSW算法理論的來龍去脈
8.高維空間最近鄰逼近搜索算法評測
本文由作者授權AINLP原創發佈於公衆號平臺，歡迎投稿，AI、NLP均可。 原文鏈接，點擊"閱讀原文"直達：
https://zhuanlan.zhihu.com/p/152522906
推薦閱讀
模型壓縮實踐系列之——bert-of-theseus，一個非常親民的bert壓縮方法
徵稿啓示| 200元稿費+5000DBC（價值20個小時GPU算力）
文本自動摘要任務的“不完全”心得總結番外篇——submodular函數優化
斯坦福大學NLP組Python深度學習自然語言處理工具Stanza試用
關於AINLP
AINLP 是一個有趣有AI的自然語言處理社區，專注於 AI、NLP、機器學習、深度學習、推薦算法等相關技術的分享，主題包括文本摘要、智能問答、聊天機器人、機器翻譯、自動生成、知識圖譜、預訓練模型、推薦系統、計算廣告、招聘信息、求職經驗分享等，歡迎關注！加技術交流羣請添加AINLPer(id：ainlper)，備註工作/研究方向+加羣目的。
閱讀至此了，分享、點贊、在看三選一吧:pray:
