NVIDIA 深入探讨 RAPIDS cuVS IVF-PQ 以加速向量搜索
realtime news Jul 18, 2024 20:29
NVIDIA 探索 RAPIDS cuVS IVF-PQ 算法，通过压缩和 GPU 加速提升向量搜索性能。
在一篇详细的博客文章中，NVIDIA 介绍了他们的 RAPIDS cuVS IVF-PQ 算法，该算法旨在通过利用 GPU 技术和先进的压缩技术来加速向量搜索。这是一个两部分系列的第一部分，延续了他们对 IVF-Flat 算法的先前探索。
IVF-PQ 算法介绍
博客文章介绍了 IVF-PQ（带产品量化的倒排文件索引），这是一种通过以压缩形式存储数据来提升搜索性能并减少内存使用的算法。然而，这种方法以一定的准确性为代价，第二部分的系列将进一步探讨这个权衡。
IVF-PQ 基于 IVF-Flat 的概念，后者使用倒排文件索引通过聚类将搜索复杂性限制在一个较小的数据子集内。产品量化（PQ）通过对数据库向量进行编码，添加了另一层压缩，使得该过程在处理大数据集时更加高效。
性能基准测试
NVIDIA 使用包含十亿条记录和 96 个维度、总计 360 GiB 的 DEEP 数据集分享了基准测试结果。典型的 IVF-PQ 配置将其压缩成 54 GiB 的索引，而不显著影响搜索性能，或者压缩到 24 GiB 略有减速。这种压缩允许索引适合 GPU 内存。
将其与流行的 CPU 算法 HNSW 在 DEEP 数据集的一亿子集上的比较表明，cuVS IVF-PQ 可以显著加速索引构建和向量搜索。
算法概述
IVF-PQ 遵循两步法：粗略搜索和精细搜索。粗略搜索与 IVF-Flat 一致，而精细搜索则涉及计算查询点与探测到的聚类中向量之间的距离，但向量以压缩格式存储。
这种压缩是通过 PQ 实现的，它使用两级量化来近似向量。这使得 IVF-PQ 能够将更多数据适配到 GPU 内存中，从而增强内存带宽利用率并加快搜索过程。
优化和性能
NVIDIA 在 cuVS 中实施了各种优化，以确保 IVF-PQ 算法在 GPU 上高效运行。这些优化包括：
- 融合操作以减少输出大小并优化内存带宽利用。
- 在可能时将查找表（LUT）存储在 GPU 共享内存中以加快访问速度。
- 在 LUT 中使用自定义 8 位浮点数据类型以加快数据转换。
- 将数据对齐为16字节块以优化数据传输。
- 实施“提前停止”检查以避免不必要的距离计算。
NVIDIA 对一亿规模的数据集进行的基准测试显示，IVF-PQ 尤其在较大批次情况下，相较于 IVF-Flat 显著提高了每秒查询数，达到3-4倍。
结论
IVF-PQ 是一个稳健的近似最近邻搜索（ANN）算法，通过聚类和压缩来提升搜索性能和吞吐量。NVIDIA 博客系列的第一部分提供了算法工作原理及其在 GPU 平台上的优势的全面概述。欲了解更详细的性能调优建议，NVIDIA 鼓励读者探索他们系列的第二部分。
欲了解更多信息，请访问 NVIDIA 技术博客。
Image source: Shutterstock