大規模向量檢索與量化方法

ZacksTang發表於2024-11-18

1. 向量檢索

在向量檢索中，KNN（K-Nearest Neighbors）和ANN（Approximate Nearest Neighbor）是兩種最常見的方法，它們都用於根據特徵向量找到資料點之間的相似性，但它們在精確度和效率上有所不同。

KNN是一種基本的分類和迴歸方法，它根據一個樣本在特徵空間中的K個最近鄰樣本的類別，來預測該樣本的類別。對於迴歸問題，則可能是根據K個最近鄰樣本的值來預測該樣本的值。在向量檢索中，KNN演算法做的即是暴力檢索，計算Query向量與向量庫中每個向量的距離（或是方向，取決於衡量指標），來找出最相似的K條向量。KNN的優點是簡單易懂，不需要訓練，但它的缺點是計算成本高，尤其是在大資料集上，因為它需要計算測試樣本與資料集中所有樣本之間的距離。

基於KNN在大規模資料量上的不足，ANN對此做了進一步改進。它一種用於在大規模資料集中快速找到與給定查詢向量近似最近的鄰居向量的方法。與KNN不同，ANN並不總是返回確切的最近鄰，而是在一定的誤差範圍內返回近似的最近鄰，以換取更快的檢索速度。ANN演算法的步驟通常包括：

  1. 索引構建：構建資料集的索引，以加快搜尋速度。常用的索引方法包括IVF（Inverted File）、HNSW（Hierarchical Navigable Small World）。
  2. 查詢處理：接收一個查詢向量，並使用索引快速找到一組候選的最近鄰。
  3. 驗證和細化：對候選的最近鄰進行精確的距離計算，以確定最終的最近鄰。

ANN的優點是速度快，適合大規模資料集，但它的缺點是結果可能不是完全精確的，並且構建索引會需要引入額外的空間儲存。ANN在許多應用中非常有用，如語義檢索、影像檢索、推薦系統等，這些應用可以容忍一定程度的近似，但需要快速響應。

1.1. ANN大規模向量檢索的資源消耗

在大規模向量檢索中，雖然我們可以利用ANN的方式，實現更快的TopK的向量召回，但是對資源的消耗也是巨大的。以OpenSearch中的ANN檢索為例，在使用nmslib實現的HNSW演算法下，對於10億條128維的向量進行檢索，構建HNSW索引所需的儲存空間為704GB[1],如果加上一個副本實現更高的資料可用性，則所需儲存空間還需要乘以2，為1408GB。在檢索時，仍需要將索引載入到記憶體進行查詢，對應即需要1408GB的記憶體空間（多個節點提供，並非單一節點）儲存HNSW索引。

不過雖然HNSW索引所需資源消耗巨大，但也能得到很高的效能收益，在記憶體資源給足的情況下，對10億條向量進行top10的檢索，可以達到99%的top10的召回，同時p50延遲僅為23ms左右。

2. 量化方法

對於ANN在大規模向量檢索下所需資源量較大的情況，業內也提出了各種量化方法來對向量進行量化，節省所需的資源的同時，仍然能夠實現ANN檢索。其中比較有代表性的包括乘積量化（PQ：Product Quantization）、二進位制量化（BQ：Binary Quantization）和標量量化（SQ：Scalar Quantization）等。下面主要介紹PQ和BQ。

3. Product Quantization

乘積量化（Product Quantization，簡稱PQ）是一種用於壓縮高維向量以減少記憶體佔用並提高最近鄰搜尋速度的方法。它的核心思想是將原始的高維向量空間分解為若干個低維向量空間的笛卡爾積，並對這些低維向量空間分別進行量化。

乘積量化的核心思想是聚類。在構建索引時，主要分為2個步驟：Cluster和Assign。下面是流程示意圖：

PQ提供一個引數m_split（有些演算法實現裡也稱m），這個引數控制向量被切分的段數。Cluster過程為：

  1. 假設向量池有N條向量，每個向量是128維，將每個向量被切分為4個“段”，這樣就得到了N *（4個32維的向量）
  2. 對每段的向量進行聚類（例如k-mean），假設聚類的cluster個數為256，由於有4個段，所以每個段（N條32維的向量）都會分為256個簇，併產生256個“中心點”（聚類的中心向量），最終簇的總個數為4 * 256個。

Assign過程為：

  1. 先對每個“簇”進行編碼，一共有4 * 256 個“簇”及對應編碼，由於“簇”為256個，所以僅需要8位的編碼id即可表示一個簇。
  2. 對現有每個向量，先切分為4個段，然後對每個段的小向量，分別計算其對應的最近的“簇”（計算與“中心點”的距離），然後用這個“簇”的id進行編碼。
  3. 這樣每條向量即可用4個段的 0-255之間的編號進行表示。也即為4個位元組即可儲存。這樣就壓縮了向量，極大的節省了記憶體

接下來是查詢：

在查詢一個向量時，和訓練時一樣，先切為4個段，然後計算每個子向量與對應的256個簇中心之間的距離，然後儲存在一個距離矩陣或陣列中。接下來即可透過查表，來計算query向量和每個向量之間的距離。計算方式就是累加每個子向量之間的距離之和，這樣就算出了query向量與庫裡所有向量的距離，最後做TopK即可。

這種方法可以極大減少計算量。但乘積量化是一種有失真壓縮方法，它在預測精度上可能有所降低，但可以大幅壓縮高維向量以減少記憶體使用，同時保持足夠的資訊以進行高效的相似性搜尋。仍然以OpenSearch中做十億條128維向量查詢為例，使用IVF+PQ的方式，在top10召回率為0.66的情況下，僅需要114GB記憶體即可完成ANN查詢索引的構建，但p50查詢延遲在117ms左右，極大節省了查詢成本，但也伴隨著明顯的查詢損失。

4. Binary Quantization

二進位制量化（Binary Quantization）在最近（本文寫於2024年11月初）得到了較多的關注。OpenSearch在最近的 2.17版本[2]中釋出了Binary Quantization（BQ）的功能，用於壓縮向量，以更低成本的方式實現向量的儲存與檢索。另一方面，Elasticsearch 最近的8.16版本[3]也釋出了Better Binary Quantization（BBQ）的功能，對原始的BQ的方式做了進一步提升。

BQ的想法非常直接，將浮點數表示的嵌入向量轉換為二進位制形式（即0和1的序列）來實現資料的壓縮和加速檢索。具體的說，對向量每一維，僅保留1（若原始值為正數）和0（原始值為負數），0 保持不變。如下圖所示：

在計算相似度時，BQ使用的方式是漢明距離（Hamming Distance）。舉個例子，假設有2條向量分別為[0.1, 0.3, -0.1, 0.2, -0.23, 0.12] 和 [0.2, -0.1, 0.21, 0.11, -0.32, 0.01]。在轉為BQ編碼後的結果即為 [1, 1, 0, 1, 0, 1] 和 [1, 0, 1, 1, 0, 1]。在計算兩者距離時，按位取異或，得到的距離結果即為2。在衡量距離時，我們預期是距離越接近，則值越低。所以兩條向量方向若是越接近，則取異或的距離值便越低。

可以看到，BQ除了可以極大地減少記憶體空間（將float32轉位了1bit儲存，減少32倍記憶體空間），還可以極大減少計算複雜度。因為它使用漢明距離（Hamming Distance）進行相似性度量，非常高效，僅涉及簡單的異或操作和位計數，這使得檢索速度得到顯著提升。初次之外，也不難發現，BQ是一種損失編碼，從向量角度來看，它僅保留了向量的方向，損失了向量的長度資訊（或者說在每個維度上的移動距離），並且編碼後無法還原出原始資訊。儘管聽起來這個方式損失資訊較多，但實踐證明在高維向量空間中，它的效果非常好。

4.1. 低維空間BQ的侷限性

在解釋為什麼BQ在高維向量空間中效果非常好之前，我們先看看BQ為什麼在低維空間效果較差。

以二維空間為例，假設原始向量分佈為下圖左部分所示，BQ後的分佈為下圖右部分：

Fig: BQ with 2 dimensional vectors[3]

可以明顯看到，在上圖這個例子中，二維空間中的所有點在BQ後，都分佈到4個“區域”內。也即表示同一區域內的點都存在衝突，且距離關係被抹除。例如對於右上角的區域，每個向量的每一維度的值都大於0，但是最終都會對映為[1,1]的向量。

對於一些特定分佈的向量池，這個問題會更明顯，例如下圖所示的向量分佈：

Fig: Displays some counter-intuitive facts about hamming distances between 2d binary vectors[3]

從直觀上來看，紅色的點和黃色的點距離是非常接近的，但在經過BQ後，黃色點全部量化為[0,1]點，紅色點全部量化為[1,0]點。這時經過漢明距離計算後，它們反而成了最遠的點。

透過這兩個例子可以看到，BQ對於低維向量空間表現非常差，主要原因在於向量在BQ後的衝突機率太高。例如二維空間內，任意2個點經過BQ後的衝突機率為1/4。

4.2. 高維空間的BQ

在低維空間中向量衝突的問題，在隨著維度空間的提升，會帶來極大的緩解。隨著維度的增加，上述提到的“區域”數量會呈指數級增長，減少了向量表示之間的衝突機率。

在高維空間中，例如756維，區域數量變得極其龐大（2^756），這使得即使有數十億或數萬億個向量，它們之間發生碰撞的可能性也非常低。在1.5K維的情況下，區域數量足以容納任何實際數量的向量，而不會有一次碰撞。所以才使得高維空間下的BQ達到了非常好的效果。

需要注意的，在實際使用中，除了1-bit的量化外，也可以使用2-bits和4-bits的量化模式。更多bit的量化一般可以達到更好的準確率，但也需要更多的記憶體進行查詢。

4.3. Rerank

在實際應用中，在對BQ後的結果進行檢索時（仍然可以基於HNSW演算法進行檢索），還會加上一個額外的rerank步驟，以提高效能。具體的說，在檢索時，先以BQ的方式進行（使用查詢向量的BQ表示，檢索向量池裡的BQ表示向量），檢索出topK * rescore_multiplier個候選向量。rescore_multiplier為可調整的引數，此值越高，則考慮的候選向量越多，召回率有可能更高，但也會導致檢索時間延長。拿到這些候選向量後（仍然為BQ表示），再使用查詢向量原始的float32位表示，與候選向量進行點積後做重排序。

二進位制量化在多個領域都有應用，包括但不限於文字檢索、影像識別等，它透過減少資料的維度和複雜度，提高了資料處理的效率和速度。

5. Better Binary Quantization

更好的二進位制量化（BBQ：Better Binary Quantization），在Elasticsearch 8.16和Lucene中被引入。是受到新加坡南洋理工大學研究人員提出的RaBitQ技術的啟發而開發的演算法。

在ElasticSearch （ES）釋出的文件中表示[4]：簡單的二進位制量化會導致大量資訊丟失，為了達到足夠的召回率，需要額外獲取10倍甚至100倍的鄰居進行重排，而這並不理想。基於此限制，ES引入了BBQ，下面是它與簡單BQ的顯著區別：

  1. 所有向量都圍繞一個“質心”進行標準化：這解鎖了一些有利於量化的特性
  2. 儲存多個誤差校正值：部分校正值用於質心標準化，部分用於量化校正
  3. 非對稱量化：在這種方法中，向量自身儲存為單位元值，而查詢向量僅量化到int4。這顯著提升了搜尋質量，且不增加儲存成本
  4. 使用按位操作實現快速搜尋：查詢向量被量化並轉換，以便能夠進行高效的按位操作

下面我們以ES官方部落格中給出的案例[5]，詳細解釋BBQ的過程。

5.1. 構建向量

在將原始向量量化為bit級別的向量時，我們的目標是將這些向量轉換為更小的表示，同時還要：

  1. 提供一個快速估計距離的合理近似值
  2. 對向量在空間中的分佈提供某種保證，以便更好地控制更好地控制召回真實最近鄰所需的資料向量總數

可以透過以下方式實現這兩點：

  1. 將每個向量平移到一個超球體內。例如，如果是二維向量，則這個超球體就是單位圓
  2. 將每個向量固定到圓中每個區域內的一個代表點上
  3. 保留校正因子，以更好地近似向量池中每個向量與查詢向量之間的距離

接下來我們逐步解析這些步驟。

5.2. 圍繞中心點的標準化

為了對每個維度進行劃分，我們需要選擇一箇中心點。為簡化操作，我們將選擇一個點來轉換所有的資料向量。

以二維空間為例，假設我們有3個向量，分別為：

v1: [0.56, 0.82]，v2: [1.23, 0.71]，P3: [-3.28, 2.13]

對每個維度取平均值，得到中心點 c：

X維度: (0.56 + 1.23 + (-3.28)) / 3 = -0.49

Y維度: (0.82 + 0.71 + 2.13) / 3 = 1.22

因此,中心點向量c為: [-0.49, 1.22]，如下圖所示：

然後對原始向量做標準化，也就是做 (v-c)/||v-c|| 的操作。從公式我們可以瞭解到，其結果是一個單位向量，其方向與v-c的方向相同（即兩條向量的夾角方向），但長度為1。

以計算v_c1為例：

  1. 計算v1 – c = [0.56, 0.28] – [-0.49, 1.22] =[1.05, -0.39]
  2. 計算||v1-c|| = 1.13
  3. 計算v_c1 = (v1 – c)/ ||v1-c|| = [0.94, -0.35]

最終得到圍繞中心點做了標準化之後的3個點：

v_c1=[0.94, -0.35]，v_c2= [0.96, -0.28]，v_c3=[-0.95, 0.31]

到目前為止，我們可以看到將原始點投射到了一箇中心點為(0,0)且半徑為1的“圓”上。

5.3. 二進位制量化

在做完標準化後，接下來先對每個向量做BQ，分別得到這些向量所屬的3個區域：r1=[1, 1]，r2=[1, 1]，r3=[0, 0]

接下來，將每個向量固定到每個區域內的一個代表點來完成量化。具體地說，選擇單位圓上與每個座標軸等距的點，即。

將量化後的向量表示為v1_q, v2_q, v3_q，分別將這3個點投射到每個區域的代表點上。以v1為例，計算方式為：

得到結果為：

此時，我們對每個向量得到了BQ的近似值，儘管比較模糊，但可以用於距離比較。當然，也可以看到此時v1和v2在BQ後是同一個點，這個也在前面解釋過，隨著維度的提升，這種“碰撞”的機率會大大降低。

5.4. 儲存誤差矯正值

和原始BQ一樣，這種編碼也損失了大量的資訊，所以我們需要一些額外的資訊來補償這種損失並校正距離估算

為了恢復精度，我們會儲存2個float32值：

  1. 每個原始向量到中心點的距離
  2. 該向量做了標準化後與其量化後的形式的點積

向量到質心的歐氏距離很簡單，在量化每個向量時我們已經計算過：

||v1-c|| = 1.13

||v2-c|| = 1.79

||v3-c|| = 2.92

預先計算每個向量到中心點的距離可以恢復向量做中心化的轉換。同樣，我們也會計算查詢向量到中心點的距離。直觀上，中心點在這裡充當了中介，而不是直接計算查詢向量與資料向量之間的距離。

標準化後向量與其量化後的向量的點積如下：

v_c1 * v1_q = 0.90

v_c2 * v2_q = 0.95

v_c3 * v3_q = 0.89

量化向量與原始標準化後的向量之間的點積作為第二個校正因子，反映了量化向量與其原始位置的偏離程度。

需要注意的是，我們進行此轉換的目的是減少資料向量的總大小並降低向量比較的成本。這些校正因子雖然在我們的二維示例中看起來較大（需要1個float儲存），但隨著向量維度的增加，其影響會變得微不足道。例如，一個1024維的向量如果以float32儲存需要4096位元組，而使用這種位元量化和校正因子後，僅需136位元組（1024 bit = 128位元組，再加上2個float 8個位元組）。

5.5. 查詢

假設查詢向量q=[0.68, -1.72]，在進行查詢時，首先需要將它按照同樣的步驟進行量化：

  1. 計算q – c = [0.68- (-0.49), -1.72-1.22]=[1.17, -2.95]
  2. 計算||q-c|| = 3.17
  3. 計算投射到“圓”後的向量q_c = (q – c)/ ||q-c|| = [0.37, -0.92]

接下來，我們對查詢向量執行標量量化SQ（Scalar Quantization），將其量化到4位元，我們稱這個向量為q_scalar。值得注意的是，我們並沒有將其量化為bit表示，而是保留了一個int4標量量化形式，即以int4位元組陣列的形式儲存q_scalar，用於估算距離。我們可以利用這種非對稱量化方法在不增加儲存成本的情況下保留更多資訊。Scalar Quantization在此不做介紹，有興趣可以參考文件[6]。

進行標量量化（SQ）的步驟為：

  1. 獲取q_scalar向量的維度範圍（即lower和upper標量值）：由於只有2維，所以lower=-0.92，upper=0.37
  2. 計算量化步長：width=(upper-lower)/(24-1)=0.08
  3. 得到SQ結果：q_scalar = |(q_c – min) / width| = |[0.37, -0.92] – [-0.92, -0.92]/0.08| = [15, 0]

由於我們只有兩個維度，量化後的查詢向量現在由int4範圍的上限和下限兩個值組成。

在計算距離時，使用向量池中每個向量與這個查詢向量進行比較。我們透過對量化查詢向量中與給定資料向量共享的每個維度求和來實現這一點。基本上，這就是一個普通的點積運算，只不過使用的是位元和位元組。

例如，計算查詢向量q與向量池中v1的距離時，首先使用的是量化過的q_scalar和量化後的r1進行點積：

q_scalar * r1 = [15, 0] * [1, 1] = 15

然後使用修正因子來展開量化結果，從而更準確地反映估計的距離。其中一個修正因子便是到中心點的距離，前面已經計算過：||q-c|| = 3.17。

5.6. 預估距離

到目前為止，我們已經對查詢向量做了量化，並得到了修正因子。現在可以計算查詢向量q和向量P1的近似距離。使用歐式距離公式並展開：

在這個公式中，大部分值在前面已經計算過，例如||v1-c||。但我們仍然需要計算部分值例如q_c *v_c1。它的值可以基於前面得到的修正因子和量化的二進位制距離度量來合理且快速地估算該值：

後續計算比較複雜，感興趣可以參考原文[5]。最終可以得到近似的距離結果：

est_dist(v1, q) = 2.02

est_dist(v2, q) = 1.15

est_dist(v3, q) = 6.15

相較於量化之前的點之間的距離：

real_dist(v1, q) = 2.55

real_dist(v2, q) = 2.50

real_dist(v3, q) = 5.52

可以看到此方法在量化後的距離遠近上的排名與原始距離排名保持一致。

5.7. Rerank

這些估計的距離確實只是估計值。即使加上額外的修正因子，二進位制量化生成的向量的距離計算也僅是向量之間距離的近似值。在ES的實驗中，透過多階段處理流程，能夠實現較高的召回率。為了獲得高質量的結果，需要對透過二進位制量化返回的合理樣本向量集合進行更精確的距離計算重新排序。一般來說這些候選集的規模可以很小。在大型資料集（>100 萬）中，通常用 100 個或更少的候選集即可實現超過 95% 的高召回率。

在 RaBitQ 中，結果在搜尋操作中會被持續重新排序。在ES的實驗中，為了實現更具可擴充套件性的二進位制量化，我們將重新排序步驟與搜尋分離開來。儘管 RaBitQ 透過在搜尋過程中重新排序能夠維持更優的前 N 名候選列表，但代價是需要不斷載入完整的 float32 向量。這對於某些更大規模的生產類資料集而言是不現實的。

6. ES官方測試BBQ的效能表現

ES官方對BBQ在Lucene和Elasticsearch中的效能表現進行了測試。

Lucene基準測試

測試覆蓋了三個資料集：E5-small、CohereV3和CohereV2。每個資料集的測試結果都顯示了在不同過取樣率（1, 1.5, 2, 3, 4, 5）下的召回率@100。

E5-small

l 500k個向量，基於quora資料集。

l BBQ量化方法在索引時間和記憶體需求上都優於4位和7位量化方法，且召回率表現良好。

量化方法

索引時間

Force Merge時間

記憶體需求

bbq

161.84

42.37

57.6MB

4 bit

215.16

59.98

123.2MB

7 bit

267.13

89.99

219.6MB

raw

249.26

77.81

793.5MB

僅使用單位元精度，就能達到 74% 的召回率。由於維度數量較少，BBQ 的距離計算速度相比最佳化的 int4 並沒有快多少。

CohereV3

l 1M個1024維向量，使用CohereV3模型。

l BBQ在索引時間和記憶體需求上同樣優於4位和7位量化方法，且召回率超過90%。

量化方法

索引時間

強制合併時間

記憶體需求

bbq

338.97

342.61

208MB

4 bit

398.71

480.78

578MB

7 bit

437.63

744.12

1094MB

raw

408.75

798.11

4162MB

在這種情況下，1 位元量化結合 HNSW 僅透過 3 倍過取樣就能實現超過 90% 的召回率。

CohereV2

1M個768維向量，使用CohereV2模型。

BBQ在索引時間和記憶體需求上與4位量化方法相當，且召回率表現良好。

量化方法

索引時間

強制合併時間

記憶體需求

bbq

395.18

411.67

175.9MB

4 bit

463.43

573.63

439.7MB

7 bit

500.59

820.53

833.9MB

raw

493.44

792.04

3132.8MB

在這個基準測試中，BBQ 和 int4 的表現幾乎同步。並且BBQ 僅透過 3 倍過取樣，就能在內積相似度上實現如此高的召回率。


總結

在大規模向量檢索中，使用傳統全精度的hnsw演算法需要消耗大量的記憶體資源。而如果不將索引載入到記憶體，又會導致查詢延遲較高。在乘積量化之後，OpenSearch與ElasticSearch均推出了優秀的二進位制量化的方法，可以以更低的成本實現大規模向量檢索的同時還能保障較好的查詢延遲與召回，解決了乘積量化效果較差的問題。而對於Better Binary Quantization，透過向量標準化、誤差校正和非對稱量化等技術，實現了在保持高召回率的同時顯著減少記憶體使用和提高檢索速度。

透過Elasticsearch官方的測試結果，我們看到了BBQ在不同資料集上的優秀表現，無論是在召回率、索引時間還是記憶體需求方面，BBQ都展現出了超出預期的效能。這些測試不僅驗證了BBQ技術的實用性，也為我們在實際應用中選擇適合的向量檢索技術提供了參考。

總的來說，BBQ作為一種新的量化技術，它在大規模向量檢索中的應用潛力是巨大的。它不僅能夠幫助我們有效地處理和檢索海量資料，還能夠在保持高效率的同時，大幅度降低資源消耗。隨著技術的不斷進步和最佳化，我們可以預見，量化將在未來的資料分析和檢索領域扮演越來越重要的角色。

References

[1] Choose the k-NN algorithm for your billion-scale use case with OpenSearch: https://aws.amazon.com/blogs/big-data/choose-the-k-nn-algorithm-for-your-billion-scale-use-case-with-opensearch/

[2] OpenSearch Binary Quantization: https://opensearch.org/docs/latest/search-plugins/knn/knn-vector-quantization/#binary-quantization

[3] 32x Reduced Memory Usage With Binary Quantization: https://weaviate.io/blog/binary-quantization#the-importance-of-data-distribution-for-bq

[4] Better Binary Quantization (BBQ) in Lucene and Elasticsearch: https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch

[5] Scalar quantization 101：https://www.elastic.co/search-labs/blog/scalar-quantization-101

[6] Better Binary Quantization (BBQ) in Lucene and Elasticsearch: https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch#alright,-show-me-the-numbers

相關文章