【PaperNotes】Embedding-based Retrieval in Facebook Search

語言: CN / TW / HK
時間 2021-05-24 13:27:38
主題:

論文地址:

1. ACM： https:// dl.acm.org/doi/10.1145/ 3394486.3403305

2. arXiv： https:// arxiv.org/abs/2006.1163 2

這是一篇Applied Data Science Track Paper，非Research Track Paper，側重於工業界技術的落地。

從標題可以看出，論文介紹了Facebook（以下簡稱FB）的向量召回技術，也就是Embedding-Based Retrieval（EBR）。讓我們開始吧，祝開卷有益！

1. 預備知識

首先，搜尋引擎的匹配策略按照查詢關鍵詞（Query）是否完全命中文件（Document，以下簡稱doc），可以分為兩大類： term matching（以下譯作文字匹配） 與 semantic matching（語義匹配） 。

文字匹配，能夠做到完全精確的匹配（Exact Match）。如果是中文，通常先對query進行分詞，得到更細粒度的詞（Term）；再用term去檢索doc，term出現於doc則命中；然後取各個term檢索結果的交集返回。

我之前看過一點《Introduction to Information Retrieval》，它的第一章就在講布林召回（Boolean Retrieval），正是上面這一套傳統的做法。

語義匹配，不再追求完全精確的匹配，而是力圖滿足使用者的搜尋意圖（Search Intent），也就是 意會 。使用者輸入的query，只是Ta搜尋意圖的一種表達形式而已，比如“美國前總統”、“唐納德・特郎普”、“川普”都是“懂王”。要表示使用者意圖，就用到了本文的主角embedding，一種用稠密向量（Dense Vector，沒有特殊說明，後文的向量都指稠密向量）來表徵物件的形式（Representation）。然後基於query embedding與doc embedding來計算結果。

一個有趣的說法：萬物皆可embedding。沒那麼高深，就是用向量來表示物件，然後用向量之間的計算來表示物件之間的關係。

按階段劃分，搜尋引擎通常有兩個大的步驟：召回（Retrieval）與排序（Ranking）。召回，算是一個比較生僻的詞，新聞裡偶爾報道“問題產品的召回”，搜尋中召回是同樣的意思，就是 儘可能把涉及的相關的doc一個不漏找全了，寧抓錯不放過 ，前文的兩類匹配策略也可以說是召回策略；排序的話，就是在召回的基礎上，把更相關的、更符合使用者意圖的doc排在更顯眼的位置。

說到這兒，論文將要介紹的EBR應該比較清晰了：用embedding來表示query與doc，得到query embedding與doc embedding，將召回問題轉化為 向量空間中的最近鄰搜尋問題 。

在搜尋召回時應用EBR的一個挑戰是： 召回是萬里挑一的工作，從海量的資料——上千萬甚至億級的doc中，找出成百上千相關的doc。這對於embeddings的訓練與使用都是極其嚴峻的考驗 。

且看，FB是如何迎難而上的。

2. 本文工作

FB按照階段劃分了面臨的挑戰，並提出了不同的解決方案：

  1. 建模（Modeling）。他們提出了unified embedding——一個雙塔模型，一端是query，另一端是doc；
  2. 服務（Serving）。針對多通道召回（EBR 與布林召回）的問題，他們開發了一個混合召回框架；
  3. 全棧優化。為了對整個搜尋系統進行全面的優化，他們將embeddings 整合到了排序層，構造了一個數據閉環，以學習更好的embeddings。
Facebook的向量召回系統

上圖是FB的EBR召回系統的概覽，各位不妨先花兩分鐘研究一下，再繼續往下看。

2.1 建模

搜尋召回任務可以公式化為 召回優化問題 ，即 給定query、與query相關的doc全集T、模型召回的topK個結果dk ，最大化：

也就是，儘可能將相關的doc找全了。

前文提到EBR將召回問題轉化為向量空間中的最近鄰搜尋問題，EBR模型返回的topK個結果就是“向量空間中，與query embedding‘距離’最近的doc embeddings表示的K個docs”。

建模的目的就在於：如何構造query embedding與doc embedding。在此，FB並沒有太多的創新，使用了目前業界常用的雙塔模型——用兩個神經網路分別作為query與doc的編碼器，如下所示。

使用Unified Embedding的雙塔模型

與常規方法（僅文字進行編碼）不同的是，他們在query側與doc側都加入了輔助特徵，比如query側加入了使用者的位置、社交關係等，是為 unified（大一統的） embedding 。

值得注意的是，對於類別特徵（Categorical Feature），他們的處理與文字一樣，先做embedding，再將得到的特徵向量輸入編碼器，同其他特徵一起計算unified embedding。

從資訊的視角來看，unified embedding的有效性高度依賴於新增的特徵補充的資訊量。

文章簡要介紹了幾點特徵工程的工作：

  1. 文字特徵：通過char n -gram補充了subword資訊；
  2. 位置特徵：在query側，補充了搜尋發起者的城市、地域、國家、使用語言；在document側，補充了開放性的位置資訊，比如群組的位置；
  3. 社交embedding特徵：為了更好地利用社交網路資訊，基於社交圖譜訓練了一個embedding模型。

2.2 訓練

訓練使用的損失函式則是一個三元（Triplet）的間隔損失（Margin Loss）：

其中，q、d+、d- 分別表示query、與query相關的doc（正樣本）、不相關的doc（負樣本），D( u,v )是距離函式，m則是間隔（Margin）。

距離函式與相似度函式是一組相對的概念。作者使用常規的 餘弦相似度（Cosine Similarity） 來度量query embedding與doc embedding的相似度：

因此，距離函式D( u, v )=1-cos( u, v )。

根據作者們的說法，損失函式中，m的選擇很重要，會導致5-10%的離線召回偏差。

此外， 訓練樣本的選擇 對於訓練效果同樣至關重要。本文的做法是， 對於負樣本，使用隨機樣本而不是展現未點選的樣本 。

這一點，有資料的支撐——使用展現未點選樣本的召回效果明顯更差。不過也很好理解：模型工作在召回層，召回的目的是海量樣本中找全相關的樣本。使用隨機樣本，召回候選集是全體樣本，能夠模擬真實自然的召回過程；而使用展現未點選樣本，召回候選集是既有召回演算法過濾後的、是有偏的，將導致模型習得的embeddings也是有偏的，無法習得未展現樣本的特徵。

至於正樣本的選擇，文章實驗了點選樣本與展現樣本，在資料量相當的情況下，兩種策略的效果相近。

2.3 線上服務

2.3.1 近似近鄰搜尋

向量召回的線上服務主要依賴於近似近鄰搜尋演算法（Approximate Near Neighbor，ANN）來建立倒排索引。這一技術方案的優點在於：

  1. 向量量化（Quantization of embedding vectors）使得儲存的成本更低；
  2. 保留了倒排索引，易於整合到當前的檢索系統中。

具體的實現則依賴於FB自家的 Faiss 。

向量量化有兩大元件：Coarse Quantization（CQ）和Product Quantization（PQ），前者主要使用 K-means 將向量轉成相對粗糙的聚類，後者則通過更加精細的量化來支援高效的距離計算。

CQ、PQ都有一些演算法以及引數可供選擇與調整（統稱ANN的引數調整），對此感興趣的同學可以去看原文或 Faiss的Wiki 。以下是幾點tricks：

  1. 引數調整時，以 召回率/掃描的文件數 為指標，比如recall@10；
  2. 最後再調整ANN的引數，或者說一旦向量召回的模型重大變更，ANN的引數務必重新調整；
  3. 務必嘗試OPO（PQ的一種演算法）；
  4. 將pq_bytes（PQ演算法的一項引數，指定了位元組數）設為d/4（d是向量的維數）；
  5. 以線上實驗效果為準，調整nprobe（決定了將為query_emb分配的聚類數），num_clusters，pq_bytes，從而更好地去理解它們對真實效能的影響。

2.3.1 系統實現

在整合向量召回之前，讓我們先來簡單地看一看FB原來的檢索系統（被稱為Unicorn，獨角獸）是怎樣的。

以bag of terms來表示一篇doc，term並不限於文字，而是可以表徵任意的二值屬性，比如“北京的趙喧典”擁有的屬性就包括“text:趙喧典”和“location:北京”。

類似地，以上述形式來表示一個query，比如“北京或浙江的趙喧典”，它的布林表示式就是 (and (or (term location:北京)(term location:浙江)) (term text:趙喧典)) 。利用布林召回，將返回所有使得表示式為真的doc。

為了支援向量召回，檢索系統為每一篇doc添加了額外的embedding欄位。考慮到不同模型的向量召回，一篇doc可以繫結多個embedding，以不同的 <model> 作為區分。相應地，在query側增加了形如 (nn <model> :radius <radius>) 的運算元，表示doc需要滿足在 <model> 向量空間中，doc_emb與query_emb的餘弦距離小於 <radius> 。

上例中，“北京或浙江的趙喧典”，在追加向量召回通道之後，其布林表示式就變為了：

(and
    (or (term location:北京) (term location:浙江))
    (or (term text:趙喧典) (nn M1 :radius 0.2333))
)

此處，FB的工程師們對比過按距離召回與按TOP K召回兩種形式。前者在系統性能與召回結果的質量之間能取得更好的平衡，因為它對於搜尋半徑以及近鄰的相似性都有所限制。

此外，為了提高向量召回的效率與質量，使用了query selection和index selection。

  1. 所謂query selection，就是 有選擇性地觸發向量召回 ，對於向量召回不的，比如搜尋意圖與模型訓練的初衷不同的，或者過於簡單的query，比如使用者剛剛搜尋過或點選過的，不進行向量召回。
  2. 所謂index selection，實際上是構建倒排索引時，只使用月活使用者而不是全部使用者，近期的事件，熱門群組等。

通過上述描述，我們能夠知道，在雙塔模型訓練完畢之後，query embedding模型與doc embedding模型是分開使用的，這也是業界對雙塔模型的常規使用方式。具體地，query_emb是線上實時推斷的，而doc embedding模型是離線部署的，批量地為各個doc生成embedding，併發布到Faiss中。

2.4 其他

2.4.1 整體優化與訓練閉環

考慮到既有的排序演算法是為非向量召回的結果設計的，其排序結果對於向量召回通道的結果可能是次優的。FB的工程師們提出了以下兩點解決思路：

  1. 向量召回的embedding作為排序演算法的特徵。具體地，就是將向量相似度，或者向量召回模型吐出的原始向量透傳給排序演算法。實驗表明，在FB的搜尋場景下，透傳 餘弦相似度 對排序演算法的提升最佳。
  2. 訓練資料反饋閉環。相對於文字召回，向量召回提高了召回率（recall），但是精準率（precision）有所不如。財大氣粗的FB引入了人工評估——將向量召回的結果落日誌，人工評估召回結果的相關性，並用人工標註的結果重新訓練模型，真正形成資料的閉環。

2.4.2 Hard Mining

因為諸如文字召回、向量召回以及其他不同召回技術的混合使用（多通道召回），使得召回結果的資料分佈極其多樣。這對於向量召回模型的訓練是一個不小的挑戰——embedding的學習更加困難。

Hard mining是資訊檢索領域進行資料探勘、構造合適訓練集的技術的統稱。

本文研究探討了兩類hard mining技術：

  1. Hard negative mining 。分析發現：當用戶搜尋人名時，同名情況下，向量召回模型沒有充分利用樣本的社交特徵。這很可能是因為以隨機樣本作為負樣本，負樣本都是異名的，模型只需要關注文字特徵即可。因此，工程師們提出使用與正樣本更加相似的樣本作為負樣本來訓練模型。
    1. Online hard negative mining（此處的online並非線上服務的意思，而是訓練時構造hard negative samples的意思）。具體的做法是，對於一個query，維護了一個正樣本的池子，訓練時，使用與當前正樣本最相似的2條作為負樣本（該數量來自於實驗觀察），並補充上隨機負樣本。
      1. 優點：帶來了相當顯著的召回率的提升；
      2. 缺點：擁有hard negative樣本的比例相對隨機負樣本而言，很低。
    2. Offline hard negative mining（訓練前，在構造訓練樣本時挖掘hard negative樣本）。具體的做法是，先利用某一演算法（KNN、ANN）得到與query最相似的N個doc，再基於一個基於業務特點的挑選策略選擇若干負樣本，然後對模型進行迭代訓練。挑選負樣本-模型迭代，是一個不斷迭代的過程。以下是幾點實驗洞察：
      1. 簡單地用困難負樣本訓練的模型比隨機負樣本還差。分析發現，前者為非文字特徵賦予了過多的權重，在文字匹配能力上差於後者；
      2. 挑選排序模型輸出的位置在100->500（適合）的doc作為負樣本，模型召回效果最佳；
      3. 保留隨機負樣本是極其必要的，它還原了召回的自然過程。可供參考的兩種混合機制：1、混合隨機樣本與困難負樣本，FB的最佳實踐比例是100:1；2、遷移學習，先難後易，反之低效；
      4. 使用KNN來生成困難樣本候選集的時間成本太高，可以直接上ANN。
  2. Hard positive mining 。上文提到，使用點選樣本或展現樣本作為正樣本的效果相似；而此處的做法是，從搜尋日誌中挖掘未展現的正樣本。有效，但是想來成本是相當高的。

2.4.3 Embedding Ensemble

Embedding Ensemble類似於model ensemble，就是整合多個不同的向量召回模型。在前文的介紹中，不同模型將關注到不同的特徵，比如完全使用隨機負樣本的模型將更關注文字特徵、召回能力更強，而使用困難負樣本的模型的精準率更高、對排序更友好，因此embedding ensemble是一個自然能想到的優化的步驟。

是故，文章提出了多階段召喚的方法：第一步關注擴召回，第二步在第一步的基礎上關注區分相似的召回結果。並且，提供了兩種思路：

  1. Weighted concatenation ：
    1. 實際上該方式只有一步，以 加權求和 的方式對不同模型算得的餘弦相似度進行彙總，作為query與doc最終的相似度，再取最相似的作為召回結果。
    2. 線上服務時，為了服務效能，將先計算餘弦相似度再加權求和的方式，近似地變形為先拼接向量再計算且只計算一次餘弦相似度。
    3. 因此，操作流程就是，不同向量加權拼接，以拼接後的向量在Faiss中進行查詢。
  2. 瀑布模型 ：
    1. 第一步使用簡易的召回模型，可以是完全使用隨機負樣本訓練的模型，甚至是text embedding而非unifed embedding，目的是觸達更海量的候選集，不遺漏；
    2. 第二步使用精細一些的模型，比如使用離線困難負樣本訓練的模型，或者使用unified embedding代替text embedding，此時的目的是做一個類似rerank的動作，重新評估第一步得到的候選集中的doc與query的相似度，再按預設的召回量取TOP N條結果。

3. 尾聲

我在半年前看完這篇論文，一直想寫這篇筆記，與諸君交流，但總是寫寫停停；半年後，我有了更多的工作積累，這篇文章反而讀厚了：剛畢業那會兒，我大概只能看到建模與訓練部分，現在也會關注serving以及其他的工程實踐。

原文真是一篇相當nice的論文，強烈建議翻一翻原文！

「」