语义向量召回之 ANN 检索

作者：蘑菇先生学习记

本篇文章主要介绍 KDD 2020 Applied Data Science Track Papers 中的一篇来自 Facebook 的语义检索文章，Embedding-based Retrieval in Facebook Search。关于样本、模型、训练等细节可以参考：负样本为王：评 Facebook 的向量化召回算法。本文重点关注其工程实践上的经验，也借此机会对基于 PQ 量化的近似近邻搜索 (ANN) 的原理做个梳理。

架构总览

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

EBR 系统架构图示意图

左上角是查询处理模块，是双塔模型中的 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 维向量和集合 ,需要找到与距离最短的 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；一个向量，可以划分为 m 组向量，第i组向量形如：[x_{i\_1},x_{i\_2},…,x_{i\_D^{\star}}]，每组的 codebook 为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 个分块聚类中心向量】。示意图如下：

PQ 量化过程示意图

如上图所示，将原始向量等分为 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，量化器为 qcqc。这个就是所谓的粗糙量化，或者称为粗糙聚类。接着，每个向量 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。

做个小结，量化的结果：

本文地址：语义向量召回之 ANN 检索
本文版权归作者和AIQ共有，欢迎转载，但未经作者同意必须保留此段声明，且在文章页面明显位置给出

聚类量化	乘积量化	粗糙量化 + 残差量化
k 个聚类簇中心下标 id	m x k个分块聚类中心下标组成的code*	k 个聚类簇中心下标 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将向量 yy 映射到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 量化的方式来构建索引，调参等。