【Faiss】PQ和IVF介紹

Faiss是什麼

Faiss是FAIR出品的一個用於向量k-NN搜索的計算庫,其作用主要在保證高準確度的前提下大幅提升搜索速度,根據我們的實際測試,基於1600w 512維向量建庫,然後在R100@1000 (即召回top 1000個,然後統計包含有多少個實際距離最近的top 100)= 87%的前提下單機15線程可以達到1000的qps,這個性能應該是可以滿足大部分的推薦系統召回模塊性能需求了。

 

向量搜索一例:把圖片轉換成向量然後進行k-NN搜索從而實現相似圖片匹配功能

Faiss 的原理

首先來介紹一下Faiss使用時候的數據流:

在使用Faiss的時候首先需要基於原始的向量build一個索引文件,然後再對索引文件進行一個查詢操作,在第一次build索引文件的時候,需要經過Train和Add兩個過程,後續如果有新的向量需要被添加到索引文件的話還可以有一個Add操作從而實現增量build索引,但是如果增量的量級與原始索引差不多的話,整個向量空間就可能發生了一些變化,這個時候就需要重新build整個索引文件,也就是再用全部的向量來走一遍Train和Add,至於具體是怎麼Train和Add的,就關係到Faiss的核心原理了。

Faiss的核心原理其實就兩個部分:

  1. Product Quantizer, 簡稱PQ.
  2. Inverted File System, 簡稱IVF.

Produce Quantizer

Quantizer是通信領域經常出現的一個名詞,在這裏我覺得PQ有一點CS裏面的HashMap的意思,抽象來講就是把連續的空間離散化,這麼做的目的就是爲了優化距離計算的速度,在這篇博客裏面有具體描述PQ的計算過程。PQ有一個Pre-train的過程,一般分爲兩步操作,第一步是Clustering,第二部是Assign,這兩步合起來就是對應到前文提到Faiss數據流的Train階段,可以以一個128維的向量庫爲例:

PQ Clustering && ssign

在做PQ之前,首先需要指定一個參數M,這個M就是指定向量要被切分成多少段,所以M一定要能整除向量的維度,在上圖中M=4,所以向量庫的每一個向量就被切分成了4段,然後把所有向量的第一段取出來做Clustering得到256個簇心(256是一個作者拍的經驗值),再把所有向量的第二段取出來做Clustering得到256個簇心,直至對所有向量的第N段做完Clustering,從而最終得到了256*M個簇心,做完Clustering就開始對所有向量做Assign操作。這裏的Assign就是把原來的N維的向量映射到M個數字,以N=128,M=4爲例,首先把向量切成四段,然後對於每一段向量,都可以找到對應的最近的簇心 ID,4段向量就對應了4個簇心 ID,一個128維的向量就變成了一個由4個ID組成的向量,這樣就可以完成了Assign操作的過程。

 

PQ Search

完成了PQ的Pre-train,就可以看看如何基於PQ做向量檢索了。

同樣是以N=128,M=4爲例,對於每一個查詢向量,以相同的方法把128維分成4段32維向量,然後計算每一段向量與之前預訓練好的簇心的距離,得到一個4*256的表,就可以開始計算查詢向量與庫裏面的向量的距離,而PQ優化的點就在這裏,在計算查詢向量和向量庫向量的距離的時候,向量庫的向量已經被量化成M個簇心 ID,而查詢向量的M段子向量與各自的256個簇心距離已經預計算好了,所以在計算兩個向量的時候只用查M次表,比如的庫裏的某個向量被量化成了[124, 56, 132, 222], 那麼首先查表得到查詢向量第一段子向量與其ID爲124的簇心的距離,然後再查表得到查詢向量第二段子向量與其ID爲56的簇心的距離。。。最後就可以得到四個距離d1、d2、d3、d4,查詢向量跟庫裏向量的距離d = d1+d2+d3+d4。所以在提出的例子裏面,使用PQ只用4×256次128維向量距離計算加上4xN次查表,而最原始的暴力計算則有N次128維向量距離計算,很顯然隨着向量個數N的增加,後者相較於前者會越來越耗時。

由於上面的量化,粗心一共有256*4個128維的向量(可能有重複)。

接着輸入一條查詢向量,查詢向量與4*256條向量比一下,但是這是隔斷開的,依舊分爲4個,也就是查詢向量與第一段256條向量進行比較,離得近就被量化爲誰;接着與第二段的256條向量比較,離得近就被量化爲誰。以此類推。

這樣新輸入的查詢向量就被量化成一個1*4的向量。

此時該1*4的向量與庫中被量化後的N*4的向量進行比較。比如的庫裏的某個向量被量化成了[124, 56, 132, 222], 那麼首先查表得到查詢向量第一段子向量與其ID爲124的簇心的距離,然後再查表得到查詢向量第二段子向量與其ID爲56的簇心的距離。。。最後就可以得到四個距離d1、d2、d3、d4,查詢向量跟庫裏向量的距離d = d1+d2+d3+d4。

PQ和直接暴力計算的比較

Inverted File System

PQ優化了向量距離計算的過程,但是假如庫裏面的向量特別多,每一個查詢向量依舊要進行很多次距離計算,效率依舊還是不夠高,所以這時就有了Faiss用到的另外一個關鍵技術——Inverted File System。

IVF本身的原理比較簡單粗糙,其目的是想減少需要計算距離的目標向量的個數,做法就是直接對庫裏所有向量做KMeans Clustering,假設簇心個數爲1024,那麼每來一個查詢向量,首先計算其與1024個粗聚類簇心的距離,然後選擇距離最近的top N個簇,只計算查詢向量與這幾個簇底下的向量的距離,計算距離的方法就是前面說的PQ,Faiss具體實現有一個小細節就是在計算查詢向量和一個簇底下的向量的距離的時候,所有向量都會被轉化成與簇心的殘差,這應該就是類似於歸一化的操作,使得後面用PQ計算距離更準確一點。使用了IVF過後,需要計算距離的向量個數就少了幾個數量級,最終向量檢索就變成一個很快的操作。

 

IVF-PQ

Faiss本身的索引格式有很多種,原理大都基於PQ和IVF中的兩個或者一個,不同的索引格式對應不同的應用場景,官方給出了一個如何選擇索引格式的guideline,在具體應用的時候可以根據自己的數據量級來參照實驗。

Tips

關於cosine距離:Faiss中的PQ目前還都是基於L2距離(歐式距離),並不支持cosine距離;

關於源代碼閱讀:可以從AutoTone.cpp這個文件開始閱讀;

關於矩陣計算框架:Faiss外部依賴只有一個矩陣計算框架,這個框架可以用OpenBlas也可以用Intel的MKL,使用MKL編譯的話性能會比OpenBlas穩定提升30%,在發佈Faiss的時候MKL還是商用License,所以官方並沒有直接使用,但是現在MKL已經免費了,所以建議使用MKL;

關於OpenMP:Faiss內部實現使用了大量的OpenMP來提高計算效率,其默認的向量檢索也是batch,如果應用場景是單條向量查詢,建議把環境變量OMP_NUM_THREADS設爲1,避免使用OpenMP帶來的多餘性能開銷,這樣可以將單條查詢的latency減少至原本的20%;

默認返回的結果是有可能重複的:要想保證結果不重複就在IndexPQ.cpp:927中MinSumK <float, SemiSortedArray<float>, false>中把第三個參數改成true。

References

  1. https://github.com/facebookresearch/faiss
  2. Product quantization for nearest neighbor search
  3. https://yongyuan.name/blog/ann-search.html
  4. http://vividfree.github.io/%E6%9C%BA%E5%99%A8%E5%AD%A6%E4%B9%A0/2017/08/05/understanding-product-quantization
  5. https://zhuanlan.zhihu.com/c_159623040
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章