实例理解product quantization算法

近几年,深度学习技术被广泛用于图像识别、语音识别、自然语言处理等领域,能够把每个实体(图像、语音、文本)转换为对应的embedding向量。如这里千人千面智能淘宝店铺背后的算法研究登陆人工智能顶级会议AAAI 2017。而对于推荐、搜索或者广告投放问题,都可以描述为从大规模候选中给用户提供有限的展现结果。那么,这里就会涉及到向量检索的问题。

向量检索最简单的想法是暴力穷举法,如果全部实体的个数是$n$$n$是千万量级甚至是上亿的规模,检索实体对应的向量是D,那么当要从这个实体集合中寻找某个实体的相似实体,暴力穷举的计算复杂度是$O(n×D)$,这是一个非常大的计算量,该方法显然不可取。所以对大数据量下高维度数据的相似搜索场景,我们就需要一些高效的相似搜索技术,而PQ就是其中一类方法。

1. Product Quantization算法

假设有50,000张图片组成的图片集,使用 CNN 提取特征后,每张图片可以由1024维的特征表示。那么整个图片集由50000*1024的向量来表示。如下图。

image.png

然后我们把1024维的向量平均分成m=8个子向量,每组子向量128维。如下图。

image.png

对于这8组子向量的每组子向量,使用 kmeans 方法聚成k=256类。也就是说,每个子向量有256个中心点(centroids)。如下图。

image.png

在product quantization方法中,这256个中心点构成一个码本。这些码本的笛卡尔积就是原始D维向量对应的码本。用$q_j$表示第$j$组子向量,用$C_j$表示其对应学习到的码本,那么原始D维向量对应的码本就是$C=C_1×C_2×…×C_m$,码本大小为$k^m$

注意到每组子向量有其256个中心点,我们可以中心点的 ID 来表示每组子向量中的每个向量。中心点的 ID只需要 8位(=$log_2{256}$)来保存即可。这样,初始一个由32位浮点数组成的1,024维向量,可以转化为8个8位整数组成。如下图。

image.png

2. 最近邻搜索

对向量压缩后,有2种方法作相似搜索。一种是SDC(symmetric distance computation),另一种是ADC(asymmetric distance computation)。SDC算法和ADC算法的区别在于是否要对查询向量x做量化,参见公式1和公式2。如下图所示,x是查询向量(query vector),y是数据集中的某个向量,目标是要在数据集中找到x的相似向量。

image.png

SDC算法:先用PQ量化器对x和y表示为对应的中心点q(x)和q(y),然后用公式1来近似d(x,y)。这里 q 表示 PQ量化过程。

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

ADC算法:只对y表示为对应的中心点q(y),然后用下述公式2来近似d(x,y)。

$\widetilde{d}(x,y)=d(x,q(y))=\sqrt{\sum_j{d(u_j(x),q_j(u_j(y)))^2}} $ (1)

3. python例子

如果我们以纯 python来模拟最近邻搜索的话,可以分这样几步。

step 1. 对初始向量按列平均分成 m 组,以每组作 kmeans 聚类。
vecs_sub = vecs[:, m * self.Ds : (m+1) * self.Ds]
self.codewords[m], _ = kmeans2(vecs_sub, self.Ks, iter=iter, minit='points')
step2. 得到各分组子向量的每个向量的聚类的 ID 号。这样向量量化完成。
codes[:, m], _ = vq(vecs_sub, self.codewords[m])
step3. 对于查询向量$q$,将$q$按列也分成 m 组后,计算与聚类中心的距离。
for m in range(self.M):
            query_sub = query[m * self.Ds : (m+1) * self.Ds]
            dtable[m, :] = np.linalg.norm(self.codewords[m] - query_sub, axis=1) ** 2
step 4. 按组将与聚类中心的距离求和后,取其距离之和最小值所对应的聚类中心。这些聚类中心所对应的初始向量,即最近邻向量。
dists = np.sum(self.dtable[range(M), codes], axis=1)

完整代码见这里https://github.com/matsui528/nanopq

4. Optimized product quantization(OPQ)方法

Product Quantization的本质是将原始高维空间分解为有限数量的低维子空间的笛卡尔积,然后分别量化。OPQ 试图寻找一个正交矩阵,将原始矩阵旋转后再行分解,以使量化后的向量重建后,其误差最小。如下图。

image.png

其中的一种寻找旋转矩阵$R$的方法如下。

step1. 固定旋转矩阵$R$,使用 PQ 方法,求解量化后的向量$c^1,c^2,...c^M$。即

$min_{C^1,C^2,...C^M}{\sum_x{||\hat{x}-\hat{c}_i(x)||^2}}$

这里$ \hat{x}=Rx, \hat{c}=Rc $

step2. 固定$c^1,c^2,...c^M$,求解 $R$

$min_{R}{\sum_x{||Rx-\hat{c}_i(x)||^2}}$

s.t. $R^TR=I$

$y=\hat{c}_i(x)$,则上式写为

$min_{R}{||RX-Y||^2_F}$

可得到 $R=VU^T, [U,V]=svd(XY^T)$

5. 参考链接