Faiss PQ 进阶：GPU 加速与 HNSW 融合的深度探索

你好！如果你正在处理海量的向量数据，并且希望在速度、内存和精度之间找到那个“甜蜜点”，那么你一定对 Faiss 不陌生。而在 Faiss 的众多索引技术中，乘积量化（Product Quantization, PQ）无疑是压缩和加速近似最近邻（ANN）搜索的利器。不过，基础的 PQ 用法可能已经无法满足你对极致性能的追求了。今天，我们就来深入聊聊 PQ 的一些“高阶玩法”：如何在 GPU 上榨干并行计算能力，以及如何将 PQ 与强大的 HNSW 图索引结合起来，看看它们能碰撞出怎样的火花。

我们假设你已经对 PQ 的基本原理有所了解，知道它是将高维向量切分成子向量，然后对每个子空间独立进行 k-means 聚类，用聚类中心的 ID 来表示原始子向量，从而实现压缩。我们也知道基于 PQ 码本可以快速估算距离，比如非对称距离计算（Asymmetric Distance Computation, ADC）——查询向量是原始向量，数据库向量是 PQ 编码；或者对称距离计算（Symmetric Distance Computation, SDC）——查询向量和数据库向量都被 PQ 编码（虽然 SDC 精度损失更大，但在某些内存极度受限场景下有用）。

这篇文章的目标，就是带你超越 IndexPQ 的基础应用，探索如何在 GPU 上为 PQ 的距离计算插上翅膀，以及 IndexHNSWPQ 这个混合索引背后的奥秘、优势和挑战。

一、GPU 火力全开：加速 PQ 距离计算

当你的数据集达到数十亿甚至更多时，单靠 CPU 来计算查询向量与海量 PQ 编码向量之间的距离，即使是近似距离，也会成为性能瓶颈。尤其是在需要低延迟响应的在线服务场景中，这简直是不可接受的。怎么办？答案是：召唤 GPU！

GPU 拥有数千个核心，天然适合执行大规模并行计算。PQ 的距离计算过程，特别是 ADC，具有良好的并行性。回忆一下 ADC 的计算公式：

d(q, y) ≈ d_hat(q, y) = Σ_{i=1}^{m} || q_i - c_i(y_i) ||^2

其中 q 是查询向量，y 是数据库中的 PQ 编码向量，m 是子空间数量，q_i 是查询向量的第 i 个子向量，c_i(y_i) 是 y 的第 i 个子编码 y_i 对应的第 i 个子空间的聚类中心（码字）。

这个计算过程主要包含两个步骤：

1. 查找子距离： 对于查询向量 q 的每个子向量 q_i，计算它与对应子空间 i 中所有 k* （通常是 256，因为常用 8 比特编码）个码字 c_i(j) 的距离 || q_i - c_i(j) ||^2。这一步通常在查询时预计算出来，形成一个 m x k* 的距离表 T_q。
2. 累加子距离： 对于每个数据库 PQ 编码 y = (y_1, y_2, ..., y_m)，通过查表 T_q 得到每个子编码对应的子距离 T_q[i][y_i]，并将这 m 个子距离相加，得到近似距离 d_hat(q, y)。

在 CPU 上，这个过程通常是串行地处理每个数据库向量 y，或者使用 SIMD 指令进行有限的并行化。但在 GPU 上，我们可以做得更多！

CUDA 核函数设计思路

要在 GPU 上高效实现 PQ 距离计算，关键在于设计合适的 CUDA 核函数（Kernel）来最大化并行度并优化内存访问。

1. 数据布局：

   • PQ 编码： 海量的数据库 PQ 编码通常存储在 GPU 的全局内存（Global Memory）中。为了实现内存合并（Memory Coalescing），最好将同一个向量的 m 个子编码连续存储。例如，一个 uint8_t 数组，每个向量占用 m 个字节。
   • 预计算距离表 T_q： 这个表的大小是 m x k* （例如 m x 256），对于每个查询 q 都是独立的。它可以存储在常量内存（Constant Memory）或只读缓存（Read-Only Cache，通过 __ldg() 内置函数访问）中，因为所有线程块（Thread Blocks）和线程束（Warps）都需要频繁读取它，而它在单次查询中是不变的。对于非常大的 m，可能需要存储在全局内存，但利用好 L1/L2 缓存至关重要。
   • 码本 c_i(j)： 码本本身在训练后是固定的，用于预计算 T_q。在 GPU 上执行预计算时，码本也需要存储在 GPU 内存中，同样可以考虑常量内存或纹理内存（Texture Memory），后者对空间局部性不强的访问模式有优化。

2. 并行化策略：

   • 基本思路： 最直接的想法是让 GPU 上的每个线程负责计算查询向量 q 与一个数据库 PQ 编码 y 之间的近似距离。如果有 N 个数据库向量，我们可以启动 N 个线程（或者更实际地说，启动足够多的线程块，每个线程块处理一部分向量）。
   • 核函数内部（单个线程）： 每个线程 tid 对应一个数据库向量 y_{tid}。该线程需要：
     a. 从全局内存读取 y_{tid} 的 m 个子编码 (y_{tid,1}, ..., y_{tid,m})。
     b. 根据这 m 个子编码，从预计算距离表 T_q 中查找对应的 m 个子距离：T_q[0][y_{tid,1}], T_q[1][y_{tid,2}], ..., T_q[m-1][y_{tid,m}]。
     c. 将这 m 个子距离累加起来，得到最终的近似距离 d_hat(q, y_{tid})。
     d. 将结果写入全局内存中的输出缓冲区。

3. 优化技巧：

   • 内存访问优化：
     • 确保对 PQ 编码的读取是合并的。如果一个 Warp (32 个线程) 同时读取连续的 32 个数据库向量的第一个子编码，然后是第二个子编码……这样访问全局内存效率较高。
     • 对于距离表 T_q 的访问，由于所有线程都在读取相同的数据（或者说是根据自己的 y_{tid,i} 读取不同的列，但访问模式相对固定），利用常量内存广播或只读数据缓存可以显著提升性能。
   • 计算优化：
     • Warp 内求和 (Reduction): 如果 m 不是特别大（例如，小于 32），一个线程可以直接完成 m 次查表和累加。如果 m 较大（例如 64, 128），可以将这 m 次累加分配给同一个 Warp 内的多个线程，或者让一个线程分多次完成。更高级的技巧是使用 Warp 内置函数（如 __shfl_down_sync, __reduce_add_sync 等，需要 Volta 或更新架构）来并行执行 Warp 内的求和，这通常比单个线程循环累加更快。
     • 共享内存 (Shared Memory): 如果一个线程块负责处理多个数据库向量，并且 m 很大，可以考虑将查询 q 的预计算距离表 T_q 的一部分加载到共享内存中。这样，块内所有线程访问 T_q 时就能命中高速的共享内存，而不是每次都去访问相对较慢的常量缓存或全局内存。但这会增加实现的复杂性，并可能受限于共享内存的大小。
   • 指令级并行： 编译器通常会自动进行指令调度，但理解计算依赖关系有助于编写更优化的代码。例如，查表和累加操作可以流水线化。

Faiss 中的 GPU 实现

Faiss 的 GPU 模块 (faiss/gpu/) 提供了 GpuIndexPQ 类，它内部就实现了高效的 PQ 距离计算核函数。当你调用 GpuIndexPQ 的 search 方法时，它会：

1. 将查询向量 q 传输到 GPU。
2. 在 GPU 上启动一个核函数来预计算距离表 T_q (通常称为 distance_tables 或类似名称)。这个核函数通常会让每个线程块负责计算一部分子空间的距离表。
3. 启动另一个（或一系列）主要的搜索核函数，该核函数实现上述的并行化策略，让大量线程并发地计算 q 与数据库中所有（或经过预筛选的）PQ 编码向量的近似距离。
4. 这些核函数内部大量运用了我们讨论的优化技巧，如内存合并、共享内存/只读缓存利用、Warp 内操作等。
5. 计算出的距离（以及对应的向量 ID）会被用来维护一个 Top-K 结果堆，这也可以在 GPU 上高效完成，例如使用并行化的堆操作或归并排序的变种。

Faiss GPU 代码是高度优化的，考虑了不同 GPU 架构的特性，并且使用了像 CUTLASS 这样的底层库来进一步提升性能。深入研究 faiss::gpu::GpuIndexPQ::searchImpl_ 和相关的核函数（如 pq_scan_kernels.cu 中的函数）可以获得更具体的实现细节。

性能考量： GPU 加速 PQ 并非没有代价。数据传输（CPU 到 GPU 的查询向量，GPU 到 CPU 的结果）会带来开销。核函数启动本身也有延迟。因此，GPU 加速通常在以下情况效果显著：

• 数据库规模极大 (N 非常大)。
• 查询是批量进行的 (Batch Processing)，可以摊销数据传输和核函数启动开销。
• d (向量维度) 和 m (子空间数) 适中，使得计算量足够大，能够体现 GPU 的优势。

二、PQ 与 HNSW 的联姻：内存与速度的权衡

HNSW (Hierarchical Navigable Small World graphs) 是目前最先进的基于图的 ANN 搜索算法之一。它通过构建一个多层导航图，能够在对数复杂度内（近似地）找到最近邻，同时保持很高的召回率。然而，标准的 HNSW 索引需要存储所有原始的全精度向量，这在处理超大规模数据集时会导致惊人的内存消耗。

比如，存储 10 亿个 128 维的 float32 向量，光是向量数据就需要 10^9 * 128 * 4 bytes ≈ 512 GB 的内存！这还没算 HNSW 图结构本身的开销。

这时候，我们就想到了 PQ 的看家本领——压缩。能不能把 PQ 编码塞进 HNSW 图里，用近似距离来指导图的遍历呢？

答案是肯定的，这就是 IndexHNSWPQ 背后的思想。

HNSWPQ 的运作机制

IndexHNSWPQ 结合了 HNSW 的图结构和 PQ 的向量压缩与近似距离计算：

1. 训练与构建：

   • 首先，像训练普通 IndexPQ 一样，需要对数据集（或其样本）进行训练，得到 PQ 的码本。
   • 然后，构建 HNSW 图。但在图的节点中，我们不再存储原始向量，而是存储它们的 PQ 编码。原始向量可能被丢弃（如果内存极度宝贵），或者存储在磁盘/其他地方，或者也存储一份在内存中用于后续的精确重排（Faiss 的 IndexHNSWFlat 会存原始向量，IndexHNSWPQ 则存 PQ 码）。

2. 搜索过程：

   • 搜索从 HNSW 图的顶层入口点开始。
   • 在每一层，算法需要找到当前节点邻居中离查询向量 q “最近”的若干个节点，作为下一层搜索的入口或最终候选集的一部分。
   • 关键区别： 在 IndexHNSWPQ 中，计算 q 与邻居节点 y 之间距离时，使用的不是精确距离（比如 L2），而是基于 y 的 PQ 编码计算出的近似距离，通常是 ADC (d_hat(q, y))。
   • 算法维护一个基于近似距离排序的候选列表（或优先队列）。
   • 根据这个近似距离驱动的搜索路径，逐步下降到图的底层（第 0 层）。
   • 在第 0 层，执行更精细的邻居探索，同样使用 PQ 近似距离。

3. （可选但强烈推荐）重排 (Re-ranking)：

   • 通过 HNSW 图遍历和 PQ 近似距离，我们得到了一个较小的候选集（例如，efSearch 个节点）。
   • 由于 PQ 距离是近似的，这个候选集里可能包含一些“假近邻”，并且排序可能不完全准确。
   • 为了提高最终结果的精度（召回率），通常需要一个重排步骤：获取这个候选集中每个节点的原始全精度向量（如果存储了的话），计算它们与查询向量 q 的精确距离，然后根据精确距离重新排序，最终返回 Top-K 结果。
   • Faiss 的 IndexHNSWSQ (Scalar Quantizer) 和 IndexHNSW2Level (两层量化，第二层可以是 PQ) 都体现了类似的思路，并内置了重排机制。IndexHNSWPQ 本身如果只存储 PQ code，则需要外部配合获取原始向量来重排。

优势

• 显著的内存压缩： 这是最大的优点。用 PQ 编码代替全精度向量可以极大地降低内存占用。例如，使用 m=16, nbits=8 的 PQ，一个 128 维 float32 向量（512 字节）可以被压缩到 16 字节，压缩比高达 32 倍！这使得在单机内存中容纳数十亿甚至上百亿向量成为可能。
• 可能加速图遍历（特定情况）： 计算 PQ 近似距离（特别是高度优化的 ADC）可能比计算全精度 L2 距离更快，尤其是在维度 d 很高的情况下。这可能对图遍历的速度有一定正面影响，但这通常不是选择 HNSWPQ 的主要原因，HNSW 本身的图结构带来的加速远比距离计算本身的优化要显著。

挑战与权衡

• 近似误差对图导航的影响： 这是最核心的挑战。由于使用 PQ 近似距离来指导 HNSW 的搜索方向，累积的误差可能导致搜索路径偏离最优路径，错过真正的最近邻。这会直接损害搜索的召回率。PQ 的精度（取决于 m, nbits 和码本质量）对最终搜索效果至关重要。
  • 思考一下： 如果某个节点的 PQ 近似距离偏小，它可能会被错误地优先选择；如果某个真正近的邻居因为 PQ 近似距离偏大而被忽略，搜索可能就走错了方向。
• 精度 vs. 内存/速度的权衡：
  • 增加 PQ 的子空间数 m 或每个子空间的比特数 nbits 可以提高近似距离的精度，从而改善 HNSW 的导航准确性，提高召回率。但这会增加 PQ 编码的存储大小（降低压缩率），并可能增加近似距离的计算时间。
  • 选择合适的 HNSW 参数 (M, efConstruction, efSearch) 与 PQ 参数 (m, nbits) 是一个复杂的调优过程，需要根据具体的应用场景和对内存、速度、精度的不同要求来权衡。
• 重排的必要性与开销： 为了弥补 PQ 近似距离带来的精度损失，基于精确距离的重排几乎是必须的。这意味着你可能还是需要存储或能够快速获取原始向量。重排本身也会增加额外的计算开销（精确距离计算）和潜在的 I/O 开销（如果原始向量存储在较慢的介质上）。
• 训练开销： 除了构建 HNSW 图，还需要额外的时间来训练 PQ 码本。

实践建议：

• 通常，IndexHNSWPQ 适用于内存是首要瓶颈，且可以接受一定精度损失（或者有能力进行有效重排）的场景。
• 仔细选择 PQ 参数。m 的值（通常是 d 的因子，如 d/2, d/4）和 nbits（通常是 8）需要根据数据特性和精度要求来定。更大的 m 通常更好，但会增加内存和计算量。
• efSearch 参数（搜索时维护的动态列表大小）对 HNSWPQ 的召回率影响很大，通常需要设置得比 IndexHNSWFlat 更大一些，以补偿近似距离导航带来的不确定性。
• 务必评估是否需要以及如何实现重排。如果原始向量存储在内存中，重排相对容易；如果存储在磁盘，需要考虑 I/O 成本。

三、更进一步：优化与未来

PQ 与 GPU、HNSW 的结合并非终点，还有一些值得关注的方面：

• 优化 PQ (OPQ)： OPQ 在 PQ 之前应用一个旋转矩阵，使得子空间内的方差更加均衡，可以提高量化精度。将 OPQ 与 HNSW 结合（IndexHNSWOPQ，虽然 Faiss 中没有直接命名为此，但可以通过 IndexPreTransform + IndexHNSWPQ 实现）可能会比 HNSWPQ 获得更好的精度/内存权衡。
• GPU 上的 HNSWPQ？ 能否在 GPU 上加速 HNSWPQ 的搜索过程？这比较复杂。HNSW 的图遍历本质上是内存依赖性强、访问模式不规则的操作，这与 GPU 擅长的规则、大规模并行计算模式不太匹配。虽然可以在 GPU 上加速其中的 PQ 距离计算部分，但图导航本身的并行化和优化是难点。Faiss GPU 目前对 HNSW 的支持有限，但这是研究的热点方向。
• 参数自动调优： HNSWPQ 涉及 HNSW 和 PQ 两组参数，手动调优非常耗时。利用自动化方法（如贝叶斯优化）来寻找最佳参数组合是一个有前景的方向。
• 与其它量化方法比较： 除了 PQ，还有标量量化 (SQ)、残差量化 (RQ)、局部敏感哈希 (LSH) 等。将这些量化方法与 HNSW 结合（如 Faiss 的 IndexHNSWSQ），或者探索 RQ+HNSW 等组合，提供了更多的选择空间，需要根据具体需求进行评估。

结语

我们深入探讨了 Faiss 中 PQ 的两种高级应用：利用 GPU CUDA 核心并行计算 ADC/SDC 距离以实现大规模搜索加速，以及将 PQ 编码嵌入 HNSW 图节点以大幅降低内存占用。这两种技术都展示了在追求极致性能和效率时，算法与硬件、不同算法组件之间协同工作的重要性。

GPU 加速 PQ 距离计算为处理海量数据提供了强大的算力支持，尤其适用于需要高吞吐量和低延迟的场景。而 HNSWPQ 则是在内存极度受限的情况下，通过牺牲一定的导航精度（通常需要重排补偿）来换取存储百亿甚至千亿级别向量数据的能力。

理解这些技术的内部机制、优势和挑战，对于选择合适的索引策略、进行有效的参数调优至关重要。记住，没有万能的索引，只有最适合你特定数据、特定场景、特定约束（内存、速度、精度）的解决方案。希望这次的深度探索能帮助你更好地驾驭 Faiss，解决更大规模的向量检索问题！