離散與提煉——一些關於向量召回算法優化方法的思考

✏️  做者介紹：

周語馨，高級雲智能工程師

最近作的不少向量召回的相關工做，主要集中在優化 Faiss 裏面經常使用的幾個算法，包括 IVFFlat 和 IVFPQ，而且針對這兩個算法都作出了專門的優化。

前一陣子靈光乍現，想到了一種與具體算法無關的（或者更嚴格地說，與具體算法相關性較小的）優化方法，能夠優化諸如 Flat、IVFFlat 或者 HNSW 等算法。我稱之爲「方法」而不是「算法」，是由於它聽從原有算法的邏輯，只是在計算過程當中大幅下降了內存帶寬，從而提高性能。所以通過優化的算法在召回率等方面徹底不變。我把這種方法稱爲「離散與提煉（Discretize and Refine)」，以類比於 Google 的「Map and Reduce」，強調其方法論的地位。

開門見山，先來講說「離散與提煉」最核心的思想：

• 離散：使用 int8 和 bfp16 來表達原始向量，計算距離過程當中使用這些壓縮後的向量，從而下降內存帶寬；
• 提煉：使用壓縮後的向量計算獲得的距離存在偏差，可是足以剔除距離太遠的點，以後在很小的一個候選集中使用原始向量計算精確距離。

接下來開始讓我幫你一步步理解。

  1. 背景

向量召回其實就是經典的 KNN（k-NearestNeighbor）問題。KNN 問題的表述爲：假設如今有 N 個向量 yi，當給定一個查詢向量 x 時，找出距離 x 最近的 K 個 yi。爲了方便理解，我以最簡單的暴力搜索算法（Flat）爲例。那麼當給定一個查詢向量 x 時，須要把 x 到每個 yi 的距離都計算出來。 

在 Faiss 以及其餘的 ANN（近似近鄰）搜索庫中，向量都是使用 32 位浮點數（下文中簡稱爲「fp32」）表達的，所以每一個維度佔用 4 個字節。好比一個 128 維度的向量就佔用 512 字節。把向量維度記做 d。那麼一次查詢須要訪問的內存容量有 d*4*N，並且是順序訪問。無論是經過實驗仍是 Faiss wiki 均可以得知，Flat 類（純 Flat 或者 IVFFlat 等）算法的性能瓶頸就在於內存帶寬，大多數時候 CPU 的計算單元都在等待 yi 被送到 cache 中。並且經過統計發現，當壓力較大時，內存讀的時間佔到了整個算法的 90% 以上。 

所以，若是可以使用 int8 來表達向量的每一維，那麼一次查詢的內存訪問量會下降到 d*1*N，即原本的四分之一，理論上性能會提高到接近原來的 4 倍。這就是 「Discretize and Refine」 的動機。 

  2. 最簡單的離散化

先來考慮最容易處理的狀況，即全部 yi 的每一維都落在區間 [-128.0, 127.0) 中。

此時，把 yi 的每一維都四捨五入到最接近的整數上，即 -128 到 127 這 256 個整數，那麼就可使用 int8 存儲。把 yi 通過四捨五入（下稱爲「離散化」）後使用 int8 表達的向量記做 zi，而且計算 yi 到 zi 的距離，記做 ei。下圖是一個簡單的例子來理解 zi 和 ei 的幾何含義。

很直觀，yi 就是實數（暫且把 fp32 看做實數）空間中的點，而 zi 就是其最接近的格點（座標均爲整數的點），而 ei 就是二者的距離。

那麼通過這樣的預處理後，咱們有了 N 個 yi，N 個 zi 和 N 個 ei，分別使用 fp32、int8 和 fp32 表達，佔用 d*4*N 字節、d*1*N 字節和 4*N 字節。  當給定一個 x 時，先計算 x 到每個 zi 的距離，記做 bi。這裏須要強調兩點：

1. x 是用戶傳入的、未經離散化的、使用 fp32 表達的原始向量；
2. zi 的每一維 int8 被依次讀入寄存器後，再經過 CPU 指令在寄存器中轉成 fp32，與 x 的每一維 fp32 計算差的平方，進行累加，最後開方。

所以，從內存到 CPU 的總線上傳輸的都是 int8，而計算的結果是 fp32，即 x 到 zi 的精確距離。把 bi 減去 ei，獲得 Li。至此，搜索的第一步能夠用數學語言表達：

L_i=|x−z_i|−e_i

 

把 x 到yi 的精確距離記做 di，那麼 Li 是 di 的下界。爲何呢？看下圖：

圖中的綠線即爲  bi = |x − zi| ，是 x 到 zi 的精確距離。紅線爲   ei = |yi − zi| ，是 yi 到 zi 的精確距離。根據三角不等式，|x − yi| ≥ |x − zi| − |yi − zi| ，所以有 di  ≥ bi − ei = Li 。至此，咱們使用離散化後的數據獲得了 N 個 Li。

  3. 最簡單的提煉

如何根據 Li 來獲得最小的 K 個 di 呢？這就是整個方法的第二步——提煉（Refine）。

先來看一個生活中的例子。有 100 個運動員賽跑，選出前三名。終點處的裁判手裏有個秒錶，每個運動員到達終點線，裁判都會掐表來記錄時間。因爲裁判是人，存在幾毫秒的反應時間，所以第 i 個運動員過線的客觀時間 Ti，與秒錶記錄下的測量時間 ti 之間存在偏差。可是，ti 與 Ti 之間存在很是大的正相關性，若是 ti 遠比 tj 小，那麼幾乎能夠判定 Ti 比 Tj 小。只有在 ti 和 tj 都比較小且很接近時（好比第三名、第四名幾乎同時過線），才須要調取錄像來獲取 Ti 和 Tj 再進行比較。這個例子很直觀地表達了提煉的思路——經過估算距離就能排除掉絕大部分不可能入選 topK 的點，而在剩餘的少數可能入選 topK 的點中使用精確距離進行最後的篩選。

這裏，咱們選用 Li 做爲估算距離進行提煉。爲何不選 bi 呢？後續算法會展示 Li 的妙用。

既然是最簡單狀況，這裏的提煉算法也是力求最容易理解的，只爲講清數學原理，在性能上並非最優的，後續會給出更優化的實現方法。

首先，將 N 個 i 與 Li 打包成(i,Li)這樣的元組，而且按照 Li 從小到大排序，把排好序的元組數組記做 S。好比 S = [ (14,21.5),(31,25.2),(16,33,7), ...]，含義即爲 到 x 的距離至少爲 21.5，到 x 的距離至少爲25.2，到 x 的距離至少爲 33.7……

而後，初始化一個最大堆 topK。topK 的最大容量爲 K，往裏面丟入若干個(label,   distance)，只會保留最多 K 個 distance 最小的(label,   distance)。而且把 topK 的門檻值記爲 T（當 topK 中的項數小於 K 時，T 爲+∞，不然 T 爲 topK 中最大的 distance）。T 的數學含義爲：若是試圖丟入(label, distance)，而 distance>=T，那麼(label, distance)不可能成爲 topK 其中一項。

for (label, lower_bound) in S:  if lower_bound >=topK.T:      break     distance = |x -y[label]|     topK += (label,distance)

最後，對  S  和  topK  執行以下循環：

代碼依次遍歷 S 中的每一對(label, lower_bound)，若是 lower_bound 不小於 topK 的門檻值，那麼算法終止，此時的 topK 即爲所求。不然，計算 x 與 yi 的精確距離 distance，而且把(label, distance)丟入 topK 中。

代碼很是簡單。可是，爲何代碼第 2 行的判斷能夠確保當前的 topK 就是最後的結果呢？

咱們知道，S 是根據(i, Li)中的 Li 從小到大排序的，若是其中第 i 項知足S_i. lower_bound ≥ topK. T ，那麼第 i 項後續的任一項都大於 T，即：

 S_j.lower_bound≥S_i.lower_bound≥topK.T,∀j≥i

 

又由於 L 做爲 d 的下界，有：

dSj.label  ≥S_j.lower_bound

 

聯立兩式，有：

dSj.label  ≥topK.T,∀j≥i

 

因此，一旦代碼第 2 行的判斷成立，那麼 S 中後續的全部 yi 都不可能入選 topK。

計算 N 個 Li 時，順序訪問了 N 個 zi 和 ei，共 d*1*N+4*N 字節。而提煉過程當中，只須要訪問 Li 最小的 M 個 yi，共 d*4*M 字節。一次查詢的內存訪問量，從 d*4*N 字節變成了 d*N+4*N+d*4*M 字節。經驗上，M 大約在 K 到 4*K 之間（因不一樣數據集而異）。以 SIFT1M 數據集(d=128)、K=100 爲例，內存訪問量從 128*4*1M=512MB 降 到 128*1M+4*1M+128*4*4*100≈132MB 。事實上，計算 Li 的過程當中，由於須要在寄存器中把 int8 轉回 fp32 ，因此計算指令反而變多了。可是因爲 Flat 類算法瓶頸幾乎徹底在內存訪問上，所以以增長 CPU 指令爲代價來換內存帶寬的下降，是很是值得的。至此，離散與提煉的加速原理已經明瞭了。

  4. 基於 heap 的提煉

提煉步驟中有一個明顯能夠優化的地方，即對 N 個 Li 排序。上文提到，S 中只有前 M 項會被訪問，而 M 一般遠小於 N。算法只要求前 M 項按照 Li 排序便可，而以後的 N-M 項既然都不會被訪問，那麼花費大部分 CPU 時間排序後 N-M 項就很不划算。

改進方法是使用堆（優先隊列同理）。若是把 N 個 (i, Li) 放入最小堆中，每次從中取出一個當前  L 最小的 (i, Li) ，那麼就能夠避免徹底排序。用僞代碼表述就是：

min_heap = build_min_heap ((1, L[1]), (2, L[2]), ... , (N, L[N]))while True:  (label, lower_bound) = min_heap.pop()  if lower_bound >= topK.T:  break     distance = |x - y[label]|        topK += (label, distance)

建堆過程複雜度是 O(N)，進行 M 次取出操做複雜度是 O(M*logN)，加起來就是 O(N+MlogN)，比本來的徹底排序的 O(NlogN) 的複雜度要低不少。

  5. 基於 nth_element 的提煉

仔細觀察，會發現基於最小堆的提煉仍然不是最優的。在提煉過程當中，後 N-M 徹底不須要排序，然而堆算法會建堆，儘管不是徹底排序，可是會調整至小根二叉樹的結構，能夠看做一種「淺排序「，也就是仍然浪費了沒必要要的 CPU 時間。

既然經驗上 M 大約爲 K 到 4K，那麼是否可使用 nth_element 算法把 N 個項中最小的 4K 個項取到數組最開頭處，而後只對這 4K 個項作徹底排序呢？若是很不巧，遍歷完這 4K 個仍然沒法結束算法，那麼就繼續從剩下的 N-4K 個項中抽取出 4K 個……如此循環。算法用僞代碼表示以下：

  6. 使用線性變化的離散化

上文中，咱們假定 yi 的每一維都落在 int8 的表達訪問以內，所以離散化操做只是將浮點數近似到最接近的整數。可是，若是 yi 的各個維度的取值區間很大呢，好比[-10000,     10000)？或者不是關於原點對稱的，好比[0, 10000)？這個時候，咱們須要利用歐幾里得距離的兩個特性：

1. 座標系的平移不會改變兩點之間的距離；
2. 全部的維度同時作相同比例的線性縮放，則兩點之間的距離也被按該比例縮放。

這兩點都很直觀，固然，若是須要數學證實也很是容易。

所以，對於一個超出 int8 表達範圍的空間，先對每個維度分別使用一個恰當的偏置，把取值區間平移到關於原點對稱。以後選取一個恰當的比例，把每一個維度作一個縮放，使得每一維都落入 int8 的表達範圍。該操做用數學語言表達就是：

y′ = ky + b

實際操做中，對於每個加入的 yi，都先經過線性變換獲得 yi'，而後對 yi' 作離散化獲得 zi 和 ei。

搜索時，對於用戶給定的 x，也要先執行 x′ =  kx + b ，將 x 映射到 int8 空間中。以後計算 Li 的過程當中使用 x'、zi 和 ei 做爲參數。須要注意的時候，最後 Li 須要除以 k。搜索的第一步用數學語言表達爲：

因爲後續的提煉過程當中，計算的精確距離都是基於原始空間的，因此須要將 Li 除以 k，映射回原始空間中去。

  7. 使用殘差的離散化

真實的數據集每每有必定的聚類特徵，極可能表現爲下圖所示的樣子：

這樣的聚類特徵會致使上文中的離散化方法產生較大的信息損失。好比，上圖中的全部綠色點可能都被離散化到了(100, 80)這個格點上，或者附近少數幾個格點上。那麼搜索時，Li 就過於「粗糙」，點與點之間的區分度不大，以致於不能高效地提煉。

在 ANN 算法中，IVF 類算法特別適用於這種具備聚類特徵的數據集。IVF 算法在構建索引時，將原始數據聚類成 nlist 個類（每一個類的聚類中心記做 Ci)，每一個點屬於其中一個類。當給定一個待搜索的 x 時，找到距離 x 最近的 nprobe 個 Ci，在這 nprobe 個聚類中的點中執行暴力搜索。IVF 算法利用了聚類特徵大幅減小了候選集的大小（本來的 nprobe/nlist），從而犧牲必定精度來換取性能的大幅提高。

基於 IVF 算法，離散化操做能夠進一步優化。對於每個 yi，假設其所在聚類的中心點爲 C_y 。使用 r_i = y_i − C_y ，即 yi 的殘差做爲「基於線性變換的離散化」的輸入，獲得 zi 與 ei。該過程的數學本質是，以每一個聚類中心爲原點創建一個座標系，在該「局部座標系」中對屬於該聚類的點作離散化。如此便可解決信息損失的問題。

當給定 x 時，按照 IVF 算法找出最近的 nprobe 個聚類。對於每個聚類，計算 x 的殘差 q =  x − C_x  ，而後在該聚類中計算 Li。該操做本質上就是把 x 分別映射到每個「局部座標系」中計算距離下界。

  8. 使用 bfp16 應對出界點

上文中全部的討論，都是基於一個假設：向量的取值範圍都是有界的。在有界的狀況下，能夠經過線性變換的方法將向量映射到 int8 的表達範圍中。可是，若是：

1. 向量的取值範圍是無界的；
2. 絕大多數向量都落在某個有限範圍內，但是極少數向量遠遠地出界。

第一種狀況顯然沒法利用 int8 來離散化。

第二種狀況中，若是選用絕大多數向量所在的那個有限範圍，就不得不剔除那些出界的向量。而若是爲了這極少數的向量而擴大取值範圍，就會使得離散化的信息損失變大，從而下降提煉效率。

這兩種狀況，均可以利用 bfp16 來離散化。fp32 使用最高位表達符號，以後的 8 位表達指數，最低的 23 位表達小數部分。而 bfp16 的符號位、指數位與 fp32 相同，惟一的區別是隻使用 7 位來表達小數。

bfp16 擁有與 fp32 同樣的很是寬闊的表達範圍，只是精度會更低。把 fp32變成 bfp16 的過程，其實就是把一個 23 位小數保留到一個 7 位小數（固然，這裏的小數是指二進制中的概念）。使用 bfp16以後，因爲帶寬能夠下降到本來的一半，所以性能仍然有明顯提高（接近本來的兩倍）。

在上述的第一種狀況中，能夠所有使用 bfp16 作離散化，雖然性能不及使用 int8，可是性能提高到接近兩倍仍然很是有誘惑力。

而在上述的第二種狀況中，能夠混合使用 int8 和 bfp16，即絕大多數可以落入有限範圍的向量，使用 int8 作離散化，而少部分「出界」的向量，使用bfp16 作離散化。這樣，即可以享受到 int8 的性能，又能保證"出界點"也被正確處理。

  9. 實測性能

測試平臺爲 Intel(R) Xeon(R) Platinum 8268 CPU @ 2.90GHz，單個 socket 上 6 個內存通道所有插滿了內存條。選用的算法是 IVF1024,Flat (nprobe=20)，測試數據集分別爲 SIFT1M、GIST1M。測試結果以下：

橫座標爲執行搜索的線程數。因爲 8268 單 socket 有 24 個物理核心，48 個超線程，所以線程數從 1 逐漸增長到 48。圖中共用 6 條線，深藍色線、灰色現和淺藍色線分別是原始的  IVF1024, Flat  算法在使用 SSE4 、 AVX2 和 AVX-512 時的總吞吐量（ QPS ，即 query per second ），而橙色線、黃色線和綠色線分別是優化後 IVF1024, Flat 算法在使用 SSE4 、 AVX2 和  AVX-512  時的總吞吐量。觀察這六條線的走勢，能夠得出以下結論：

1. 對於原始算法，隨着線程數的增長，QPS 很快就觸頂，再也不增長；
2. 對於原始算法，更先進的計算指令集幾乎發揮不了優點；
3. 對於優化後的算法，隨着線程數增長，總吞吐量逐步上升，最多能達到 3~3.5 倍的性能；
4. 對於優化後的算法，更先進的計算指令集能夠充分發揮優點。

很顯然，優化後的算法緩解了內存帶寬瓶頸問題，使得多核平臺與先進的向量化指令集能夠充分發揮性能優點。事實上，對於 SIFT1M 數據集，原始算法須要的內存帶寬高達 91GB/s，而優化後的算法僅須要 20GB/s，這使得使用 DCPMM 等廉價存儲器來下降內存採購成本成爲可能。

   歡迎加入 Milvus 社區

github.com/milvus-io/milvus | 源碼

milvus.io | 官網

milvusio.slack.com | Slack 社區

zhihu.com/org/zilliz-11| 知乎

zilliz.blog.csdn.net | CSDN 博客

space.bilibili.com/478166626 | Bilibili

本文分享自微信公衆號 - ZILLIZ（Zilliztech）。
若有侵權，請聯繫 support@oschina.cn 刪除。
本文參與「OSC源創計劃」，歡迎正在閱讀的你也加入，一塊兒分享。