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

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

高维数据的最近邻搜索是许多应用(如图像检索、多媒体数据库)中的一个重要特征。产品量化(Product Quantization, PQ)是一种广泛使用的高性能解决方案，它可以在保持高精度的同时降低响应时间。PQ通过紧凑编码表示高维向量(例如，图像描述子)。因此，可以将非常大的数据库存储在内存中，从而允许NN查询而无需求助于缓慢的I/O操作。PQ使用缓存查找表来计算与邻居的距离，因此它的性能仍然受到限制(i)算法需要大量的缓存访问，以及(ii)它无法利用现代cpu上可用的SIMD指令。

本文认为缓存局部性不足以提高效率。为了解决这些限制，我们设计了一种新的算法，PQ快速扫描，将cacherresident查找表转换为小表，大小适合SIMD寄存器。这种转换允许(i)寄存器内查找代替缓存访问，以及(ii)高效的SIMD实现。PQ快速扫描与PQ具有完全相同的精度，但响应时间比PQ低4到6倍(例如，对于2500万向量，扫描时间从74ms减少到13ms)。 

一研究问题

为了回答ANN查询，PQ使用缓存驻留查找表计算查询向量和大量数据库向量之间的距离。这个过程称为PQ扫描，具有很高的CPU成本，是本文的重点。 

通过PQ Scan的例子，我们研究了为什么基于查找表的算法不能利用SIMD指令。我们提出了缓存驻留查找表的替代方案，以实现更好的性能。在此基础上，设计了一种新的PQ快速扫描算法。PQ快速扫描的性能比PQ扫描好46倍，同时返回完全相同的结果。

1. PQ SCAN LIMITATIONS

PQ扫描在cpu上无法高效实现。我们确定了限制其性能的两个基本瓶颈:(i)大量的cache访问和(ii)无法使用SIMD高效地实现它。识别这些瓶颈是设计克服PQ扫描限制的新算法的关键。

2. PQ FAST SCAN 

本文提出了PQ快速扫描算法，它克服了PQ扫描的局限性。PQ快速扫描对每个扫描向量的L1 cache访问次数少于2次，并且允许使用SIMD指令高效地实现加法运算。通过设计，PQ快速扫描返回与PQ扫描PQ8×8量化器完全相同的结果，同时执行4-6倍更快。

PQ快速扫描背后的关键思想是使用小表，其大小适合SIMD寄存器，而不是驻留在缓存中的距离表。这些小表用于计算距离的下界，不需要访问L1缓存。因此，下界计算速度很快。此外，它们使用SIMD加法实现，进一步提高了性能。我们使用下界计算来修剪缓慢的pqdistance计算，这些计算需要访问L1缓存，无法从SIMD加法中获益。

为了计算下界，我们需要在SIMD寄存器中存储的小表中查找值。为此，我们使用pshufb指令，这是PQ快速扫描性能的关键。