[论文解读] RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search

Jianyang Gao, Long Cheng|arXiv (Cornell University)|May 21, 2024

Advanced Image and Video Retrieval Techniques被引用 4

一句话总结

RaBitQ 是一种新颖的随机量化方法，通过随机投影和双值表示将高维向量编码为 D 位字符串，实现具有可证明理论误差边界的高效近似最近邻搜索。与产品量化（PQ）及其变体相比，它通过支持无偏、有界误差的快速按位和 SIMD 基距离估计，在准确率-效率权衡上表现更优。

ABSTRACT

Searching for approximate nearest neighbors (ANN) in the high-dimensional Euclidean space is a pivotal problem. Recently, with the help of fast SIMD-based implementations, Product Quantization (PQ) and its variants can often efficiently and accurately estimate the distances between the vectors and have achieved great success in the in-memory ANN search. Despite their empirical success, we note that these methods do not have a theoretical error bound and are observed to fail disastrously on some real-world datasets. Motivated by this, we propose a new randomized quantization method named RaBitQ, which quantizes $D$-dimensional vectors into $D$-bit strings. RaBitQ guarantees a sharp theoretical error bound and provides good empirical accuracy at the same time. In addition, we introduce efficient implementations of RaBitQ, supporting to estimate the distances with bitwise operations or SIMD-based operations. Extensive experiments on real-world datasets confirm that (1) our method outperforms PQ and its variants in terms of accuracy-efficiency trade-off by a clear margin and (2) its empirical performance is well-aligned with our theoretical analysis.

研究动机与目标

• 解决现有量化方法（如产品量化，PQ）缺乏理论误差边界的缺陷，这些方法在真实数据集上可能不可预测地失效。
• 开发一种量化方案，确保在距离估计上具有精确的理论误差边界，同时保持高经验准确性。
• 通过利用按位操作和 SIMD 指令的优化实现，实现高效的内存内近似最近邻（ANN）搜索。
• 确保距离估计误差对每个数据向量均有界（以高概率），而不仅是在期望值上。
• 提供一种兼具理论严谨性和实际效率的方法，克服以往方法为追求速度而牺牲理论保证的局限性。

提出的方法

• 在量化前对 D 维向量应用随机旋转和投影，生成双值表示（±1）。
• 通过将投影值在零处进行阈值化，将每个 D 维向量编码为 D 位字符串，形成二进制码。
• 使用查询向量与量化数据向量之间的内积作为真实内积的无偏估计器。
• 利用集中不等式推导估计距离的理论误差边界，确保误差以高概率有界。
• 通过按位操作（XOR 和 popcount）实现高效的单次查询距离估计，以及通过 SIMD 加速实现批量处理。
• 将查询向量量化为 4 位整数，以最小化查询端的误差，同时保持数据向量完全量化为位。

实验结果

研究问题

• RQ1 能否设计一种量化方法，为高维向量的距离估计提供可证明的理论误差边界？
• RQ2 如何使该方法在仅使用按位或 SIMD 操作的情况下，保持在内存内近似最近邻搜索中的高效率？
• RQ3 该理论误差边界是否与真实世界数据集上的实际性能一致，特别是在 PQ 失效的情况下？
• RQ4 该方法能否在保证每个向量误差有界的前提下，超越 PQ 及其变体在准确率-效率权衡上的表现？
• RQ5 仅量化数据向量（不量化查询）对系统整体误差和效率有何影响？

主要发现

• RaBitQ 在距离估计上实现了精确的理论误差边界，确保对每个数据向量以高概率具有有界误差。
• 在 MSong 数据集上，RaBitQ 将距离估计的平均相对误差降低至 10% 以下，而 PQ 在使用 SIMD 加速时误差超过 50%。
• 即使不进行重排序，RaBitQ 的召回率也高于 PQ（例如，在 k=10 时召回率 >90%），而 PQ 在相同条件下召回率低于 60%。
• 该方法仅通过按位操作（XOR 和 popcount）即可实现高效的单次查询距离估计，使现代 CPU 上每查询延迟低于 100ns。
• 使用 SIMD 指令进行批量距离估计相比基线实现可实现 3–5 倍的加速，且准确率损失极小。
• RaBitQ 的实际性能与理论预测高度一致，证实了在多样化真实世界数据集上误差边界的可靠性。