螞蟻金服 ZSearch 在向量檢索上的探索

時間 2019-12-20

原文原文鏈接

圖爲 ZSearch 基礎架構負責人十倍 2019 Elastic Dev Day 現場分享html

引言

ElasticSearch（簡稱 ES）是一個很是受歡迎的分佈式全文檢索系統，經常使用於數據分析，搜索，多維過濾等場景。螞蟻金服從2017年開始向內部業務方提供 ElasticSearch 服務，咱們在螞蟻金服的金融級場景下，總結了很多經驗，這次主要給你們分享咱們在向量檢索上的探索。git

ElasticSearch 的痛點

ElasticSearch 普遍應用於螞蟻金服內部的日誌分析、多維分析、搜索等場景。當咱們的 ElasticSearch 集羣愈來愈多，用戶場景愈來愈豐富，咱們會面臨愈來愈多的痛點：github

如何管理集羣；
如何方便用戶接入和管理用戶；
如何支持用戶不一樣的個性化需求；
...

爲了解決這些痛點，咱們開發了 ZSearch 通用搜索平臺：算法

基於 K8s 底座，快速建立 ZSearch 組件，快捷運維，故障機自動替換；
跨機房複製，重要業務方高保；
插件平臺，用戶自定義插件熱加載；
SmartSearch 簡化用戶搜索，開箱即用；
Router 配合 ES 內部多租戶插件，提升資源利用率；

向量檢索需求

基於 ElasticSearch 的通用搜索平臺 ZSearch 日趨完善，用戶愈來愈多，場景更加豐富。數組

隨着業務的飛速發展，對於搜索的需求也會增長，好比：搜索圖片、語音、類似向量。網絡

爲了解決這個需求，咱們是加入一個向量檢索引擎仍是擴展 ElasticSearch 的能力使其支持向量檢索呢？架構

咱們選擇了後者，由於這樣咱們能夠更方便的利用 ElasticSearch 良好的插件規範、豐富的查詢函數、分佈式可擴展的能力。app

接下來，我來給你們介紹向量檢索的基本概念和咱們在這上面的實踐。框架

向量檢索基本概念

向量從表現形式上就是一個一維數組。咱們須要解決的問題是使用下面的公式度量距離尋找最類似的 K 個向量。less

歐式距離：
- 兩點間的真實距離，值越小，說明距離越近；
餘弦距離：
- 就是兩個向量圍成夾角的 cosine 值，cosine 值越大，越類似；
漢明距離：
- 通常做用於二值化向量，二值化的意思是向量的每一列只有0或者1兩種取值。
- 漢明距離的值就兩個向量每列數值的異或和，值越小說明越類似，通常用於圖片識別；
傑卡德類似係數：
- 把向量做爲一個集合，因此它能夠不只僅是數字表明，也能夠是其餘編碼，好比詞，該值越大說明越類似，通常用於類似語句識別；

由於向量檢索場景的向量都是維度很高的，好比256，512位等，計算量很大，因此接下來介紹相應的算法去實現 topN 的類似度召回。

向量檢索算法

KNN 算法

KNN 算法表示的是準確的召回 topK 的向量，這裏主要有兩種算法，一種是 KDTtree，一種是 Brute Force。咱們首先分析了 KDTree 的算法，發現 KDTree 並不適合高維向量召回，因而咱們實現的 ES 的 Brute Force 插件，並使用了一些 Java 技巧進行加速運算。

KDTree 算法

簡單來說，就是把數據按照平面分割，並構造二叉樹表明這種分割，在檢索的時候，能夠經過剪枝減小搜索次數。

構建樹

以二維平面點(x,y)的集合(2,3)，(5,4)，(9,6)，(4,7)，(8,1)，(7,2)爲例：

圖片來源

按照 x 排序，肯定中間值7，其餘座標分兩邊；
第二層按照 y 排序，第三層按照 x 排序；
而且在構建時維護每一個節點中的 x 最大最小，y 最大最小四個值；

查找最近點

圖片來源

搜索(3,5)的最近鄰：

到根節點距離爲5；
遍歷到右節點(9,6)，發現整棵右子樹的x軸，最小值是8，因此全部右子樹的節點到查詢節點的距離必定都大於8-3=5，因而全部右子樹的節點都不須要遍歷；
同理，在左子樹，跟(5，4)節點比較，(7,2)被排除；
遍歷完(2,3),(4,7)，最近點(5,4) 返回；

結論

Lucene 中實現了 BKDTree，能夠理解爲分塊的 KDTree，而且從源碼中能夠看到 MAX_DIMS = 8，由於 KDTree 的查詢複雜度爲 O(kn^((k-1)/k))，k 表示維度，n 表示數據量。說明 k 越大，複雜度越接近於線性，因此它並不適合高維向量召回。

Brute Force

顧名思義，就是暴力比對每一條向量的距離，咱們使用 BinaryDocValues 實現了 ES 上的 BF 插件。更進一步，咱們要加速計算，因此使用了 JAVA Vector API 。JAVA Vector API 是在 openJDK project Panama 項目中的，它使用了 SIMD 指令優化。

結論

使用 avx2 指令優化，100w 的 256 維向量，單分片比對，RT 在 260ms，是常規 BF 的 1/6。ElasticSearch 官方在7.3版本也發佈了向量檢索功能，底層也是基於 Lucene 的 BinaryDocValues，而且它還集成入了 painless 語法中，使用起來更加靈活。

ANN 算法

能夠看到 KNN 的算法會隨着數據的增加，時間複雜度也是線性增加。例如在推薦場景中，須要更快的響應時間，容許損失一些召回率。

ANN 的意思就是近似 K 鄰近，不必定會召回所有的最近點。ANN 的算法較多，有開源的 ES ANN 插件實現的包括如下這些：

基於 Hash 的 LSH；
基於編碼的 IVFPQ；
基於圖的 HNSW；

ZSearch 依據本身的業務場景也開發了 ANN 插件（適配達摩院 Proxima 向量檢索引擎的 HNSW 算法）。

LSH 算法

Local Sensitive Hashing 局部敏感 hash，咱們能夠把向量經過平面分割作 hash。例以下面圖例，0表示點在平面的左側，1表示點在平面的右側，而後對向量進行屢次 hash，能夠看到 hash 值相同的點都比較靠近，因此在 hash 之後，咱們只須要計算 hash 值相似的向量，就能較準確的召回 topK。

IVF-PQ 算法

PQ 是一種編碼，例如圖中的128維向量，先把向量分紅4份，對每一份數據作 kmeans 聚類，每份聚類出256個聚類中心，這樣，原始向量就可使用聚類中心的編號從新編碼，能夠看出，如今表示一個向量，只須要用4個字節就行。而後固然要記錄下聚類中心的向量，它被稱之爲碼本。

圖片來源

PQ 編碼壓縮後，要取得好的效果，查詢量仍是很大，因此前面能夠加一層粗過濾，如圖，把向量先用 kmeans 聚類成1024個類中心，構成倒排索引，而且計算出每一個原始向量與其中心的殘差後，對這個殘差數據集進行 PQ 量化。用 PQ 處理殘差，而不是原始數據的緣由是殘差的方差能量比原始數據的方差能量要小。

這樣在查詢的時候，咱們先找出查詢出靠近查詢向量的幾個中心點，而後再在這些中心點中去計算 PQ 量化後的 top 向量，最後把過濾出來的向量再作一次精確計算，返回 topN 結果。

圖片來源

HNSW 算法

講 HNSW 算法以前，咱們先來說 NSW 算法，以下圖，它是一個順序構建圖流程：

例如第5次構造 D 點的流程；
構建的時候，咱們約定每次加入節點只連3條邊，防止圖變大，在實際使用中，要經過自身的數據；
隨機一個節點，好比 A，保存下與 A 的距離，而後沿着 A 的邊遍歷，E 點最近，連邊。而後再從新尋找，不能與以前重複，直到添加完3條邊；

查找流程包含在了插入流程中，同樣的方式，只是不須要構建邊，直接返回結果。

HNSW 算法是 NSW 算法的分層優化，借鑑了 skiplist 算法的思想，提高查詢性能，開始先從稀疏的圖上查找，逐漸深刻到底層的圖。

以上這3類算法都有 ElasticSearch 的插件實現：

	LSH 插件	IVSPQ 插件	HNSW 插件
概述	在建立 index 時傳入抽樣數據，計算出 hash 函數。寫入時增長 hash 函數字段。召回用 minimum_should_match 控制計算量	在建立索引時傳入聚類點和碼本，寫入數據就增長聚類點和編碼字段，召回先召回編碼後距離近的再精確計算	擴展 docvalue，在每次生成 segment 前，獲取 docvalue 裏的原始值，並調用開源 hnsw 庫生成索引
優勢	1.實現簡單，性能較高；2.無需藉助其餘 lib 庫；3.無需考慮內存；	1.性能較高；2.召回率高 >90%；3.無需考慮內存；	1.查詢性能最高；2.召回率最高 >95%；
缺點	1.召回率較其餘兩種算法較差，大概在85%左右；2.召回率受初始抽樣數據影響；3.ES 的 metadata很大；	1.須要提早使用 faiss 等工具預訓練；2. ES 的 metadata很大；	1.在構建的時候，segment 合併操做會消耗巨大的 CPU；2.多 segment 下查詢性能會變差；3.全內存；

ZSearch HNSW 插件

咱們根據本身的場景（輕量化輸出場景），選擇了在 ES 上實現 HNSW 插件。由於咱們用戶都是新增數據，更關心 top10 的向量，因此咱們使用了經過 seqNo 去 join 向量檢索引擎方式，減小 CPU 的消耗和多餘 DocValues 的開銷。

對接 Porxima 向量檢索框架：

Proxima 是阿里內部達摩院開發的一個通用向量檢索引擎框架，相似與 facebook 開源的 faiss；
支持多種向量檢索算法；
統一的方法和架構，方便使用方適配；
支持異構計算，GPU；

實現 ProximaEngine

寫入流程，擴展 ElasticSearch 自己的 InternalEngine，在寫完 Lucene 之後，先寫 proxima 框架，proxima 框架的數據經過 mmap 方式會直接刷到磁盤，一次請求的最後，Translog 刷入磁盤。就是一次完整的寫入請求了。至於內存中的 segment，ElasticSearch 會異步到達某個條件是刷入磁盤。

Query 流程

查詢的時候，經過 VectorQueryPlugin，先從 proxima 向量檢索引擎中查找 topN 的向量，得到 seqNo 和類似度，再經過構造 newSetQuery 的 FunctionScoreQuery，去 join 其餘查詢語句。

這裏的數字型 newSetQuery 底層是經過 BKDTree 去一次遍歷所得，性能仍是很高效的。

Failover 流程

固然咱們還要處理各類的 Failover 場景：

數據從遠端複製過來時：
- 咱們攔截了 ElasticSearch 的 recovery action；
- 而後生成 Proxima 索引的快照，這個時候須要經過寫鎖防止數據寫入，快照生成因爲都是內存的，其實很是快；
- 把 Proxima 快照複製到目的端；
- 再進行其餘 ElasticSearch 本身的流程；

數據從本地 translog 恢復時，咱們會記錄快照的 LocalCheckPoint，若是當前 CheckPoint 小於等於 LocalCheckPoint，能夠直接跳過，不然咱們會回查 proxima 檢索引擎，防止數據重試；
目前還有一個狀況，數據會有重複，就是主副分片所有掛掉時，translog 還未刷盤，數據可能已經寫入 proxima 了。

對比

sift-128-euclidean 100w 數據集(https://github.com/erikbern/a...

	HNSW 插件	ZSearch-HNSW 插件
數據寫入（單線程1000個 bulk 寫入）	1.初始寫入 5min，25個 segment，最大 CPU 300%；2.合併成1個 segment 5min，最大 CPU 700%(本地最大)；	構建索引 15min，由於單線程，因此最大CPU 100%
查詢	1. Top 10，召回率98%；2.rt 20ms，合併成1個 segment 後，5ms；	1. Top 10，召回率98%；2.rt 6ms；
優勢	兼容 failover	CPU 消耗少，無額外存儲
缺點	CPU 消耗大，查詢性能跟 segment 有關	主副分片全掛的狀況下會有少許數據重複

總結

ES 參數配置最佳實踐

100w 256維向量佔用空間，大概是0.95GB，比較大：
- 因此更大的堆外內存分配給 pagecache；
- 例如 8C32G 的機器，JVM 設置 8GB，其餘 24GB 留給系統和 pagecache；
設置 max_concurrent_shard_requests：
- 6.x 默認爲節點數*5，若是單節點 CPU 多，能夠設置更大的 shards，而且調大該參數；
BF 算法使用支持 AVX2 的 CPU，基本上阿里雲的 ECS 都支持；

算法總結

KNN 適合場景：
- 數據量小(單分片100w如下)；
- 先過濾其餘條件，只剩少許數據，再向量召回的場景；
- 召回率100%；
ANN 場景：
- 數據量大(千萬級以上)；
- 先向量過濾再其餘過濾；
- 召回率不須要100%；
- LSH 算法：召回率性能要求不高，少許增刪；
- IVFPQ 算法：召回率性能要求高，數據量大(千萬級)，少許增刪，須要提早構建；
- HNSW 算法：召回率性能要求搞，數據量適中(千萬如下)，索引全存內存，內存夠用；