大规模向量检索与量化方法

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)/(2^4^-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