【论文阅读】Cache locality is not enough

Cache locality is not enough:High-Performance Nearest Neighbor Search with Product Quantization Fast Scan

作者: Fabien Andre´、Anne-Marie Kermarrec、Nicolas Le Scouarnec
关键词: PQ, SIMD, 量化
发表年份: 2015

Abstract

PQ使用驻留在cache中的距离表来计算，其性能受到以下限制：

• 需要大量的高速缓存访问
• 无法利用现代CPU上可用的SIMD指令

由此设计了算法PQ Fast Scan：将驻留在cache中的距离表转换成大小适合SIMD寄存器的small tables。就可以

• 直接寄存器内查找代替高速缓存访问
• 有效地利用SIMD实现

Intro

PQ scan：PQ使用驻留在缓存中的距离表计算查询向量和大量数据库向量之间的距离。——花费很高的CPU成本、依赖主存、无法充分利用现代CPU性能、无法利用SIMD指令集。

贡献：

• 分析了PQ scan性能。研究表明调整距离表的大小以适应L1缓存对于性能来说是不够的。PQ scan无法通过SIMD指令有效实现，即使使用最新一代英特尔CPU中可用的收集指令也是如此
• 设计了PQ fast scan 来解决这些问题。背后的关键思想是建立小表，大小适合SIMD寄存器，可以使用快速SIMD指令查找。使用这些小表来计算距离的下限，并避免对驻留在高速缓存中的查找表进行不必要的访问（这种技术允许修剪95%以上的L1缓存访问，提升了4-6倍速度）。
• 构建小表需要: (I)相似向量的分组，(ii)最小表的计算，以及(iii)浮点距离到8位整数的量化。

Back

PQ

IVFADC：在PQ和ADC之上增加了一个倒排索引(IVF)。索引是从一个基本的矢量量化器(称为粗量化器)建立的。粗略量化器的每个Voronoi单元形成数据库的一个分区。用IVFADC回应ANN查询包括三个步骤:

1. 使用索引选择数据库的分区。所选分区对应于查询向量所在的粗量化器的Voronoi单元；
2. 计算距离表，其用于加速ADC计算；
3. 扫描分区，即计算查询向量和分区的所有向量之间的ADC。此步骤命名为PQ scan。

其中第一步和第二步只占不到1%的CPU时间，而步骤3占用了99%以上的CPU时间。

PQ scan 的两个基本瓶颈 :(I)它执行的许多缓存访问 (ii)不可能用SIMD有效地实现它。

PQ scan的局限

Memory accesses

每个 pqdistance计算包括：

• 加载簇心索引$p[j]$的 $m$ 次内存访问 [mem1]
• 从距离表加载$D_j[index]$ 值的 $m$ 次内存访问 [mem2]
• m次加法

mem1 访问质心索引，mem2 访问距离表，它们可能命中不同的高速缓存级别。

由于现代CPU中包含的硬件预取器，mem1访问总是命中L1 cache（硬件预取器能够检测连续的存储器访问，并将数据预取到L1高速缓存）。所以顺序地访问$p[j]$值，即首先访问$p[0]$（p是第一个向量），然后是$p[1]$直到$p[m-1]$。接下来，对第二个数据库向量执行相同的访问…….

mem2 所访问的高速缓存级别取决于乘积量化器的m和k *参数（m是子量化器数、k是单位子量化器的簇心数）。较大的距离表存储在较高的高速缓存级别中。

（PQ 8×8提供了最佳的性能权衡，是文献中最常用的配置）

使用性能计数器来实验性地分析不同PQ scan实现的性能。

• 对于所有实现，具有未决加载操作的周期数(周期w/load)几乎等于周期数，这证实了PQ扫描是一种存储器密集型算法。

无法利用SIMD指令

8条指令，但每条指令都在一个表中进行查找的。如果是a0、a1、…. 、a7连续加，则要在8个距离表中查询。显然SIMD更好。

但使用SIMD加法带来的周期增益被逐个设置SIMD寄存器的需要所抵消。对于PQ scan来说，除了内存访问之外，这样做还需要许多SIMD指令，其中一些具有很高的延迟。总的来说，这抵消了SIMD附加物带来的好处。这解释了为什么依赖于查找表的算法，如PQ扫描，很难从SIMD处理中受益。（具体见3.2）

PQFS

PQ fast scan对每个扫描向量执行不到2次L1缓存访问，并允许使用SIMD指令高效地执行加法。

背后的关键思想：使用大小适合SIMD寄存器的小表，而不是驻留在缓存中的距离表。小表用于计算距离的下限，而无需访问L1缓存。下界计算是快速的（用SIMD addition实现、在寄存器中实现），使用下限计算来修剪缓慢的pqdistance计算（访问L1缓存，且不能从SIMD加法中受益。）

（上图表示每个数据库向量p的处理步骤。⊗意味着丢弃向量p并移动到下一个数据库向量。最小值是查询向量到当前最近邻居的距离。）

在SIFT数据上的实验结果表明，PQ快速扫描能够减少95%的PQ距离计算。

使用pshufb指令：在SIMD寄存器中存储的小表中查找值。但得将每个小表的大小限制为16个元素，每个元素8位(16×8位，128位)。（原始PQ scan的一个距离表是256x32位）

为建立适合于计算距离下限的8个小表，结合了三种技术:

• 向量分组
• 最小表的计算
• 距离的量化

前两种技术，用于建立16个元素(16×32位)的表。第三种技术，用于将每个元素缩小到8位(16×32位→16×8位)。总的来说是：

• 对向量进行分组并量化距离，以构建前四个小表 $S_0,\dots ,S_3$ 。
• 计算最小表和量化距离，以建立最后四个小表 $S_4,\dots ,S_7$xi

向量分组

对向量进行分组，使得属于一个组的所有向量命中距离表的16个划分中的相同部分。具体方法为：

• 将几个距离表分别（以$D_0$ 为例）分割成16个部分，每个部分包含16个元素。计算pq距离时，第一分量（子量化器）$p[0]$在00和0f (0到15)之间的所有数据库向量p将触发在D0的第0部分中的查找，组成组0。
• 以此类推，对前四个子量化器对应的距离表$D_0,D_1,D_2,D_3$ ——进行划分，每个距离表中有16个组，则一共有$16^4$ 个组。
• 每个向量的因前四个子码被划到组 $(i_0,i_1,i_2,i_3)$ ——表示属于第1个距离表的第 $i_0$ 组，第2个距离表的第 $i_1$ 组，第3个距离表的第 $i_2$ 组，第4个距离表的第 $i_3$ 组。（如下图的组（3，1，2，0）) 。因为每次扫描的是同一个组的向量，所以向量的四个组号（共4x4bit）可以省略不记（下图灰色数组为无需记录的组号）。
• 在扫描一个组前，加载距离表$D_0,D_1,D_2,D_3$ 的相关部分（比如$D_0$的第 $i_0$ 部分，$D_1$的第 $i_1$ 部分，$D_2$的第 $i_2$ 部分，$D_3$的第 $i_3$ 部分）到四个SIMD寄存器中，作为小表$S_0,S_1,S_2,S_3$。

如果有n个数据，c个距离表需要分组的话，那么平均每组的数据量大小为$s=n/16^c$。

在扫描一个组前，需要把距离表的部分内容加载到SIMD寄存器中，是十分耗时的，所以s应该超过50个向量（否则CPU大部分时间被用来加载距离表上），这就要保证n在3百万以上。

最小表

• 把原始距离表$D_4,D_5,D_6,D_7$ 分别分成16份，每份16个元素。取每个部分的最小值，得到4个包含16个元素的表$S_4,S_5,S_6,S_7$。

（只使用最小表技术会产生包含低值，这对扫描性能是有害的。如果这些值太低，则计算出的下限不紧密，即远离实际的pqdistance）

• 为获得具有更高值的小表，引入子量化器簇心索引的优化分配：

由于学习子量化器时，簇心索引被任意分配。因此，具有对应于距离表的一部分的索引(例如，00到0f) 的簇心之间没有特定的关系，所以优化分配要确保对应于给定部分(例如，00到0f)的所有索引被分配给彼此靠近的簇心（因为靠近给定簇心的查询子向量也可能靠近附近的质心，距离表中给定部分的所有值都将是接近的。这就允许计算具有更高值的最小表，因此具有更紧的下限。）

具体方法：将质心分成16组，每组16个元素，使用一种k-means的变体，迫使组具有相同的大小。对应于距离表的一部分，相同簇中的簇心被给予连续的索引。

距离量化

将32位浮点距离量化为8位整数。（由于没有SIMD指令来比较无符号的8位整数，需要将距离量化为有符号的8位整数。）

将$qmin$和$qmax$之间的浮点距离量化为n =127个区间。每个区间的大小为$(qmax-qmin)/n$，区间号(0-126)用作量化浮点的表示值。$qmax$以上的所有距离都量化为127。

$qmin$ ：为所有距离表中的最小值

$qmax$：使用原始PQscan算法在数据库的前 keep%向量(通常keep ≈1%)，中找到查询向量的临时最近邻。然后使用这个临时最近邻之间的距离作为$qmax$

（最后为了避免整数溢出问题，使用饱和SIMD加法）

在小表中查找

在扫描每组之前，我们将这些量化后部分距离表 $S_0,...,S_3$ 载入SIMD寄存器（实线箭头所示），中途需要改变。最后四张小表$S_4,...,S_7$，在扫描过程开始时被载入SIMD寄存器，之后一直不变。

使用p[0]p[3]的4个最低有效位来索引表 $S_0,...,S_3$ 中的值，4个最高有效位来索引表$S_4,...,S_7$的值。（在小表中的查找由虚线箭头表示。虚线箭头所示的查找是使用pshufb执行的。）

将8个查找值相加得到距离下限。将下限与min的量化值进行比较，min是查询向量和当前最近邻居之间的距离。

实验