如何利用SIMD加速量化码扫描50%?

向量检索实验室 2024-09-10

468

今天为大家分享一篇ICMR 2017年的向量检索论文：

Accelerated Nearest Neighbor Search with Quick ADC

使用快速不对称距离计算技术加速最近邻搜索

论文作者来自法国Technicolor公司和Inria研究所。

本文技术已在FAISS库中支持，即IVFPQfs (fast scan)索引，代码具体位置如下。

论文地址：https://arxiv.org/pdf/1704.07355

代码地址：https://github.com/facebookresearch/faiss/blob/main/faiss/IndexIVFPQFastScan.h

论文概述

基于量化的索引，比如IVFPQ，在查询过程中计算距离时是通过查表完成的，而非对原始向量完整计算距离，比如欧式或内积距离。而在ANN算法的实现中，SIMD（Single Instruction Multiple Data,单指令多数据）技术（也称向量化执行技术）对于距离计算（比如L2距离、内积距离）的加速至关重要。如何使用SIMD技术加速PQ索引的查表式距离计算是论文研究的问题。

PQ索引查询时，首先会构建查找表，即预计算查询向量与量化码本中每个（子）码字的距离。在这之后，每次扫描到一个码字，会根据其中的子码字去查找表中取出预计算的距离，各个子码字对应的距离相加，即为PQ中的距离（估计距离，ADC）。

在这个过程中，包含取码字、读距离、距离加和三个操作。论文对这三个阶段分别设计了高效的SIMD算法，包括计算方式的重组、码字的内存重排和量化参数的设定。论文结果显示，这种方式可以获得近50%左右的性能提升。

关键算法

PQ中的参数影响

PQ索引中有两个关键参数，一个是分段数，一个是子码字的比特数。每个数据库中向量会被编码成比特的码字。同等内存占用下，分段数越大，量化损失越高，距离计算速度越慢，但是查找表会更小，构建查找表的速度也会更快。

下面的表格描述了64bit下，三种参数配置的性能表现。可以看到，当分段数小到刚好把查找表放进L1 cache时，查询性能最高。

128bit SIMD寄存器

现有的CPU的每个核中，都会有一个SIMD寄存器用以执行SIMD指令。不同指令集的寄存器大小不同，比如SSE是128bit的，AVX是256bit的。由于查找表是PQ扫描码字时频繁访问的数据结构，如果能将查找表（每个段的子表）放入SIMD寄存器，并行化查找距离，则可以减少串行访问cache带来的开销。

因此，论文采取的策略是固定参数，此时每个子查找表则有个浮点数构成，再将浮点数做8比特的量化，则每个子查找表只有比特构成，可以放入寄存器。

查找表读取

从查找表中读取数据可以使用pshufb
指令，可以同时在SIMD寄存器中执行16次表查找的操作。因此，可以同时执行16个向量的距离计算（查表）。也就是说，读取16个向量的第一个子码字，注意每个子码字只有比特，如下图->共16个向量的第一段子码字。

码字读取

由于我们每次会读取16个码字的第一段，因此这些数据最好能在内存（cache）中连续排布以降低读取代价。论文将向量的码字表翻转了一下，使得每128bit是16个向量->的前两个段的码字，也就是说每16个向量被组织成一个block进行存储，是最小的距离计算单位。

距离计算流程

综合以上方法，quick ADC在进行计算流程如下：

读码字：每次读取一个block的一行128bit，mask掉每个字节的高4位；
查距离：在SIMD寄存器中读取查找表，执行pshufb
读取距离；
距离加和：读取的距离与每个向量累积的距离做加和（）。

由于每次读取一个向量的两个子码字，所以计算完一个子码字后，可以右移以获得下一个子码字，减少一次码字读取。

精彩段落

<<< 左右滑动见更多 >>>

总结

论文提出的quick ADC是非常实用的对于PQ索引进行SIMD计算加速的优化方法，已经在FAISS中得到支持。但同时也对参数进行了限制，即4bit量化，这意味着在某些数据集上带来了较大的精度损失。更为常见的参数是8bit与16bit量化，这在后续的工作中有所补充。

延伸阅读

[1] TPAMI'21：AVX512量化加速 Quicker ADC : Unlocking the Hidden Potential of Product Quantization With SIMD (https://arxiv.org/pdf/1812.09162)

[2] VLDB'15: PQ暴力扫描SIMD优化方法 Cache locality is not enough: high-performance nearest neighbor search with product quantization fast scan (https://inria.hal.science/hal-01239055/file/p288-andre.pdf)

编者简介

王泽宇

复旦大学与巴黎西岱大学联合培养博士生，研究领域为高维数据（向量、序列等）管理和分析。以第一作者在SIGMOD,VLDB,VLDBJ,TKDE等数据库领域会议/期刊发表多篇论文，并担任审稿人。

个人主页：https://zeyuwang.top/

谷歌学术：https://scholar.google.com.hk/citations?hl=zh-CN&user=XXGhABIAAAAJ

技术博客：https://www.jianshu.com/u/d015902c6d09

👆 关注 AI 搜索引擎，获取更多专业技术分享 ~

查找算法量化管理

文章转载自向量检索实验室，如果涉嫌侵权，请发送邮件至：contact@modb.pro进行举报，并提供相关证据，一经查实，墨天轮将立刻删除相关内容。