IVF-PQ

IVF-PQ 是一种倒排文件索引 (IVF) 算法，它是 IVF-Flat 算法的扩展（例如，数据点首先被分成聚类），在每个聚类内执行乘积量化，以缩小索引的内存占用。乘积量化是一种有损压缩方法，它能够通过将原始向量卸载到主内存来在 GPU 上存储更多向量，但更高的压缩级别通常会导致召回率降低。通常采用称为精炼重排的策略来弥补召回率损失，方法是对 IVF-PQ 索引查询大于所需数量的 k 个结果，然后根据未量化向量的距离进行重新排序并减少到 k 个。不幸的是，这意味着需要提供未量化的原始向量，并且通常可以使用多个 CPU 线程有效地完成此操作。

[ C API | C++ API | Python API | Rust API ]

配置参数

构建参数

名称	默认值	描述
n_lists	sqrt(n)	用于划分索引的粗粒度聚类数量。该值的一个良好启发式是 sqrt(索引中的向量数量)
kmeans_n_iters	20	搜索 k-means 中心时的迭代次数
kmeans_trainset_fraction	0.5	用于迭代 k-means 构建的训练数据比例
pq_bits	8	使用 PQ 压缩后每个向量元素的位长。可能的值是 4 到 8 之间的任何整数。
pq_dim	0	使用 PQ 压缩后每个向量的维度。当为 0 时，维度将根据启发式算法设置。
codebook_kind	per_subspace	码本的创建方式。per_subspace 在一些子维度上训练 k-means，而 per_cluster
force_random_rotation	false	即使 dim % pq_dim == 0，也对输入数据和查询应用随机旋转矩阵
conservative_memory_allocation	false	为了支持动态索引（预计稍后会添加点），可以预先故意过度分配单个 IVF 列表，以减少列表大小增加的数量和影响，因为增加列表大小需要分配更多内存并将旧列表复制到新的更大的列表中。
add_data_on_build	True	在构建索引后是否应将训练点添加到索引中？
max_train_points_per_pq_code	256	在 PQ 码本训练期间，每个 PQ 码使用的最大数据点数。

搜索参数

名称	默认值	描述
n_probes	20	对每个查询点要扫描的最接近的 IVF 列表数量。
lut_dtype	cuda_r_32f	用于存储 pq 查找表的数据类型。也可以使用 cuda_r_16f 表示半精度，使用 cuda_r_8u 表示 8 位精度。较小的查找表可以放入共享内存并显着提高搜索时间。
internal_distance_dtype	cuda_r_32f	搜索时计算的距离/相似度的存储数据类型。也可以使用 cuda_r_16f 表示半精度。
preferred_smem_carveout	1.0	首选的 SM 统一内存 / L1 缓存比例，用作共享内存。默认值为 100%

调优注意事项

IVF-PQ 与 IVF-Flat 有类似的调优注意事项，尽管 PQ 压缩率增加了权衡索引大小与搜索质量的额外变量。

重要的是要注意，IVF-PQ 会很快变得非常具有损耗性，因此通常需要精炼重排来获得合理的召回率。此步骤通常包括最初搜索比所需数量更多的 k 个近邻，然后通过计算精确距离将结果邻域减少到 k 个。此步骤可以在 CPU 或 GPU 上高效执行，并且通常对搜索延迟只有边际影响。

内存占用

索引（设备内存）：

简单的近似公式：\(n\_vectors * (pq\_dim * \frac{pq\_bits}{8} + sizeof_{idx}) + n\_clusters\)

IVF 列表最终由一个稀疏数据结构表示，该结构存储指向每个列表的指针、一个包含每个列表中每个向量索引的索引数组，以及一个包含每个列表的编码（和交错）数据的数组。

IVF 列表指针：\(n\_clusters * sizeof_{uint32\_t}\)

索引：\(n\_vectors * sizeof_{idx}\)

编码数据（交错）：\(n\_vectors * pq\_dim * \frac{pq\_bits}{8}\)

按子空间方法：\(4 * pq\_dim * pq\_len * 2^{pq\_bits}\)

按聚类方法：\(4 * n\_clusters * pq\_len * 2^{pq\_bits}\)

额外项：\(n\_clusters * (20 + 8 * dim)\)

索引（主机内存）：

当在主机上使用数据集进行精炼时，需要原始的未量化向量：\(n\_vectors * dims * sizeof_{float}\)

搜索峰值内存使用量（设备）：

总使用量：\(index + queries + output\_indices + output\_distances + workspace\)

工作空间大小并非微不足道，一个启发式算法控制批处理大小，以确保工作空间适合 raft::resource::get_workspace_free_bytes(res)`。

构建峰值内存使用量（设备）：

\[ \begin{align}\begin{aligned}\frac{n\_vectors}{trainset\_ratio * dims * sizeof_{float}}\\+ \frac{n\_vectors}{trainset\_ratio * sizeof_{uint32\_t}}\\+ n\_clusters * dim * sizeof_{float}\end{aligned}\end{align} \]

注意，如果工作空间内存资源中剩余空间不足，IVF-PQ 构建会自动切换到托管内存用于训练集和标签。