語義向量召回之ANN檢索

語言: CN / TW / HK
時間 2020-11-16 09:47:26

語義向量召回之ANN檢索

本篇文章主要介紹KDD 2020 Applied Data Science Track Papers中的一篇來自Facebook的語義檢索文章，Embedding-based Retrieval in Facebook Search。關於樣本、模型、訓練等細節可以參考：負樣本為王：評Facebook的向量化召回演算法。本文重點關注其 工程實踐 上的經驗，也藉此機會對基於PQ量化的近似近鄰搜尋 (ANN) 的原理做個梳理。

架構總覽

首先從總體上預覽一下paper提出的EBR系統的架構 (Embedding-based Retrieval System):

EBR系統架構圖示意圖

  • 左上角是 查詢處理模組 ，是雙塔模型中的query-side embedding model

  • 右上角 索引構建模組 ，是用document embedding model推斷doc embedding後構造向量索引。

    構造過程中，先拓展doc的metadata，加入doc embedding，並匯入doc的正排索引中 (比如用於ranking的特徵)；

    同時，通過向量量化技術來降低索引儲存和計算距離代價，並將 量化的結果 存在倒排索引中。

    這一部分也是paper中關注的重點，下文會重點介紹。

  • 左側中間部分是 檢索模組 ，拿到query embedding後，可以通過精心為語義召回設計的 NN運算元(Near Neighbor Operator) 計算距離並進行 語義召回 。

    召回的過程中，既支援原來的布林匹配召回，也支援向量語義召回，還支援混合召回。

    這部分下文也會重點介紹。

  • 左下角是 排序模組 ，得到召回結果後，再經過排序模組進行實時排序，排序時會利用到 embedding特徵 。

paper中涉及工程實現細節的主要是 索引構建模組 和 檢索模組 。下文我會先介紹索引構建模組，這裡頭涉及很多向量量化的概念，比如PQ量化、粗糙量化、殘差量化等，我會先介紹一些量化的背景知識，核心的索引過程和搜尋過程，然後介紹線上檢索模組，認識基於構建好的向量索引，線上是 如何運轉 和 實現召回 的；接著探討paper中提到的工程優化點和調參經驗。最後，做個總結。

索引構建模組

基於Product Quantization的近似最近鄰搜尋，核心的一些問題預覽一下：

  • 問題描述，解決什麼問題？

  • 傳統方法存在什麼問題？

  • 什麼是向量量化？為什麼要量化？量化場景下距離怎麼計算？

  • 什麼是乘積量化？為什麼要乘積量化？

  • 什麼是粗糙量化+殘差量化？為什麼要殘差量化？

  • 索引過程？搜尋過程

問題描述

給定D維向量和集合,需要找到與距離最短的 k 個最近鄰。距離的衡量可以是歐式距離、餘弦距離等。

暴力搜尋問題

如果以最粗暴的方法進行窮舉搜尋，構造 距離矩陣 的複雜度為：。從距離矩陣中查詢到 k個最近鄰 ，最小堆，則複雜度為。

舉個例子：假設，則構造的距離矩陣包含個元素，假設每個距離值32bit，至少佔用1600TB空間，構建距離矩陣的時間是量級，查詢Top-K搜尋時間是量級。

向量量化

所謂向量量化，就是將原來無限的空間對映到一個有限的 向量集合 中，其中是一個自然數。將這個從到集合的函式記為，則，在資訊理論中稱為codebook。即：通過來 近似 代表。

聚類量化

最常用的就是k-means聚類，通過聚類後的 聚類中心向量 來近似 量化 原始的向量。 聚類中心的個數 即為 codebook大小 。因為量化的存在， 任意兩個向量之間的距離 可以通過對應的 量化向量的距離 進行近似，也就是聚類中心向量的距離。因為聚類中心的個數 小了 很多，故計算距離的複雜度也下降了很多。（顯然，這種方式太粗糙了，誤差很大。除非聚類數非常大，極端情況下，聚類數等於樣本數時，每個樣本一個聚類簇，此時無誤差，但是沒有起到減少計算複雜度的目的）

聚類量化的結果：產出聚類中心向量的過程對應train的過程的一部分。 量化後，每個向量都可以用 聚類簇中心下標ID 來 標識 ，根據ID可以獲取聚類簇的中心向量。

聚類量化過程示意圖

如上圖所示，N個向量，通過聚類量化產生多個聚類中心，每個向量屬於某個聚類簇中，那麼就用該聚類簇對應的中心向量來量化該向量，可以用聚類簇中心對應的下標ID來表示，比如：C1量化vec1。

Faiss中對應的實現是 IndexIVFFlat。

乘積量化

動機： 即：PQ量化，很多時候我們向量不同部分之間的 分佈不同 的，比如下圖(3段向量)，因此可以考慮對 向量分段 ，並分別進行 分塊量化 。這個只是直覺原因，本質原因下文講。

乘積量化定義： 將向量分成 個不同的部分，對每個部分進行向量量化,假設平均劃分，則每個部分的維度大小為：D*=D/m； 一個向量 ，可以劃分為m組向量，第 組向量形如： ，每組的codebook為 ，對應的量化器記為 , ( )。 則最終的全域性codebook就是 ，乘積量化的名稱也來源於此。

分塊量化也可以採取聚類量化來實現，則 分塊聚類中心的個數 即為 分塊codebook 的大小。相當於在這個方法下，對每個向量，有【 m個分塊向量】 來量化它，即： 【m個分塊聚類中心向量】 。 示意圖如下：

PQ量化過程示意圖

如上圖所示，將原始向量等分為m組分塊向量，每組都進行聚類量化，那麼每個向量就有m個分塊聚類中心向量來表示，比如vec1用, ...,

共m個向量來量化。

乘積量化的好處 ： 假設每個子codebook大小一樣，記做 ，排列組合一下，那麼相當於能表達的向量空間容量是這麼大， ，但是隻需要 的codebook空間，這也是乘積量化大幅度降低空間佔用的本質原因。 PQ量化原論文中給出的經驗取值是：

，即：分成8塊，每個分塊的codebook大小為256，對應的向量空間大小為約，能夠表達的向量個數足夠大了。

乘積量化結果：個分塊聚類中心向量。 每個向量可以用m個 分塊聚類簇中心下標 ID來標識，所有ID連起來稱為 code 。假設每塊的聚類中心個數為256，則需要8bits，即1byte標識某分塊下哪個聚類中心，m塊則需要m bytes，即code大小為m bytes。

粗糙量化+殘差量化

核心思想： 層次化量化，這個也是Faiss中PQ索引的實現方式。其中粗糙量化使用 聚類量化 ，用劃分到的 粗糙聚類簇的中心向量 來 粗粒度 量化該向量，該結果存在較大的誤差；接著對殘差結果進行 細粒度 乘積量化。這樣的話，誤差就小了。 

1. 總體上，每個向量先進行粗糙量化劃分到某個粗糙聚類簇裡，1個向量對應1個粗糙聚類簇標識，通常稱為粗糙量化ID；
2. 然後計算殘差向量，即：向量-聚類簇中心向量，再對該殘差向量進行分塊，並進行細粒度分塊殘差量化。殘差量化的時候，每一塊對應一個細粒度聚類簇，1個向量M塊，則對應M個細粒度聚類簇標識，通常稱為殘差量化code。
3. 為什麼用殘差量化？原始向量可能會有特別大的分佈差異/不平衡，也就是說可能聚類後，不同聚類簇分佈得非常分散，每個簇所擁有的樣本數極度不平衡。但是通過殘差化後，即：每個樣本向量減去所屬的聚類簇中心向量後，殘差向量之間的差異就不太大了，然後再對殘差向量進行量化，就能更精確的近似原向量。

過程： 具體而言，向量庫先構造一個小規模codebook，量化器為。這個就是所謂的 粗糙量化，或者稱為粗糙聚類 。接著，每個向量y都會有一個殘差。具體而言：記殘差量化步驟的量化器為，則 y 可以通過來表示。

其中，是粗糙量化結果，是殘差量化結果。

這樣的話，【 查詢向量x】和y之間的距離 ：

  • term 2 與查詢向量x 無關， 可以提前計算好；

  • term 1 求x和y的粗糙量化向量的歐式距離，最多計算次，為粗糙聚類中心個數。

  • term 3 求x和y的殘差量化向量的內積，遍歷所有簇的時候，最多計算次，為分塊聚類中心向量。

對應Faiss中的實現是 IndexIVFPQ。

做個小結，量化的結果：

聚類量化 乘積量化 粗糙量化+殘差量化
k個聚類簇中心下標id m x k* 個分塊聚類中心下標組成的 code k個聚類簇中心下標ID，m x k* 個分塊聚類中心下標組成的 code

也就是說，為了表示每個doc的 量化結果 ，可以為doc可以新增兩種 結構化欄位 ： 粗糙量化id， 殘差量化code，用於實時檢索使用。

倒排索引

索引結構

建立從粗粒度聚類中心id 到 doc的對映關係，其中doc的資訊包括：向量id， 向量的殘差量化code。每個doc通過粗糙聚類中心id和殘差量化code就能知道原始向量如何對映到量化向量。

整個倒排索引不需要儲存原始向量本身，索引結構儲存的內容：

  • 儲存標識： 粗糙量化id ， doc id 和殘差量化 code。

    空間佔用很小。m bytes 殘差量化code，即：code_size或者pq_bytes。這個數越大，那麼細粒度聚類簇越大，則精度越高。

    基於id，可以找到對應的粗粒度量化向量 (共k個) ； 基於code 可以找到對應 細粒度殘差量化向量 （共 m x k* 個)。

索引構建過程

搜尋場景中，y可以理解為doc。

  • 通過粗糙量化器將向量對映到，即粗糙量化。這樣就知道掛到哪個連結串列上了。

  • 計算殘差

  • 將殘差量化到，其中包含了 m 個分組，每個分組有對應的一個細粒度聚類簇ID，用1byte表示，則共m bytes，對應code標識。

  • 構造一個entry，其中id是doc的標識，code是殘差量化標識。

搜尋過程

搜尋場景中，x可以理解為查詢向量。

  • 通過粗糙量化器來量化 查詢向量x ，即：找到離x最近的w個粗糙聚類簇。實際上是用於 限定搜尋的範圍 ，只搜尋w個粗糙聚類簇ID索引下的向量。w是個超引數，即nprobes，量化的結果 對應term 1。

  • 選定某個粗糙聚類簇，

    • 計算 x 和該粗糙聚類簇下的 k* （預設即256）箇中心點向量的內積。對應term 3，計算時間複雜度 O(m x D/m x k*) = O(Dk*)，記錄下來，下一步查表用。

    • 遍歷該聚類簇下的doc文件， 計算距離時 ，實際上全是 查表操作 ，term 2是提前算的，term 1粗糙量化時算的，term 3上一步算的。查表時間複雜度實際上是 O(m)

      w個粗糙聚類簇，搜尋時間複雜度為 O(w(Dk* + m))， 另外，返回top-k的話，要加上最小堆排序時間。

檢索模組

布林檢索

Facebook在自研的檢索引擎Unicorn中支援了第一代的近鄰搜尋。

首先簡單介紹下Unicorn。Unicorn原本的檢索方式主要以Term布林匹配為主。

  • doc側：以Term詞袋的方式來表示doc，會基於Term的語義來區分名稱空間。Term上還可以新增一些term特定的 payload 資訊。舉例：John living in Seattle會表示成【 text : john and location : seattle】。

  • query側：定義 Term-level 的布林表示式來表示query。舉例：下述query主要返回doc的文字中有john和smithe字眼，並住在seattle和menlo_park的使用者。

向量檢索

為了支援embedding，需要擴充套件 doc和query的表示

  • doc側：新增結構化欄位 key 來拓展doc詞袋錶示，形如： <key>，比如：key=model-141795009，代表了產出doc embedding model的版本。設定該欄位方便部署多種embedding版本。

  • query側：新增 nn運算元 ，即： nn < key > : radius < radius > ，使用時，通過計算query embedding和 模型產出的doc embedding的距離，來匹配距離在指定radius內的 文件 。radius此處起到 閾值約束 的作用。

索引和計算向量距離時，Facebook將 向量索引和向量線上檢索 通過某種方式轉化成上述已有的 布林檢索語言 ，很巧妙，可以完美融入 現有的布林檢索系統 ，而不需要重新寫一套系統。

先做個對應關係。

布林匹配檢索 語義向量檢索
Term 粗糙聚類中心ID 標識，Cluster-ID
Payload 細粒度聚類中心CODE 標識, Cluster-Code

doc側：每個doc embedding會被 量化 並轉成一個 term 和一個 payload ，相當於是兩個doc的結構化欄位，完美相容已有的檢索系統Unicorn的設計。

  • term，其實就是用於標識倒排索引中，該doc屬於哪個 粗粒度聚類簇 (用於粗糙聚類量化)

  • payload，用於標識每個分塊向量下的 細粒度聚類簇 （用於殘差量化）

query側：

  • term: (nn) 重寫成和粗糙聚類相關的 term。

    重寫規則：計算query embedding和所有粗糙聚類中心距離，選出nprobes個最近的，用粗糙聚類中心ID來標識 (和doc的結構化欄位 Cluster-ID 進行比較)，不同聚類中心對應的Term之間的關係就是or的關係。

  • payload: 對query進行殘差量化，得到滿足 radius約束條件 的細粒度聚類中心，用code標識。

    重寫規則：對每個粗糙聚類簇，計算query embedding和其細粒度聚類中心的距離，距離 滿足<radius約束 的細粒度聚類中心對應的Code 取值記錄一下 (和doc的 結構化欄位Cluster-Code 進行比較)，Code之間是OR關係，但是Code和ID是AND關係。

舉個query改寫的例子：

or ((and( (term(Cluster-ID, '粗糙聚類中心ID-a')), 
          (or (term(Cluster-Code, '殘差聚類中心ID-a1'), term(Cluster-Code, '殘差聚類中心ID-a3'),...)))
    ), 
    (and( (term(Cluster-ID, '粗糙聚類中心ID-b')), 
          (or (term(Cluster-Code, '殘差聚類中心ID-b1'), term(Cluster-Code, '殘差聚類中心ID-b4'),...)))
    ),
    ...
   )

另外，作者強調了radius-mode和topk-mode的差異，radius方式的效能和質量更高。radius是一種 受限制 NN搜尋。top-K需要掃描整個索引庫來找 Top-K結果 。radius效能更好，但是需要確定radius值。

正是因為在已有的布林查詢語言上融入目前的EBR語義檢索，使得混合檢索的實現非常方便。特別是在模糊匹配場景、或者query存在錯誤的場景。舉個例子：

上述查詢，使用者把smith->smithe了，這樣基於term匹配的話，可能會沒有結果。 但是基於nn檢索，只要滿足query embedding和doc embedding的餘弦相似性小於0.24的話，就會有結果。 nprobe是超引數，掃描的粗糙聚類中心個數。

調參經驗

模型改進不大的時候，多調調引數，會有奇效。

  • 調節粗糙量化聚類簇數量num_cluster和實時查詢時掃描的聚類簇數量nprobe。由於不同聚類中心的向量個數可能存在很不平衡的現象，對於兩個對比實驗(比如對比不同構建索引的方法)，即使num_cluster和nprobe設定完全一致，可能兩個實驗掃描的文件數量是不一樣的，因此要監控掃描文件數量，並通過調節這兩個引數，來保證不同對比方法掃描的文件數量一致。

  • 多嘗試使用乘積量化：包括原生PQ，OPQ，PQ with PCA等。能夠顯著 降低索引儲存空間儲存複雜度、查詢時間複雜度 。其中， pq_bytes，即殘差量化codes的大小，很重要 。決定了殘差量化分塊聚類中心的個數，個數越大越精確，比如m=8,k*=256時，pq_byte=1byte x 8= 8 byte。paper中建議採用d/4，d是向量的維度數，假設64維度，則d/4=16，是預設值的2倍。

  • 多進行 線上調參 ，nprobe, num_clusters，pq_bytes。雖然離線能夠直觀感受perf vs recall之間的tradeoff，但是多部署幾套引數進行線上調整，對於理解效能因素對EBR檢索系統的影響會更加直觀。這對於減少 容量開銷 和 離線引數搜尋成本 很有用。

總結

此次分享主要介紹了Facebook在KDD 2020發表的文章中的工程實踐經驗。首先從全域性總覽其系統架構，然後針對索引構建模組和實時檢索模組展開討論。其中，索引構建模組主要介紹了ANN中 向量量化 方法的背景知識。實時檢索模組主要介紹 系統實現 ，如何將向量檢索通過 query改寫規則 等融入現有的布林檢索系統中。最後介紹了一些 調參的經驗 。

做向量召回的初期可以先重點關注模型層面上的優化，索引上也可以先採用簡單的聚類量化的索引構建方式。當向量召回優化到一定程度時，如果想進一步提升效能和召回率，可以考慮借鑑文中的一些經驗，比如以PQ量化的方式來構建索引，調參等。

參考

KDD 2020:  Embedding-based Retrieval in Facebook Search

Faiss基於PQ的倒排索引實現

負樣本為王：評Facebook的向量化召回演算法

Faiss向量召回引擎如何做到快速查詢最近鄰

A library for efficient similarity search and clustering of dense vectors

「」
「」