Facebook： 億級向量相似度檢索庫 Faiss 原理 - 應用

作者 | Chilia    

整理 | NewBeeNLP

最近在使用 ColBERT 雙塔結構進行文本召回，其中必然要涉及到向量相似度查詢，如果只用 brute-force 方法的複雜度實在太高，無法接受。所以必須在 Faiss 上建立索引。因此，今天來學習一下 Faiss 的原理和實際應用。

在這個萬物皆可 embedding 的時代，圖像、文本、商品皆被表示爲 50-1000 維的向量。在雙塔結構中，我們把物品的 embedding 離線存好；當 query 來的時候，就要在一個巨大的商品池中去召回 top k 個商品 (k 爲千級)。這個過程必須很快的完成，來達到較大的 QPS(query per second). 那麼，怎麼去快速召回相似向量呢？

雙塔結構

Faiss 是 Facebook AI 團隊開源的針對聚類和相似性搜索庫，爲稠密向量提供高效相似度搜索和聚類，支持十億級別向量的搜索，是目前最爲成熟的近似近鄰搜索庫。

它包含多種搜索任意大小向量集（向量集大小由 RAM 內存決定）的算法，以及用於算法評估和參數調整的支持代碼。Faiss 用 C++ 編寫，並提供與 Numpy 完美銜接的 Python 接口。除此以外，對一些核心算法提供了 GPU 實現。

  1. Faiss 簡介

如果用暴力搜索的方法，能夠得到完全正確的 “標準答案”，但是其時間複雜度爲 O(mn)，這根本無法接受。如果犧牲一些精度的話，比如允許與參考結果有一點點偏差，那麼相似性搜索能快幾個數量級。加快搜索速度還涉及到數據集的預處理，我們通常把這個預處理操作稱作**「索引」**。我們主要關注三個評價指標：

通常我們都會在內存資源的限制下在速度和精準度之間權衡。Faiss 中提供了若干種方法實現數據壓縮，包括 PCA、Product-Quantization 等向量壓縮的方法 (當然還有其它的一些優化手段，但是 PCA 和 PQ 是最爲核心的)，來存儲十億級別的數據。

  1. 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.

2.1 Product Quantizer

矢量量化方法，即 vector quantization，其具體定義爲: 將向量空間的點用一個有限子集來進行編碼的過程。常見的聚類算法，都是一種矢量量化方法。而在 ANN(Approximate Nearest Neighbor, 近似最近鄰) 搜索問題中，向量量化方法又以乘積量化 (PQ, Product Quantization) 最爲典型。

2.1.1 Pretrain

PQ 有一個 Pre-train 的過程，一般分爲兩步操作，第一步 「Cluster」，第二步 「Assign」，這兩步合起來就是對應到前文提到 Faiss 數據流的 **「Train」**階段，可以以一個 128 維的向量庫爲例：

PQ 乘積量化的核心思想還是聚類，K-Means 就是 PQ 乘積量化子空間數目爲 1 的特例。

在做 PQ 之前，首先需要指定一個參數 M，這個 M 就是指定向量要被切分成多少段，在上圖中 M=4，所以向量庫的每一個向量就被切分成了 4 段，然後把所有向量的第一段取出來做 Clustering 得到 256 個簇心（256 是一個作者拍的經驗值）；再把所有向量的第二段取出來做 Clustering 得到 256 個簇心，直至對所有向量的第 N 段做完 Clustering，從而最終得到了 256*M 個簇心。

做完 Cluster，就開始對所有向量做 Assign 操作。這裏的 Assign 就是把原來的 N 維的向量映射到 M 個數字，以 N=128，M=4 爲例，首先把向量切成四段，然後對於每一段向量，都可以找到對應的最近的簇心 ID，4 段向量就對應了 4 個簇心 ID，一個 128 維的向量就變成了一個由 4 個 ID 組成的向量，這樣就可以完成了 Assign 操作的過程 -- 現在，128 維向量變成了 4 維，每個位置都只能取 0~127，這就完成了向量的壓縮。

2.1.2 查詢

完成了 PQ 的 Pre-train，就可以看看如何基於 PQ 做向量檢索了。查詢過程如下圖所示：

同樣是以 N=128，M=4 爲例，對於每一個查詢向量，以相同的方法把 128 維分成 4 段 32 維向量，「然後計算每一段向量與之前預訓練好的簇心的距離」，得到一個 4*256 的表。然後就可以開始計算查詢向量與庫裏面的向量的距離。

此時，庫的向量已經被量化成 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/4 維向量距離計算加上 4xN 次查表，而最原始的暴力計算則有 N 次 128 維向量距離計算，很顯然隨着向量個數 N 的增加，後者相較於前者會越來越耗時。

2.2 Inverted File System

PQ 優化了向量距離計算的過程，但是假如庫裏面的向量特別多，依然逃不了一個遍歷整個庫的過程，效率依舊還是不夠高，所以這時就有了 Faiss 用到的另外一個關鍵技術——Inverted File System。

倒排乘積量化 (IVFPQ) 是乘積量化 (PQ) 的更進一步加速版。其加速的本質就是在前面強調的是加速原理：brute-force 搜索的方式是在 「全空間」 進行搜索，爲了加快查找的速度，幾乎所有的 ANN 方法都是通過 「對全空間分割 (聚類)」，將其分割成很多小的子空間，在搜索的時候，通過某種方式，快速鎖定在某一（幾）子空間，然後在該（幾個）子空間裏做遍歷。

在上一小節可以看出，PQ 乘積量化計算距離的時候，距離雖然已經預先算好了，但是對於每個樣本到查詢樣本的距離，還是得老老實實挨個去求和相加計算距離。

但是，實際上我們感興趣的是那些跟查詢樣本相近的樣本 (下面稱之爲 “感興趣區域”)，也就是說老老實實挨個相加其實做了很多的無用功，如果能夠通過某種手段快速 「將全局遍歷鎖定爲感興趣區域」，則可以捨去不必要的全局計算以及排序。

倒排 PQ 乘積量化的” 倒排 “，正是這樣一種思想的體現，在具體實施手段上，採用的是通過聚類的方式實現感興趣區域的快速定位，在倒排 PQ 乘積量化中，聚類可以說應用得淋漓盡致。

要想減少需要計算的目標向量的個數，做法就是直接對庫裏所有向量做 KMeans Clustering，假設簇心個數爲 1024。那麼每來一個 query 向量，首先計算其與 1024 個粗聚類簇心的距離，然後選擇距離最近的 top N 個簇，只計算查詢向量與這幾個簇底下的向量的距離，計算距離的方法就是前面說的 PQ。

Faiss 具體實現有一個小細節，就是在計算查詢向量和一個簇底下的向量的距離的時候，所有向量都會被轉化成與簇心的殘差，這應該就是類似於歸一化的操作，使得後面用 PQ 計算距離更準確一點。使用了 IVF 過後，需要計算距離的向量個數就少了幾個數量級，最終向量檢索就變成一個很快的操作。

  1. Faiss 應用

3.1 IndexFlatL2

這種索引的方式是用暴力的 (brute-force) 精確搜索計算 L2 距離。

import faiss                   # make faiss available
index = faiss.IndexFlatL2(d)   # build the index, d = 128, 爲dimension
index.add(xb)                  # add vectors to the index, xb 爲 (100000,128)大小的numpy
print(index.ntotal)            # 索引中向量的數量, 輸出100000

k = 4                          # we want to see 4 nearest neighbors
D, I = index.search(xq, k)     # xq爲query embedding, 大小爲(10000,128)
## D爲查詢得到的物品embedding與query embedding的距離,大小(10000,4)
## I爲和query embedding最接近的k個物品id，大小(10000,4)
print(I[:5])                   # neighbors of the 5 first queries

輸出：

100000
[[ 381  207  210  477]
 [ 526  911  142   72]
 [ 838  527 1290  425]
 [ 196  184  164  359]
 [ 526  377  120  425]]

IndexFlatL2 的結果精確，可以用來作爲其他索引測試中準確性程度的參考。當數據集比較大的時候，暴力搜索的時間複雜度很高，因此我們一般會使用其他方式的索引。

3.2 IndexIVFFlat

爲了加快搜索的速度，我們可以將數據集分割爲幾部分，將其定義爲 Voronoi Cells，每個數據向量只能落在一個 cell 中。查詢時只需要查詢 query 向量落在 cell 中的數據了，降低了距離計算次數。這就是上文說的倒排索引方法。

IndexIVFFlat 需要一個訓練的階段，其與另外一個索引 quantizer 有關，通過 quantizer 來判斷屬於哪個 cell。IndexIVFFlat 在搜索的時候，引入了 nlist(cell 的數量) 與 nprob(執行搜索的 cell 數) 參數。增大 nprobe 可以得到與 brute-force 更爲接近的結果，nprobe 就是速度與精度的調節器。

import faiss                   # make faiss available
nlist = 100
k = 4

quantizer = faiss.IndexFlatL2(d)  # d = 128, dimension
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)

index.train(xb) ## xb: (100000,128)
index.add(xb)                
index.nprobe = 10              # 默認 nprobe 是1 ,可以設置的大一些試試
D, I = index.search(xq, k)
print(I[-5:])                  # 最後五次查詢的結果

3.3 IndexIVFPQ

IndexFlatL2 和 IndexIVFFlat 都要存儲所有的向量數據。對於超大規模數據集來說，可能會不大現實。因此 IndexIVFPQ 索引可以用來壓縮向量，具體的壓縮算法就是 PQ。

import faiss

nlist = 100
m = 8 ##每個向量分8段
k = 4 ##求4-近鄰
quantizer = faiss.IndexFlatL2(d)    # 內部的索引方式依然不變
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) # 每個向量都被編碼爲8個字節大小
index.train(xb)
index.add(xb)
index.nprobe = 10                
D, I = index.search(xq, k)          # 檢索
print(I[-5:])

本文參考資料

[1] 圖像檢索：再敘 ANN Search: https://yongyuan.name/blog/ann-search.html

[2] Faiss | 哈嘍哈咯: https://blog.razrlele.com/p/2594

[3] faiss: https://github.com/facebookresearch/faiss

[4] Faiss 從入門到實戰精通: https://blog.csdn.net/bitcarmanlee/article/details/106447629

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/nh5JuTWiytqh4Eo8r5F3xg

猜你喜歡