Press "Enter" to skip to content

语义向量召回之ANN检索

首先从总体上预览一下paper提出的EBR系统的架构 (Embedding-based Retrieval System):

 

左上角是 查询处理模块 ,是双塔模型中的query-side embedding model;
右上角 索引构建模块 ,是用document embedding model推断doc embedding后构造向量索引。构造过程中,先拓展doc的metadata,加入doc embedding,并导入doc的正排索引中 (比如用于ranking的特征);同时,通过向量量化技术来降低索引存储和计算距离代价,并将 量化的结果 存在倒排索引中。这一部分也是paper中关注的重点,下文会重点介绍。
左侧中间部分是 检索模块 ,拿到query embedding后,可以通过精心为语义召回设计的 NN算子(Near Neighbor Operator) 计算距离并进行 语义召回 。召回的过程中,既支持原来的布尔匹配召回,也支持向量语义召回,还支持混合召回。这部分下文也会重点介绍。
左下角是 排序模块 ,得到召回结果后,再经过排序模块进行实时排序,排序时会利用到 embedding特征 。

paper中涉及工程实现细节的主要是 索引构建模块 和 检索模块 。下文我会先介绍索引构建模块,这里头涉及很多向量量化的概念,比如PQ量化、粗糙量化、残差量化等,我会先介绍一些量化的背景知识,核心的索引过程和搜索过程,然后介绍在线检索模块,认识基于构建好的向量索引,线上是 如何运转 和 实现召回 的;接着探讨paper中提到的工程优化点和调参经验。最后,做个总结。

 

索引构建模块

 

基于Product Quantization的近似最近邻搜索,核心的一些问题预览一下:

问题描述,解决什幺问题?
传统方法存在什幺问题?
什幺是向量量化?为什幺要量化?量化场景下距离怎幺计算?
什幺是乘积量化?为什幺要乘积量化?
什幺是粗糙量化+残差量化?为什幺要残差量化?
索引过程?搜索过程

问题描述

 

给定D维向量$x$和集合$\Gamma ={ y_1,y_2 … y_N }$ ,需要找到与$x$距离最短的 k 个最近邻。距离的衡量可以是欧式距离、余弦距离等。

 

暴力搜索问题

 

如果以最粗暴的方法进行穷举搜索,构造 距离矩阵 的复杂度为:$O(D N^2)$。从距离矩阵中查找到 k个最近邻 ,最小堆,则复杂度为$O((N-k)\log k)$。

 

举个例子:假设$N=2000W, k=1000$,则构造的距离矩阵包含$400T$个元素,假设每个距离值32bit,至少占用1600TB空间,构建距离矩阵的时间是$10^{17}$量级,查找Top-K搜索时间是$10^9$量级。

 

向量量化

 

所谓向量量化,就是将原来无限的空间$R^D$映射到一个有限的 向量集合 $\mathcal{C} = {\boldsymbol{c}_i, i\in[1,l]}$ 中,其中$||\mathcal{C}||=l$是一个自然数。将这个从 $R^D$ 到集合$\mathcal{C}$的函数记为$q$,则$\forall q(y) \in \mathcal{C}$,在信息论中称$\mathcal{C}$ 为codebook。即:通过$q(y)$来 近似 代表$y$。

 

聚类量化

 

最常用的就是k-means聚类,通过聚类后的 聚类中心向量 来近似 量化 原始的向量。 聚类中心的个数 即为 codebook大小 。因为量化的存在, 任意两个向量之间的距离 可以通过对应的 量化向量的距离 进行近似,也就是聚类中心向量的距离。因为聚类中心的个数 小了 很多,故计算距离的复杂度也下降了很多。(显然,这种方式太粗糙了,误差很大。除非聚类数非常大,极端情况下,聚类数等于样本数时,每个样本一个聚类簇,此时无误差,但是没有起到减少计算复杂度的目的)

 

聚类量化的结果:产出聚类中心向量的过程对应train的过程的一部分。量化后,每个向量都可以用 聚类簇中心下标ID 来 标识 ,根据ID可以获取聚类簇的中心向量。

 

 

如上图所示,N个向量,通过聚类量化产生多个聚类中心,每个向量属于某个聚类簇中,那幺就用该聚类簇对应的中心向量来量化该向量,可以用聚类簇中心对应的下标ID来表示,比如:C1量化vec1。

 

Faiss中对应的实现是 IndexIVFFlat。

 

乘积量化

 

动机:即:PQ量化,很多时候我们向量不同部分之间的 分布不同 的,比如下图(3段向量),因此可以考虑对 向量分段 ,并 分别进行分块量化。 这个只是直觉原因,本质原因下文讲。

 

 

乘积量化定义:将向量分成$m$个不同的部分,对每个部分进行向量量化,假设平均划分,则每个部分的维度大小为$D^{\star}=D/m$

 

一个向量$[x_1,x_2,…,x_{D^{\star}},…,x_{D-D^{\star}+1},…,x_D]$,可以划分为m组向量,第$i$组向量形如:$[x_{i_1},x_{i_2},…,x_{i_D^{\star}}]$,每组的codebook为$\mathcal{C_i}$,对应的量化器记为$q_i$, ($\forall q_i(x_{i^{\star}}) \in \mathcal{C_i}$)。则最终的全局codebook就是 $\mathcal{C} = \mathcal{C_1} * \mathcal{C_2} … * \mathcal{C_m}$,乘积量化的名称也来源于此。

 

分块量化也可以采取聚类量化来实现,则 分块聚类中心的个数 即为 分块codebook 的大小。相当于在这个方法下,对每个向量,有【 m个分块向量】 来量化它,即: 【m个分块聚类中心向量】 。 示意图如下:

 

如上图所示,将原始向量等分为m组分块向量,每组都进行聚类量化,那幺每个向量就有m个分块聚类中心向量来表示,比如vec1用$C_{1-1}$, …, $C_{m-2}$共m个向量来量化。

 

乘积量化的好处:假设每个子codebook大小一样,记做$||\mathcal{C}_i||=k^{\star}$,排列组合一下,那幺相当于能表达的向量空间容量是这幺大,$||k^{\star}||^m$,但是只需要$m k^{\star}$的codebook空间,这也是乘积量化大幅度降低空间占用的本质原因。PQ量化原论文中给出的经验取值是$k^{\star}=256,m=8$,即:分成8块,每个分块的codebook大小为256,对应的向量空间大小为约$256^8=2^{64}≈1.8 \times 10^{19}$,能够表达的向量个数足够大了。

 

乘积量化结果:$m \times k^{\star}$个分块聚类中心向量。每个向量可以用m个 分块聚类簇中心下标 ID来标识 ,所有ID连起来称为code 。假设每块的聚类中心个数为256,则需要8bits,即1byte标识某分块下哪个聚类中心,m块则需要m bytes,即code大小为m bytes。

 

粗糙量化+残差量化

 

核心思想:层次化量化,这个也是Faiss中PQ索引的实现方式。其中粗糙量化使用 聚类量化 ,用划分到的 粗糙聚类簇的中心向量 来 粗粒度 量化该向量,该结果存在较大的误差;接着对残差结果进行 细粒度 乘积量化。这样的话,误差就小了。

 

1. 总体上,每个向量先进行粗糙量化划分到某个粗糙聚类簇里,1个向量对应1个粗糙聚类簇标识,通常称为粗糙量化ID;
2. 然后计算残差向量,即:向量-聚类簇中心向量,再对该残差向量进行分块,并进行细粒度分块残差量化。残差量化的时候,每一块对应一个细粒度聚类簇,1个向量M块,则对应M个细粒度聚类簇标识,通常称为残差量化code。
3. 为什幺用残差量化?原始向量可能会有特别大的分布差异/不平衡,也就是说可能聚类后,不同聚类簇分布得非常分散,每个簇所拥有的样本数极度不平衡。但是通过残差化后,即:每个样本向量减去所属的聚类簇中心向量后,残差向量之间的差异就不太大了,然后再对残差向量进行量化,就能更精确的近似原向量。

 

过程:具体而言,向量库先构造一个小规模codebook $\mathcal{C}_c$,量化器为$q_c$。这个就是所谓的 粗糙量化,或者称为粗糙聚类 。接着,每个向量y都会有一个残差$r(y)=y−q_c(y)$。具体而言:记残差量化步骤的量化器为$q_p$,则 y 可以通过$q_c(y)+q_p(y−q_c(y))$来表示。

 

$$

 

y = q_c(y) + q_p(y-q_c(y)) = y_C + y_R

 

$$

 

其中, $y_C$是粗糙量化结果,$y_R$是残差量化结果。

 

这样的话,【 查询向量x】和y之间的距离 :

 

$$

 

d(x,y) = ||x-y||^2=||x-y_C-y_R||^2=\underbrace{||x-y_C||^2}_{\text{term 1}} +

 

\underbrace{||y_R||^2 + 2 y_C y_R}_{\text{term 2}} – \underbrace{2 x y_R}_{\text{term 3}}

 

$$

term 2 与查询向量x 无关, 可以提前计算好;
term 1 求x和y的粗糙量化向量的欧式距离,最多计算$O(k)$次,$k$为粗糙聚类中心个数。
term 3 求x和y的残差量化向量的内积,遍历所有簇的时候,最多计算$O(mk^{\star})$ 次,$mk^{\star}$为分块聚类中心向量。

对应Faiss中的实现是 IndexIVFPQ。

 

做个小结,量化的结果:

 

 

聚类量化乘积量化粗糙量化+残差量化
k个聚类簇中心下标idm x k* 个分块聚类中心下标组成的 codek个聚类簇中心下标ID,m x k* 个分块聚类中心下标组成的 code

 

也就是说,为了表示每个doc的 量化结果 ,可以为doc可以添加两种 结构化字段 : 粗糙量化id, 残差量化code,用于实时检索使用。

 

倒排索引

 

索引结构

 

 

建立从粗粒度聚类中心id 到 doc的映射关系,其中doc的信息包括:向量id, 向量的残差量化code。每个doc通过粗糙聚类中心id和残差量化code就能知道原始向量如何映射到量化向量。

 

整个倒排索引不需要存储原始向量本身,索引结构存储的内容:

 

存储标识: 粗糙量化id , doc id 和残差量化 code。

空间占用很小。m bytes 残差量化code,即:code_size或者pq_bytes。这个数越大,那幺细粒度聚类簇越大,则精度越高。

基于id,可以找到对应的粗粒度量化向量 (共k个) ; 基于code 可以找到对应 细粒度残差量化向量 (共 m x k* 个)。

 

索引构建过程

 

搜索场景中,y可以理解为doc。

 

通过粗糙量化器 $q_c$将向量 $y$ 映射到$q_c(y)$,即粗糙量化。这样就知道挂到哪个链表上了。
计算残差 $r(y) = y-q_c(y)$
将残差 $r(y)$量化到 $q_p(r(y))$,其中包含了 m 个分组,每个分组有对应的一个细粒度聚类簇ID,用1byte表示,则共m bytes,对应code标识。
构造一个 $id|code$ entry,其中id是doc的标识,code是残差量化标识。

搜索过程

 

搜索场景中,x可以理解为查询向量。

 

 

通过粗糙量化器来量化 查询向量x ,即:找到离x最近的w个粗糙聚类簇。实际上是用于 限定搜索的范围 ,只搜索w个粗糙聚类簇ID索引下的向量。w是个超参数,即nprobes,量化的结果 对应term 1。

 

选定某个粗糙聚类簇,

 

计算 x的残差向量 和该粗糙聚类簇下的 k* (默认即256)个中心点向量的内积。对应term 3,计算时间复杂度 O(m x D/m x k*) = O(Dk*),记录下来,下一步查表用。

 

遍历该聚类簇下的doc文档, 计算距离时 ,实际上全是 查表操作 ,term 2是提前算的,term 1粗糙量化时算的,term 3上一步算的。查表时间复杂度实际上是 O(m)

w个粗糙聚类簇,搜索时间复杂度为 O(w(Dk* + m)), 另外,返回top-k的话,要加上最小堆排序时间。

 

检索模块

 

布尔检索

 

Facebook在自研的检索引擎 Unicorn 中支持了第一代的近邻搜索。

 

首先简单介绍下Unicorn。Unicorn原本的检索方式主要以Term布尔匹配为主。

 

doc侧:以Term词袋的方式来表示doc,会基于Term的语义来区分命名空间。Term上还可以添加一些term特定的 payload 信息。举例:John living in Seattle会表示成【 text : john and location : seattle】。

 

query侧:定义 Term-level 的布尔表达式来表示query。举例:下述query主要返回doc的文本中有john和smithe字眼,并住在seattle和menlo_park的用户。

 

向量检索

 

为了支持embedding,需要扩展 doc和query的表示

doc侧:添加结构化字段 key 来拓展doc词袋表示,形如:


比如:key=model-141795009,代表了产出doc embedding model的版本。设置该字段方便部署多种embedding版本。
query侧:添加 nn算子 ,即: nn

: radius
,使用时,通过计算query embedding和

模型产出的doc embedding的距离,来匹配距离在指定radius内的
文档 。radius此处起到
阈值约束 的作用。

索引和计算向量距离时,Facebook将 向量索引和向量在线检索 通过某种方式转化成上述已有的 布尔检索语言 ,很巧妙,可以完美融入现有的 布尔检索系统, 而不需要重新写一套系统。

 

先做个对应关系。

 

 

布尔匹配检索语义向量检索
Term粗糙聚类中心ID 标识,Cluster-ID
Payload细粒度聚类中心CODE 标识, Cluster-Code

 

doc侧:每个doc embedding会被 量化 并转成一个 term 和一个 payload ,相当于是两个doc的结构化字段,完美兼容已有的检索系统Unicorn的设计。

term,其实就是用于标识倒排索引中,该doc属于哪个 粗粒度聚类簇 (用于粗糙聚类量化)
payload,用于标识每个分块向量下的 细粒度聚类簇 (用于残差量化)

query侧:

 

term: (nn) 重写成和粗糙聚类相关的 term。

重写规则:计算query embedding和所有粗糙聚类中心距离,选出nprobes个最近的,用粗糙聚类中心ID来标识 (和doc的结构化字段 Cluster-ID 进行比较),不同聚类中心对应的Term之间的关系就是or的关系。

 

payload: 对query进行残差量化,得到满足 radius约束条件 的细粒度聚类中心,用code标识 。

重写规则:对每个粗糙聚类簇,计算query embedding和其细粒度聚类中心的距离,距离 满足<radius约束 的细粒度聚类中心对应的Code 取值记录一下 (和doc的 结构化字段Cluster-Code 进行比较),Code之间是OR关系,但是Code和ID是AND关系。

 

举个query改写的例子:

 

or ((and( (term(Cluster-ID, '粗糙聚类中心ID-a')), 
          (or (term(Cluster-Code, '残差聚类中心ID-a1'), term(Cluster-Code, '残差聚类中心ID-a3'),...)))
    ), 
    (and( (term(Cluster-ID, '粗糙聚类中心ID-b')), 
          (or (term(Cluster-Code, '残差聚类中心ID-b1'), term(Cluster-Code, '残差聚类中心ID-b4'),...)))
    ),
    ...
   )

 

另外,作者强调了radius-mode和topk-mode的差异,radius方式的性能和质量更高。radius是一种 受限制 NN搜索。top-K需要扫描整个索引库来找 Top-K结果 。radius性能更好,但是需要确定radius值。

 

正是因为在已有的布尔查询语言上融入目前的EBR语义检索,使得混合检索的实现非常方便。特别是在模糊匹配场景、或者query存在错误的场景。举个例子:

 

 

上述查询,用户把smith->smithe了,这样基于term匹配的话,可能会没有结果。但是基于nn检索,只要满足query embedding和doc embedding的余弦相似性小于0.24的话,就会有结果。nprobe是超参数,扫描的粗糙聚类中心个数。

 

调参经验

 

模型改进不大的时候,多调调参数,会有奇效。

调节粗糙量化聚类簇数量num_cluster和实时查询时扫描的聚类簇数量nprobe 。由于不同聚类中心的向量个数可能存在很不平衡的现象,对于两个对比实验(比如对比不同构建索引的方法),即使num_cluster和nprobe设置完全一致,可能两个实验扫描的文档数量是不一样的,因此要监控扫描文档数量,并通过调节这两个参数,来保证不同对比方法扫描的文档数量一致。
多尝试使用乘积量化 :包括原生PQ,OPQ,PQ with PCA等。能够显着 降低索引存储空间存储复杂度、查询时间复杂度 。其中, pq_bytes,即残差量化codes的大小,很重要 。决定了残差量化分块聚类中心的个数,个数越大越精确,比如m=8,k*=256时,pq_byte=1byte x 8= 8 byte。paper中建议采用d/4,d是向量的维度数,假设64维度,则d/4=16,是默认值的2倍。
多进行 在线调参 ,nprobe, num_clusters,pq_bytes。虽然离线能够直观感受perf vs recall之间的tradeoff,但是多部署几套参数进行在线调整,对于理解性能因素对EBR检索系统的影响会更加直观。这对于减少 容量开销 和 离线参数搜索成本 很有用。

总结

 

此次分享主要介绍了Facebook在KDD 2020发表的文章中的工程实践经验。首先从全局总览其系统架构,然后针对索引构建模块和实时检索模块展开讨论。其中,索引构建模块主要介绍了ANN中 向量量化 方法的背景知识。实时检索模块主要介绍 系统实现 ,如何将向量检索通过 query改写规则 等融入现有的布尔检索系统中。最后介绍了一些 调参的经验 。

 

做向量召回的初期可以先重点关注模型层面上的优化,索引上也可以先采用简单的聚类量化的索引构建方式。当向量召回优化到一定程度时,如果想进一步提升性能和召回率,可以考虑借鉴文中的一些经验,比如以PQ量化的方式来构建索引,调参等。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注