在推荐、RAG等业务中，如何完成亿级向量的快速检索？

发布日期：2024-12-26 07:15:03 浏览次数： 2589

作者：AI有温度

微信搜一搜，关注“AI有温度”

在推荐、RAG等业务中都需要用到ANN(Approximate Nearest Neighbor)，最常用的方法就是用户Query向量与向量库中的每个向量遍历检索得到答案，俗称暴力计算，这样做精度最高，但效率也最低。

同时还有个非常严重的问题，那就是存储，假设一个向量用1024维的float32（4字节）来表示，即一个向量所占的内存就是，那么一千万个向量所占空间就是。

暴力计算通常用于精准度要求极高的场景，比如公安局人脸搜索。但是在个性化推荐或RAG文本召回场景下，我们要做到快速且大致精准。所以我们需要解决两个问题：

如何快速找到Top相近的（毫秒级）
如何高效的存储向量

本文介绍业界常见的一种ANN算法——IVFPQ是一种用于数据检索的索引方法，它结合了倒排索引（Inverted File）和乘积量化（Product Quantization）的技术。

Inverted File倒排索引

传统文本搜索（比如ES）中的倒排索引首先需要对句子分词，然后记录所有词项的集合（文档库中的每个单词）。接着构建倒排列表，对于每个词项，存储它出现在哪些文档中，以及它在文档中的位置信息。是典型的空间换时间的思路方法。

而向量IVF与传统文本的倒排不同，ANN中倒排的含义仅仅是在聚类中心上面做一步索引，记录库中每个向量属于哪一个簇。

向量IVF本身的原理比较简单，其目的是减少需要计算距离的目标向量的个数，做法就是直接对库里所有向量做KMeans聚类，假设簇心个数为1024，就会将库中已有向量分为1024类。那么每来一个查询向量，首先计算查询向量与1024个聚类簇心的距离，然后选择距离最近的top N个簇，只计算查询向量与top N个簇下向量的距离。

而这样做，只是减少了计算次数，并没有提升计算效率与减少向量的存储。

Product Quantization

这就需要用到PQ（乘积量化）方法，它的核心思想是将全样本的距离计算，转化到聚类（子空间）中心的距离计算。PQ一般分为两步操作，第一步是聚类，第二部是量化编码。

以N个128维的向量库为例，有一个超参M，要将这些128的向量切分为N个128/M维的向量，这里我们使M=4，就会有N*4个子段是32维的向量。

对于一个子空间（这里有4个子空间），通过聚类算法将每个子空间中的向量聚成K=256个类，每个聚类中心就是一个维子向量，他们由一个id来表示，叫做clusterid，一个子空间中所有的clusterid，构造了一个属于当前子空间的codebook（码本）。

重点操作来了：对于库中每一个向量，都可以计算得出离它最近的簇心 id，那么4段向量就对应了4个簇心 id，于是我们就用这个聚类中心的clusterid来代表这个子向量，最终N个128维的向量就可以用N个4维的向量进行表示，这样就大幅度的降低了向量的储存空间。

以上就完成了PQ Pre-train的过程。在向量库检索阶段，主要有两种计算方式，分别是SDC以及ADC。

symmetric（对称的）：图中y就是向量库里面的向量(注意，y是经过量化的，就是上文中说的那个用数字id替代向量)。来一个Query，我们把它转换为128维的向量，首先我们将它分为4段32维向量，然后根据码表计算每一段向量与之前预训练好的簇心的距离，看每段向量与哪个簇更近，就用这个id代表这个向量，那么Query就能转换为4个id表示的向量，即为x。图中的查询向量x和库中向量y都是用id表示的向量，这时可以有两种操作：

直接使用id表示的向量做汉明距离，即计算两个相同长度的字符串在相同位置上不同字符的个数。使用汉明距离可以直接做召回，然后再将召回的向量做余弦计算或使用双塔模型进一步精排
将x转为q(x)，将y转为q(y)。这里x与y是用id表示的向量（4维）。函数q方法是进行升维，使用x和y向量中每个id对应的聚类中心向量（32维）进行拼接来升维表示查询向量和库中向量，那么q(x)和q(y)就是两个完整的子向量（128维），再使用q(x)和q(y)计算距离

为什么名字叫symmetric呢？因为他俩都是用对应的聚类中心来计算距离，所以是对称的。SDC有一个很大的优点，两两聚类中心之间的距离，是我们在构建向量库的时候就可以知道的，完全可以离线就计算好，在线要计算q(x)和q(y)的距离的时候，直接查表，提升了在线query的效率。但是，他的误差也比ADC来的大，因为有x到q(x)、y到q(y)两个量化和转换的误差。

asymmetric（不对称的）：来一个Query，我们把它转换为128维的向量，首先我们将它分为4段32维向量，然后根据码表计算每一段向量与之前预训练好的簇心的距离，就会得到一个4*256的表，记录了在每个子空间内，每段向量与每个簇之间的距离。

PQ优化的点就在这里，在计算查询向量和向量库向量的距离的时候，向量库的向量已经被量化成256个簇心 id，而查询向量的4段子向量与各自的256个簇心距离已经预计算好了，所以在计算两个向量的距离时只用查表4次。

比如的库里的某个向量被量化成了[124, 56, 132, 222]，那么首先查看刚计算出的查询向量的每段向量与每个簇之间的距离表，得到查询向量第一段子向量与其id为124的簇心的距离，然后再查表得到查询向量第二段子向量与其id为56的簇心的距离，最后就可以得到四个距离d1、d2、d3、d4，查询向量跟库里一个向量的距离就是d = d1+d2+d3+d4。

所以在提出的例子里面，使用PQ只用4×256次128维向量距离计算，再加上4xN次查表，而最原始的暴力计算则有N次128维向量距离计算，很显然随着向量个数N的增加，后者相较于前者会越来越耗时。

与SDC不同的是，ADC的精确度更高，因为只有y和q(y)这一个量化误差。当然必须要在线计算，计算速度不如SDC。