Efficient Large-scale Approximate Nearest Neighbor Search on the GPU 阅读笔记 – 实现与分析

由于篇幅太长,这篇文章的Introduction和Related Work在另一篇文章里。


PQT的简介

乘积量化树(Product Quantization Tree)是基于多层倒排表(Inverted Multi-Index)和分层PQ算法形成的。其核心思想将之前的一个码书(Codebook)修改为基于层次化的树形结构。树形结构可以加速实时的查询,以及离线的排序。相比于之前的多层倒排表,PQT拥有更多地簇(Bins)。最后PQT还巧妙地运用一些技巧在重排序中服用之前计算过的距离值。

右图阐述了一个PQT查询的例子。查询向量[x]_1[x]_2的各个部分一次经过4个Cluster的粗排序和5个Cluster的精排序。对于第一层而言,只有w个最佳的Cluster可以到下一层继续搜索,如图中阴影部分所示(w=2)。

PQT的结构

将样本集合中所有的向量x \in \mathcal{X}归类为簇B_l从而形成若干个互不相交的子集,\mathcal{X} = \bigcup_{k=1}^{K} B_k。接下来,我们将详述如何将一个向量映射到簇中,即如何构造映射m : \mathcal{X} \rightarrow \mathcal{I}_1 \times \mathcal{I}_2 \times \dots \times \mathcal{I}_P

索引树包含两层量化器,第一层是传统的P部分(P part)的PQ量化器,在这个量化器中,每一部分的码表含有k_1个码字。在第一层量化器之后,选出的簇进一步使用第二层k_2个码字的量化器量化。这种方法中,总共存在的簇有K = (k_1 \cdot k_2)^P个。

对于码书的训练,每一部分独立生成码书,首先量化器使用Lloyd Iteration中的Linde-Buzo-Gray算法生成,所有第一层量化器分出的簇进一步使用第二层量化。

对比在多层倒排表中,使用了两级的乘积量化,其中第二级用于排序,我们也用两级索引,至于为什么是两级,只是从经验上来看,两级能够最好得平衡簇的总个数。

PQT的询问操作

一个实际的询问操作包含四个步骤——

  • tree traversal
  • bin proposal
  • vector proposal
  • re-ranking

在tree-traversal步骤中,正如之前对PQT结构的叙述,每一个向量都能归类为一个level-2的簇中。通过剪枝的策略,tree-traversal步骤使用的向量比较次数实际上是相对较少的。在比较了所有k_1个level-1簇后,询问向量距离每一个簇的距离被排序,并且只有最近的w个簇被进一步计算。在提取出了最近的w个level-1簇,接下来这些level-1簇所对应的level-2簇的中心再和询问向量进行比较。在标准参数下,在四百万簇的数据集下,只需要80次向量比较。具体细节可参见论文。

在bin-proposal步骤中,从我们获取的“最好的”簇i=(i_{1,1},i_{2,1},\dots,i_{P,1})开始,我们仍然需要一系列临近的簇,以保证在后续的处理中,有足够多的向量可供选择。正如右图所示,我们需要找到一个顺序,合理组织“次好”的向量,使得遍历完所有较近向量的这组合。着这种便利,通常来说,可以看做是一个斜线从左下向右上方移动维护的下方区域。具体的实现方式,绿线是Dijkstra-based traverse,但是无法再GPU上面实现,而红色、绿色线则是可行的。对于每一个簇,我们都可以在离线阶段预处理出最近的簇。

使用Line Quantization进行重排序

通过上述的索引,每一个询问向量都被归类于最近的簇,并获得了一个量化损失(quantization loss)。在传统方法中,我们就需要依次对于簇所包含的所有向量进行比较,但是显然这样会有巨大的花销。而在我们的算法中,我们重用了一些信息以达到更优的效率。

离线预处理

在传统方法中,一个簇所包含的所有向量是不可分的,这就导致了我们必须花巨大的计算力去依次比较所有的向量。 文中提出的Line Quantization巧妙地解决了向量不可分的问题。即,虽然每一个向量【黑点】相对于level-2簇中心【红点】是不可区分。但是他却可以被投影到两个最近level-1簇中心的连线上【紫点】。巧妙的是,对于所有level-1中心,询问的x_1点与他们的距离都是已经算过的!

右图描述了如何计算对于任意询问向量y距离某一个数据库样本x的距离。其中,c_i,c_j为两个距离x很近的level-1簇中心。

 

(1)   \begin{align*} d(y,x) &= \sum_{p=1}^{P} d([y]_p, [x]_p)^2 \\ & \apporx \sum_{p=1}^{P} h_p(y,x)^2 \\ & \apporx \sum_{p=1}^{P} (b_p^2 + \lambda_p^2 \cdot c_p^2 + \lambda_p \cdot (a_p^2-b_p^2-c_p^2)) \end{align*}

通过这个式子,我们可以用我们之前计算过的,询问向量距离level-1簇中心的距离信息,为x找到一个更加精确的距离估计值。

在GPU上实现

这个模型非常适合在GPU上面实现,不过由于我并不太了解GPU的实现,顾就不在这里赘述,读者可以对照论文自行理解。

原创文章地址:【Efficient Large-scale Approximate Nearest Neighbor Search on the GPU 阅读笔记 – 实现与分析】,转载时请注明出处mhy12345.xyz

《Efficient Large-scale Approximate Nearest Neighbor Search on the GPU 阅读笔记 – 实现与分析》有一个想法

发表评论

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

This site uses Akismet to reduce spam. Learn how your comment data is processed.