本篇文章主要介绍KDD 2020 Applied Data Science Track Papers中的一篇来自Facebook的语义检索文章，Embedding-based Retrieval in Facebook Search。关于样本、模型、训练等细节可以参考：负样本为王：评Facebook的向量化召回算法。本文重点关注其工程实践上的经验，也借此机会对基于PQ量化的近似近邻搜索 (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个向量对应1个粗糙聚类簇标识，通常称为粗糙量化ID； |
过程：具体而言，向量库先构造一个小规模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个聚类簇中心下标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$将向量 $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改写的例子：
1 | or ((and( (term(Cluster-ID, '粗糙聚类中心ID-a')), |
另外，作者强调了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量化的方式来构建索引，调参等。
参考
KDD 2020: Embedding-based Retrieval in Facebook Search
A library for efficient similarity search and clustering of dense vectors
