天天看點

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

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

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

引言

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

ElasticSearch 的痛點

ElasticSearch 廣泛應用于螞蟻金服内部的日志分析、多元分析、搜尋等場景。當我們的 ElasticSearch 叢集越來越多,使用者場景越來越豐富,我們會面臨越來越多的痛點:

  • 如何管理叢集;
  • 如何友善使用者接入和管理使用者;
  • 如何支援使用者不同的個性化需求;
  • ...

為了解決這些痛點,我們開發了 ZSearch 通用搜尋平台:

  • 基于 K8s 底座,快速建立 ZSearch 元件,快捷運維,故障機自動替換;
  • 跨機房複制,重要業務方高保;
  • 插件平台,使用者自定義插件熱加載;
  • SmartSearch 簡化使用者搜尋,開箱即用;
  • Router 配合 ES 内部多租戶插件,提高資源使用率;
螞蟻金服 ZSearch 在向量檢索上的探索

向量檢索需求

基于 ElasticSearch 的通用搜尋平台 ZSearch 日趨完善,使用者越來越多,場景更加豐富。

随着業務的飛速發展,對于搜尋的需求也會增加,比如:搜尋圖檔、語音、相似向量。

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

我們選擇了後者,因為這樣我們可以更友善的利用 ElasticSearch 良好的插件規範、豐富的查詢函數、分布式可擴充的能力。

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

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

向量檢索基本概念

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

螞蟻金服 ZSearch 在向量檢索上的探索
  • 歐式距離:
    • 兩點間的真實距離,值越小,說明距離越近;
  • 餘弦距離:
    • 就是兩個向量圍成夾角的 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)為例:

螞蟻金服 ZSearch 在向量檢索上的探索
圖檔來源

https://blog.csdn.net/richard9006/article/details/90058465
  • 按照 x 排序,确定中間值7,其他坐标分兩邊;
  • 第二層按照 y 排序,第三層按照 x 排序;
  • 并且在建構時維護每個節點中的 x 最大最小,y 最大最小四個值;

查找最近點

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

搜尋(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 指令優化。

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

使用 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。

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

IVF-PQ 算法

PQ 是一種編碼,例如圖中的128維向量,先把向量分成4份,對每一份資料做 kmeans 聚類,每份聚類出256個聚類中心,這樣,原始向量就可以使用聚類中心的編号重新編碼,可以看出,現在表示一個向量,隻需要用4個位元組就行。然後當然要記錄下聚類中心的向量,它被稱之為碼本。

螞蟻金服 ZSearch 在向量檢索上的探索
https://yongyuan.name/blog/vector-ann-search.html

PQ 編碼壓縮後,要取得好的效果,查詢量還是很大,是以前面可以加一層粗過濾,如圖,把向量先用 kmeans 聚類成1024個類中心,構成反向索引,并且計算出每個原始向量與其中心的殘差後,對這個殘差資料集進行 PQ 量化。用 PQ 處理殘差,而不是原始資料的原因是殘差的方差能量比原始資料的方差能量要小。

這樣在查詢的時候,我們先找出查詢出靠近查詢向量的幾個中心點,然後再在這些中心點中去計算 PQ 量化後的 top 向量,最後把過濾出來的向量再做一次精确計算,傳回 topN 結果。

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

講 HNSW 算法之前,我們先來講 NSW 算法,如下圖,它是一個順序建構圖流程:

  • 例如第5次構造 D 點的流程;
  • 建構的時候,我們約定每次加入節點隻連3條邊,防止圖變大,在實際使用中,要通過自身的資料;
  • 随機一個節點,比如 A,儲存下與 A 的距離,然後沿着 A 的邊周遊,E 點最近,連邊。然後再重新尋找,不能與之前重複,直到添加完3條邊;

查找流程包含在了插入流程中,一樣的方式,隻是不需要建構邊,直接傳回結果。

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

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

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

以上這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 等工具預訓練

  1. ES 的 metadata很大 | 1.在建構的時候,segment 合并操作會消耗巨大的 CPU

    2.多 segment 下查詢性能會變差

3.全記憶體 |

ZSearch HNSW 插件

我們根據自己的場景(輕量化輸出場景),選擇了在 ES 上實作 HNSW 插件。因為我們使用者都是新增資料,更關心 top10 的向量,是以我們使用了通過 seqNo 去 join 向量檢索引擎方式,減少 CPU 的消耗和多餘 DocValues 的開銷。

對接 Porxima 向量檢索架構:

**

  • Proxima 是阿裡内部達摩院開發的一個通用向量檢索引擎架構,類似與 facebook 開源的 faiss;
  • 支援多種向量檢索算法;
  • 統一的方法和架構,友善使用方适配;
  • 支援異構計算,GPU;
螞蟻金服 ZSearch 在向量檢索上的探索

實作 ProximaEngine

寫入流程,擴充 ElasticSearch 本身的 InternalEngine,在寫完 Lucene 以後,先寫 Proxima 架構,Proxima 架構的資料通過 mmap 方式會直接刷到磁盤,一次請求的最後,Translog 刷入磁盤。就是一次完整的寫入請求了。至于記憶體中的 segment,ElasticSearch 會異步到達某個條件是刷入磁盤。

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

Query 流程

查詢的時候,通過 VectorQueryPlugin,先從 proxima 向量檢索引擎中查找 topN 的向量,獲得 seqNo 和相似度,再通過構造 newSetQuery 的 FunctionScoreQuery,去 join 其他查詢語句。

這裡的數字型 newSetQuery 底層是通過 BKDTree 去一次周遊所得,性能還是很高效的。

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

Failover 流程

當然我們還要處理各種的 Failover 場景:

  • 資料從遠端複制過來時:
    • 我們攔截了 ElasticSearch 的 recovery action;
    • 然後生成 Proxima 索引的快照,這個時候需要通過寫鎖防止資料寫入,快照生成由于都是記憶體的,其實非常快;
    • 把 Proxima 快照複制到目的端;
    • 再進行其他 ElasticSearch 自己的流程;
螞蟻金服 ZSearch 在向量檢索上的探索
  • 資料從本地 translog 恢複時,我們會記錄快照的 LocalCheckPoint,如果目前 CheckPoint 小于等于 LocalCheckPoint,可以直接跳過,否則我們會回查 Proxima 檢索引擎,防止資料重試;
  • 目前還有一個情況,資料會有重複,就是主副分片全部挂掉時,translog 還未刷盤,資料可能已經寫入 Proxima 了。

對比

sift-128-euclidean 100w 資料集(

https://github.com/erikbern/ann-benchmarks)
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 算法: 召回率性能要求搞,資料量适中(千萬以下),索引全存記憶體,記憶體夠用;

未來規劃

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

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

作者介紹

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

附件