Product Quantization乘积量化

简介

乘积量化（PQ）算法是和VLAD算法是由法国INRIA实验室一同提出来的，为的是加快图像的检索速度，所以它是一种检索算法，在矢量量化（Vector Quantization，VQ）的基础上发展而来。

综合转载以下文章：

• [乘积量化（Product Quantization）]
• 理解 product quantization 算法
• 最近邻搜索之乘积量化（Product Quantizer）理解（一）
• 最近邻搜索之乘积量化 理解（二）
• 图像检索之product quantization 算法解析

引言

Product quantization，国内有人直译为乘积量化，这里的乘积是指笛卡尔积（Cartesian product），意思是指把原来的向量空间分解为若干个低维向量空间的笛卡尔积，并对分解得到的低维向量空间分别做量化（quantization）。这样每个向量就能由多个低维空间的量化code组合表示。为简洁描述起见，下文用PQ作为product quantization的简称。

The idea is to decomposes the space into a Cartesian product of low dimensional subspaces and to quantize each subspace separately. A vector is represented by a short code composed of its subspace quantization indices.

2011年，Herve Jegou等学者在PAMI上发表了PQ方法的第一篇正式paper[1]，用于解决相似搜索问题（similarity search）或者也可以说是近邻搜索（nearest neighbor search）问题。其实这几位作者在2009年的INRIA（即法国国家信息与自动化研究所）的技术报告上已经发表PQ方法。这里插一段题外话，[1]的一作Herve Jegou和二作Matthijs Douze均在2015年跳槽去了Facebook AI research，并在今年3月份合作开源了Faiss相似搜索工具[4]。

近几年，深度学习技术被广泛用于图像识别、语音识别、自然语言处理等领域，能够把每个实体（图像、语音、文本）转换为对应的embedding向量。一般来说，相似的实体转换得到的embedding向量也是相似的。对于相似搜索问题，最简单的想法是暴力穷举法，如果全部实体的个数是$n$，$n$是千万量级甚至是上亿的规模，每个实体对应的向量是$D$，那么当要从这个实体集合中寻找某个实体的相似实体，暴力穷举的计算复杂度是$O(n×D)$，这是一个非常大的计算量，该方法显然不可取。所以对大数据量下高维度数据的相似搜索场景，我们就需要一些高效的相似搜索技术，而PQ就是其中一类方法。

PQ是一种量化（quantization）方法，本质上是数据的一种压缩表达方法（其实通信学科的一个主要研究工作就是研究信号的压缩表达），所以该方法除了可以用在相似搜索外，还可以用于模型压缩，特别是深度神经网络的模型压缩上。由于相似搜索不仅要考虑如何量化的问题，还要考虑如何检索（search）的问题，而模型压缩可能更主要的是考虑如何量化的问题，不用太关注如何检索这个问题，所以这篇文章会主要站在相似搜索上的应用来介绍PQ方法。至于模型压缩，可以找找近几年研究神经网络模型压缩的paper或者一些互联网公司（比如百度, Snap等）发出的一些资料[3]。

相似搜索的若干种方法

参考文献[5][6]很好的总结了相似搜索的几类方法，这里简要总结几个核心点。可以将方法分为三大类：

• 基于树的方法
  • KD树是其下的经典算法。一般而言，在空间维度比较低时，KD树的查找性能还是比较高效的；但当空间维度较高时，该方法会退化为暴力枚举，性能较差，这时一般会采用下面的哈希方法或者矢量量化方法。
• 哈希方法
  • LSH(Locality-Sensitive Hashing)是其下的代表算法。文献[7]是一篇非常好的LSH入门资料。
  • 对于小数据集和中规模的数据集(几个million-几十个million)，基于LSH的方法的效果和性能都还不错。这方面有2个开源工具FALCONN和NMSLIB。
• 矢量量化方法
  • 矢量量化方法，即vector quantization。在矢量量化编码中，关键是码本的建立和码字搜索算法。比如常见的聚类算法，就是一种矢量量化方法。而在相似搜索中，向量量化方法又以PQ方法最为典型。
  • 对于大规模数据集(几百个million以上)，基于矢量量化的方法是一个明智的选择，可以用Faiss开源工具。

Product Quantization

背景知识

首先介绍一下向量量化的一些背景知识。向量量化的目标是为了减少空间结构的复杂度。因为经过量化，任何空间中的点都可以用有限的几个code word来表示。假设存在量化函数$q$，和词汇$C$，则对于空间中的任何一个向量$x$都有，

$$q(x)\in \mathcal{C}=\{c_i;i\in \mathcal{I}\}$$

其中$i=1,2,…,k$。即codebook中code word的数目。

所以一个code word中所包含的特征向量集合($\mathcal{V}_i$)可以用下式表示

$$\mathcal{V}_i \triangleq \{x\in \mathbb{R}^D:q(x)=c_i\}$$

因此量化过程直观上的理解就是对特征空间进行分割。为了对量化进行评价，我们引入均方误差(Mean Squared Error)：

$$MSE(q)=\mathbb{E}_X[d(q(x),x)^2]=\int p(x)d(q(x),x)^2dx$$

式中的$p(X)$表示$X$的分布函数。因此对于特定的$x$，$p(x)$表示$X$落在该点的概率。

可以看到$MSE(q)$是通过计算空间中所有点到其量化点的距离的期望值对量化函数$q$进行评价。可以很直观地看出，对于所有的$q$，$MSE(q)$的值越小，说明此量化函数的量化误差越小，量化效果越好。

既然量化函数是可以进行量化评价的，那么肯定就存在一个量化函数，对于一些特定的点可以训练得到最小的MSE。亦即存在最优的量化函数。下面利用Lloyd optimality conditions定义量化函数的最优化条件。

$$q(x)=arg\space min_{c_i\in \mathcal{C}}d(x,c_i)\\ c_i=\mathbb{E}_X[x|i]=\int_{\mathcal{V}_i} p(x)xdx$$

• 式一是对量化函数的描述，表示向量x一定要量化到离他最近的code word上面。

• 式二是对code word的定义。可以看出这是一个需要反复修正的训练过程。

可以把满足以上两个最优化条件的量化算法称为Lloyd 量化器。我们熟知的k-means算法就是一个Lloyd量化器。它是一个接近最优化的量化算法，难怪如此深入人心。

为了后面的讨论，再定义均方失真函数，

$$\xi(q,c_i)=\frac{1}{p_i} \int_{\mathcal{V}_i}d(x,q(x))^2p(x)dx$$

式子中，$p_i(x)$表示$x$被量化到$c_i$上面的概率分布。

Product Quantizer

product quantization是一种常用的编码方法。它允许将一个向量的各个部分进行分开量化。

$$\underbrace{x_1,...,x_{D^*}}_{u_1(x)},...,\underbrace{x_{D-D^*+1},...,x_D}_{u_m(x)}\\ \rightarrow q_1(u_1(x)),...,q_m(u_m(x))$$

如图，将一个向量分成$m$等份进行分开量化。这样，每个量化函数都相对简单，因为它只需要处理维数比较少的子向量即可。算法还分别为每个子量化器分配了一个codebook，假设为量化器$q_i$，分配的codebook是$c_i$。则原向量最终的量化空间将是这些codebook的笛卡尔乘积，

$$\mathcal{C}=\mathcal{C}_1 \times ... \times \mathcal{C}_m$$

假设每个子量化器都有$k^*$个code word，则对于原向量，可选的codeword一共有

$$k=(k^ \ast)^m$$

可以看到$m=1$或者$m=D$是PQ算法的2种极端情况，对$m=1$，PQ算法就回退到vector quantization，对$m=D$，PQ算法相当于对原始向量的每一维都用kmeans算出码本。

product quantization的优势在于，通过对将较小的几个code book连接形成一个规模足够大的codebook。而且对于小的codebook，训练复杂度相对直接训练大的codebook是低得多的。所以根据这个原理，直接存储$\mathcal{C}$codebook是没必要的，且相反会降低最终的量化速率。因此在存储的时候，我们采用的是分别存储$m$个大小为$k^\ast$的codebook的方式。

笛卡儿积

一个个细胞样的叫做Voronoi cell，被定义为：

$$\mathcal{V}_i \triangleq \{y\in R^d:q(y)=c_i=(c_{i,1},c_{i,2},...,c_{i,m})\}$$

上式中每个$c_{i,j}$可能取值一个有$K_s$个，记为集合$c^i$。

上图中共有16个cell，分成$K=16$个类别，$m=4$，$K_s=2$，它是$R^d$，$d=4$空间的一个划分，这就意味着任何两个$\mathcal{V}_i$之间是正交的子空间。整个空间可表示为：

$$\mathcal{V}=\mathcal{V}_1 \cup \dotsi \cup \mathcal{V}_m, \cup是并运算\\ \mathbb{C}=c^1 \times \dotsi \times c^m, \mathbb{C}是类别空间$$

乘积量化的优势在于能从少量的类别通过笛卡儿积运算得到大量的类别。就像是一个平面空间与一个垂直平面的线笛卡儿积得到一个三维空间样。

所以剩下的问题就是如何确定参数$k^\ast$和$m$了。$k^\ast$和$m$是通过MSE函数为指标确定的，MSE函数值越小越好。因为现在量化是分步进行，最后再将量化结果整合，因此需要重新定义MSE函数，

$$MSE(q)=\sum_jMSE(q_j)$$

下图是论文中描绘的MSE函数值与($k^\ast$,$m$)的函数关系图

其中code length可由$m$,$k^\ast$计算得到: $codelength = mlog2(k^\ast)$。

通过图可以看出，对于一个 codelength，$m$越小，$k^\ast$越大，则MSE的值会越小。比如当m=1时，算法就进化为满足最优化条件的k-means算法。只是由于k-means算法对高维复杂度比较高，因此不予以考虑。所以$m$和$k^\ast$的选取是一个trade off的问题。

利用product quantization进行检索

正如前面说的一样，检索本质上是一个比较欧式距离的问题。因此现在主要看看如何通过量化后的code word 对向量进行efficient的比较。

如图所示，计算待查询向量和数据库中的向量的距离有两种方法。

一种是，先将待查询向量进行量化，然后计算对应的量化子和数据库中的量化子的距离（SDC）：

$$\hat d(x,y)=d(q(x),q(y))=\sqrt{\sum_jd(q_j(x),q_j(y))^2}$$

1. 为提高计算速度，一般会提前算好$d(c_{j,i}, c_{j,i’})^2$，然后在检索时就是查表，以$O(1)$的复杂度查出结果。
2. $\hat{d}(x,y)$是$d(x,y)$的近似计算，一般会先用相似计算方法选出top N近邻，然后再做rerank以拿到最终的近邻排序结果。

另一种是不需要对待查询向量进行量化，而直接计算它和数据库中的量化子的距离（ADC）：

$$\tilde d(x,y)=d(x,q(y))=\sqrt{\sum_jd(u_j(x),q_j(u_j(y)))^2}$$

1. 为提高计算速度，一般会在检索前提前算好$d(u_j(x), c_{j,i})^2 : j=1,\cdots, m, i=1,\cdots, k^\ast$，然后在检索时就是查表，以$O(1)$的复杂度查出结果。
2. $\widetilde{d}(x,y)$也是$d(x,y)$的近似计算，与SDC类似，一般会先用相似计算方法选出top N近邻，然后再做rerank以拿到最终的近邻排序结果。

下图对比了SDC算法和ADC算法的各阶段复杂度，当$n>k^\ast D^\ast$时，计算瓶颈存在于公式的计算上，它们的复杂度都是$O(n \times m)$。

可以看到SDC相对于ADC，只有一个优点，就是不需要在内存中存储query vector，因为它已经被量化到codeword上了，直接用这个code word表示它就可以，而code code本身是存在内存中的。所以下面只考虑ADC算法。

但是需要注意的是，这里进行距离计算时还是需要遍历数据库中的$mk^\ast$个code book的。只是相对来说，这个值要小的多。下面着重分析一下ADC算法的距离误差。(PS：这个分析方法具有普遍性，并不是只适合于分析product quantizer)

误差分析

首先定义平均距离方差指标(Mean Squared Distance Error) - MSDE，

$$MSDE(q) \triangleq \iint (d(x,y)-\hat d(x,y))^2p(x)dx\space p(y)dy$$

根据三角不等式，

$$d(x,q(y))-d(y,q(y))\le d(x,y) \le d(x,q(y))+d(y,q(y))$$

于是有，

$$(d(x,y)-d(x,q(y)))^2 \le d(y,q(y))^2$$

因此由MSDE函数定义可得到下面的不等式

$$\begin{align*} MSDE(q)&\le \int p(x)(\int d(y,q(y))^2p(y)dy)dx \\ &\le MSE(q) \end{align*}$$

式中的MSE(q)是前面已定义的。

由以上不等式表明，我们的量化方法误差上界是$MSE(q)$，而我们的quantizer是用k-means算法进行学习的，因此可以保证$MSE(q)$满足进行虽有条件。这样，我们的最终算法的误差控制是比较好的。

但是不管怎么样，用量化值进行距离计算总是免不了偏差的。如下图，分别比较了实际的距离和ADC,SDC的距离之间的函数关系，

可以看到整体来说，量化距离和实际距离是呈正相关关系的，但是误差也是不可避免的。同时也可以看到，量化距离一般比实际距离的值要偏小一些，并且SDC算法的偏差相对ADC的偏差更大。原论文对偏差做了一定的修正。在这里就不赘述了。因为在实际应用中，可以不进行校正。不过从上图也可以看出用ADC算法会更加好一些。下面我们直接看一下检索过程。

为了快速检索，不例外地这里用到了倒排索引结构(inverted index)。上图可以进行说明，在实际检索中分两个步骤：

1. 用数据库中特征建立索引结构(图的上部分)
2. 对输入的待检测特征进行检索(图的下部分)

从图中可以看到。在建立索引时，先用k-means算法对数据库特征进行粗量化，得到$k’$个索引项；粗量化之后，由于会存在较大的量化误差。为了保存这个信息，这里计算query vector和粗量化的code word之间的residual vecto-$r$。这里的residual vector就是将使用本文所讲的product quantization 算法进行量化，最后将量化后得到的quantization code 存储在该索引项对应的列表项中，同时列表项中还存储中图片的identifier。这样一个向量$y$就可以用下面的向量近似：

$$\ddot {y} \triangleq q_c(y)+q_p(y-q_c(y))$$

因此$d(x,y)$可以用下面的近似算法进行计算：

$$\ddot{d}(x,y)=d(x,\ddot{y})=d(x-q_c(y),q_p(y-q_c(y)))$$

将计算任务分配到product quantization的各个subquantizers中可得到下式：

$$\ddot{d}(x,y)^2=\sum_jd(u_j(x-q_c(y)),q_{p_j}(u_j(y-q_c(y)))))^2$$

1. $x$项是粗量化的code word，因此是precomputed，因此该式子第一部分是可以离线计算的。
2. 考虑到在coarse quantization中，$x$和它的近邻不一定落在同一个簇中，所以在查询coarse quantization时，会同时取出$w$个倒排链。
3. 对取出的每个倒排链，还是用PQ算法找出近邻。
4. 考虑当$n>k^\ast D^\ast$时，朴素的ADC算法的复杂度是$O(n \times m)$，而IVFADC算法的复杂度会降低为$O((n \times w / k’) \times m)$。
5. 这里得存储粗量化所有的code word的sub vector到所有的product quantization的code word 的距离，也是非常耗内存的。为了减少量化误差，论文中还应用了multiple assignment的方法，就是将一个特征映射到多个索引上。

Multiple Assignment

如上图所示，$x$与$y$都分到了$1$这个voronoi包，如果不采用multiple assignment方法，则只会在voronoi包$1$中搜索与$x$接近的$y$，明显voronoi包$2$，voronoi包$3$中与$x$更近的$y$不会被检索到，采用multiple assignment的方法，在$x$的$w$个最近粗糙类中心中搜索$y$，那么准确率就提高了。但这要求遍历$w$个list，费内存和时间，是以牺牲空间和时间换准确率。

关于如何选取coarse_centroids的个数，multiple assignment的$w$值和寻找的近邻$y$的个数参考下图：

总结

其他基本看懂了，但是这个$u_j(·)$是个啥啊？