离散与提炼——一些关于向量召回算法优化方法的思考

 

✏️  作者介绍:

周语馨,高级云智能工程师

 

最近做的很多向量召回的相关工作,主要集中在优化 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。至此,搜索的第一步可以用数学语言表达:

Li=|x−zi|−ei

 

把 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 项满足Si. lower_bound ≥ topK. T ,那么第 i 项后续的任一项都大于 T,即:

 Sj.lower_bound≥Si.lower_bound≥topK.T,j≥i

 

又因为 L 作为 d 的下界,有:

dSj.label  ≥Sj.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,假设其所在聚类的中心点为 Cy 。使用 ri = yi − Cy ,即 yi 的残差作为“基于线性变换的离散化”的输入,得到 zi 与 ei。该过程的数学本质是,以每个聚类中心为原点建立一个坐标系,在该“局部坐标系”中对属于该聚类的点做离散化。如此即可解决信息损失的问题。

当给定 x 时,按照 IVF 算法找出最近的 nprobe 个聚类。对于每一个聚类,计算 x 的残差 q =  x − Cx  ,然后在该聚类中计算 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 | 源码

http://milvus.io | 官网

milvusio.slack.com | Slack 社区

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

zilliz.blog.csdn.net | CSDN 博客

space.bilibili.com/478166626 | Bilibili

已标记关键词 清除标记