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

引言

ElasticSearch(簡稱 ES)是一個非常受歡迎的分佈式全文檢索系統,常用於數據分析、搜索、多維過濾等場景。螞蟻金服從2017年開始向內部業務方提供 ElasticSearch 服務,我們在螞蟻金服的金融級場景下,總結了不少經驗,此次主要給大家分享我們在向量檢索上的探索。

ElasticSearch 的痛點

ElasticSearch 廣泛應用於螞蟻金服內部的日誌分析、多維分析、搜索等場景。當我們的 ElasticSearch 集羣越來越多,用戶場景越來越豐富,我們會面臨越來越多的痛點:

  • 如何管理集羣;
  • 如何方便用戶接入和管理用戶;
  • 如何支持用戶不同的個性化需求;

爲了解決這些痛點,我們開發了 ZSearch 通用搜索平臺:

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

向量檢索需求

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

隨着業務的飛速發展,對於搜索的需求也會增加,比如:搜索圖片、語音、相似向量。

爲了解決這個需求,我們是加入一個向量檢索引擎還是擴展 ElasticSearch 的能力使其支持向量檢索呢?

我們選擇了後者,因爲這樣我們可以更方便的利用 ElasticSearch 良好的插件規範、豐富的查詢函數、分佈式可擴展的能力。

接下來,我來給大家介紹向量檢索的基本概念和我們在這上面的實踐。

向量檢索基本概念

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

  • 歐式距離:
    • 兩點間的真實距離,值越小,說明距離越近;
  • 餘弦距離:
    • 就是兩個向量圍成夾角的 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 控制計算量 在創建索引時傳入聚類點和碼本,寫入數據就增加聚類點和編碼字段,召回先召回編碼後距離近的再精確計算召回 擴展 Lucene 的 DocValues,在每次生成 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/ann-benchmarks)

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;
  • 設置 maxconcurrentshard_requests:
    • 6.x 默認爲節點數*5,如果單節點 CPU 多,可以設置更大的 shards,並且調大該參數;
  • BF 算法使用支持 AVX2 的 CPU,基本上阿里雲的 ECS 都支持;

算法總結

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

未來規劃

深度學習裏的算法模型都會轉化成高維向量,在召回的時候就需要用相似度公式來召回 topN,所以向量檢索的場景會越來越豐富。

我們會繼續探索在 ElasticSearch 上的向量召回功能,增加更多的向量檢索算法適配不同的業務場景,將模型轉化成向量的流程下沉到 ZSearch 插件平臺,減少網絡消耗。希望可以和大家共同交流,共同進步。

作者介紹

呂梁(花名:十倍),2017年加入螞蟻金服數據中間件,通用搜索平臺 ZSearch 基礎架構負責人,負責 ZSearch 組件在 K8s 上的落地及基於 ES 的高性能查詢插件開發,對 ES 性能調優有着豐富的經驗。

本文轉載自公衆號金融級分佈式架構(ID:Antfin_SOFA)。

原文鏈接

https://mp.weixin.qq.com/s?__biz=MzUzMzU5Mjc1Nw==&mid=2247485678&idx=1&sn=6d79c6a7b3eea2dac4e8d2e93f9abd16&chksm=faa0e734cdd76e22d574acd3e7178eb3a9bf0b2fa0b7e7871b527abdad6176cb2b93f581f972&scene=27#wechat_redirect

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章