FAISS内部：十亿级相似性搜索

# 深入 FAISS：十亿级相似性搜索 来源：https://fremaconsulting.ch/blog/faiss ## 1. 一切都是向量 在现代 AI 中，图像、文本和音频并不是以“猫”或“交响乐”的形式被计算机理解的。相反，它们被转换成数字列表，称为**嵌入向量**或**向量**。 这些向量存在于一个高维几何空间中。核心思想很简单：语义相似的项目在这个空间中彼此靠近。 基础 · 向量与相似性 ### 1. 挑战 计算机没有眼睛或耳朵，它们只理解数字。如何向机器解释“狗”或“猫”的概念？ ### 2. 向量化 我们使用神经网络将原始数据（像素、文本）转换为数字列表。这个列表就是“嵌入向量”，它捕获了语义含义。 ### 3. 向量空间 这些数字充当坐标（X, Y, Z...）。我们可以将每张图像或每个词绘制为高维空间中的一个点。 ### 4. 相似性即距离 在这个空间中，相似的概念彼此靠近。“猫”靠近“狗”，但远离“车”。我们使用距离（欧几里得或余弦）来衡量。 上图展示了一个简化的二维切片。设想这是在 1024 维空间中。对于一个查询点，找到其“最近邻”等价于在数据库中找到最相似的项目。 ## 2. 大规模 NN 问题 形式上，给定一个查询和包含个向量的数据库，精确最近邻搜索解决的是： 朴素**暴力**方法显式计算所有距离。每个查询成本：时间以及全精度存储整个数据库的内存。对于 SIFT 描述子（，每个浮点数 4 字节），这需要 512 GB 内存，每次搜索需要十亿次距离计算，对于大规模检索或实时 LLM 内存这样的系统来说，完全不可行。 本文其余部分探讨 FAISS 利用的两条逃逸路线： - **划分**：在查询时跳过大部分数据（IVF）。 - **压缩**：使每次比较代价低廉，并使数据库适合内存（乘积量化）。 生产系统结合两者。 ## 3. FAISS 来救援 FAISS（Facebook AI 相似性搜索）引入了**近似**搜索方法。通过牺牲少量准确性（可能得到的是第二佳匹配而非最佳），我们可以将搜索速度提高几个数量级。 让我们探讨两种关键的索引策略： - **Flat：** 基准暴力搜索。慢但准确。 - **IVF（倒排文件）：** 将空间划分为“Voronoi 单元”。我们只查看最有希望的单元。 自己试试。增加数据集大小，观察不同索引的表现。 ### 交互式索引基准测试 比较不同索引算法在随机数据集上的性能。 数据集大小 (N) 索引类型 ## 4. 使用 IVF 进行划分 **倒排文件 (IVF)** 索引的工作方式类似于图书馆分类系统。不需要查看每一本书来找到特定书名，而是先前往“科幻小说”区域。 数学上，它使用 K-Means 聚类将向量空间划分为**Voronoi 单元**。搜索时，我们先识别查询向量落入哪个单元（使用“粗量化器”），然后*仅*计算该单元（以及可能几个相邻单元）内向量的距离。 ### IVF 聚类1. 原始数据 K-means 将随机向量划分成 IVF 搜索所使用的 Voronoi 单元。 **开始：** 这里是我们的原始数据集。在二维空间中有 80 个向量。线性搜索（暴力法）需要计算到 80 个点的距离。 ## 5. 使用乘积量化进行压缩 IVF 通过跳过大部分数据库来加速搜索，但每个向量仍然未压缩。十亿个 SIFT 描述子仍然需要 512 GB 内存。**乘积量化 (PQ)**，由Jégou, Douze 和 Schmid (2011) (https://hal.inria.fr/inria-00514462/document) 提出，是 FAISS 用于将每个向量压缩到 8 字节同时保持距离估计有意义的压缩技巧。 ### 相同的质心，不同的任务 §4 (https://fremaconsulting.ch/blog/faiss#sec-4)使用质心来*划分*搜索空间。PQ 使用它们来*压缩*每个向量。相同的 K-Means 数学，不同的目的：质心的*索引*成为向量本身。如果码本包含个质心，任何向量都压缩为中的一个整数。 ### 像素级类比 一个 24 位 RGB 像素可以是 1670 万种颜色中的任意一种。GIF 放弃了其中一些范围：它选择一个 256 色**调色板**，并用一个 8 位索引替换每个像素。存储量减少 3 倍，图像看起来仍然像原图。这个调色板就是一个码本；索引就是一个码。 一个 128 维的 SIFT 描述子就是一个很长的像素。我们希望使用相同的技巧：找到一个代表性向量的调色板，用其最近的调色板索引替换每个描述子。 ### 一个索引占用多少比特？ 为了唯一标记个质心，你需要足够的比特来区分它们。每个比特使标签数量加倍，因此码长为比特。具体来说： - 个质心 →每码比特（1 字节） - 个质心 →比特 - 个质心 → 64 比特（8 字节） 质心越多，量化器分辨的细节越精细，码长也相应增长。码本设计者玩的是一个比特预算游戏。 ### 论文的激进目标 Jégou 等人希望每个 128 维 SIFT 描述子变成一个**64 比特码**，即**每维度 0.5 比特**，相比原始 512 字节实现 64 倍压缩比。每维度半比特目标宏大：意味着量化器必须用近似二进制的选择来编码每个轴上的所有变化。 使用一个平面码本达到 64 比特意味着选择个质心。这大约是 1800 万兆（18 乘以 10 的 18 次方），超过地球上沙粒的总数。在我们否定它之前，先感受一下这个数字。 ### 平面码本可行性 选择一个向量维度和每维度的比特预算，查看单个平面量化器所需的质心数量和存储空间。 每码总比特64bits 所需质心数 (k = 2^64)1.84 × 10^19个质心 码本存储9.44ZB 不可能 大于地球上所有云存储的总和。 Lloyd 算法大约需要 5.53 × 10^20个训练样本。 ### 为什么 2^64 不是一个码本 当变得如此之大时，同时遇到三个障碍： - **存储。** 每个质心是 128 个浮点数 × 4 字节 = 512 字节。乘以个质心给出约**94 泽字节**，超过 2025 年地球上所有云存储的总和。你无法持久化这个码本，更不用说在查询时分页访问它了。 - **训练数据。** Lloyd 算法需要每个质心至少几十个样本来收敛。这需要个训练向量。现有数据集中没有如此大的。 - **查询成本。** 编码一个描述子需要将其与每个质心比较。每个向量进行 18 百万兆次距离计算是不可接受的延迟。 平面码本在所有三个方面都失败了。论文的做法是保持*有效*词汇表大小不变，同时将存储的码本缩减多个数量级。 ### 分解技巧 回到 GIF 类比。不为整张图像寻找一个 256 色调色板，而是将图像切成 8 块，每块有自己的 256 色调色板。每块现在是一个字节；整张图像是 8 字节；每块都能访问一个完整 256 色调色板，且该调色板针对其自己的像素进行了优化。PQ 对向量正是这样做的。 形式化：将分割成个子向量，并为每个块学习一个小型子量化器。乘积量化器是它们输出的元组： 当每块具有个质心时，*有效*码本具有种可能的组合（与不可行的平面情况数量级相同），但我们每个子量化器只存储个质心向量。总码本存储为（论文§II.B (https://hal.inria.fr/inria-00514462/document)，表 I），足够小以放入 L2 缓存。 每个编码向量占用，每个块索引一个字节。比例：512 B 到 8 B，内存**64 倍**下降，同时码本实际适合机器。 下面的演练在论文的参考 SIFT 设置（ ）上构建完整索引，将上述每个概念转化为对具体数字的实际操作。 压缩 · 乘积量化 ### 1. 一个描述子，完全放大 一个 128 维向量是 128 个浮点数。放大到单个浮点数，你会发现 IEEE 754：1 个符号位，8 个指数位，23 个尾数位。 每维度四个字节，每个向量 512 个字节。 ### 2. 将其切成 8 个子向量 相同的 128 个数字，重新组合成 8 个连续的块，每个块 16 维。 每个块有自己的颜色，因为每个块将有自己的码本。 ### 3. 现在乘以：一个小型演示数据库 码本需要一个人口，而不是单个向量。这个演示使用 4 个示例向量；论文在数百万个上进行训练。 每个数据库向量都进行相同的 8 块分割，每个块为每个向量在其自己的子空间中贡献一个点。 ### 4. 聚焦一个码本：u1 每个子空间都有自己的码本，独立训练。这种独立性就是乘积量化中的“乘积”。 为了查看完整配方，我们放大到仅一个子向量 u1，为了可读性绘制成二维（论文的子空间存在于 d* = D/m = 16 维，其中 D = 128，m = 8）。四个训练向量是教学性的简图；真正的 PQ 在每个子空间上训练大约 10^6 个描述子。 Lloyd 算法仅在这个子空间上运行 k-means：将每个点分配给它最近的质心，将每个质心移动到其簇均值，重复直到块内失真 MSE(q1) = E‖u1(x) − q1(u1(x))‖^2 停止下降。 你最终得到的 6 个质心（论文中为 k* = 256）构成码本 C1。编码任何未来向量的第一个块就是最近质心查找：发出索引 i of c1,i。 ### 5. 相同操作，并行 8 次 所有 8 个子空间同时以相同方式量化，每个块在其自己的数据切片上使用自己的码本。 Lloyd 在每个块上独立运行（论文中每个码本 k* = 256 个簇，此处显示 6 个）。 ### 6. 编码一个全新向量 查询 x 在搜索时到达，它从未在训练集中出现过。 将其分割成 8 个子向量，并在每个图中，吸附到最近的质心。8 个索引拼接成 8 字节的码。 ### 7. 解码与每块误差 重建是表查找：用质心替换每个索引。 原始点与质心之间的差距是每块量化误差；总重建误差是它们的平方和。 ### 8. 不解码的距离计算：SDC vs ADC 从码中计算查询到数据库距离的两种方式。 SDC 对两个点都进行量化，并读取质心到质心的距离。ADC 保持查询的全精度，并读取原始点到质心的距离。 ADC 具有更紧的误差界；FAISS 默认使用它。 一旦编码，如何在不将每个向量解码回其 512 字节的情况下测量距离？论文给出了两种选择。 1. **SDC**（对称）对查询和数据库向量都进行量化，并从每个子量化器的预计算表中读取成对质心距离。 2. **ADC**（非对称）将查询保持在全精度，并仅预计算每个块的查询到质心距离，然后对每个块求和一次： 每个查询的成本相似（Jégou 等人，表 II），但**ADC**具有更紧的误差界：均方距离误差满足对于**ADC**（等式 18） vs 对于**SDC**。FAISS 默认使用**ADC**。 ## 6. 组合：IVFPQ **IVFPQ**（论文中也称为 IVFADC）堆叠了前两种技术：粗量化器将数据库裁剪为少量单元，乘积量化器将剩余部分压缩为每个向量 8 字节。一个微妙的选择使组合生效：PQ 编码的不是原始向量，而是其相对于粗质心的**残差**。 **索引**一个数据库向量： 1. 粗量化，找到最近的粗质心； 2. 计算残差； 3. 使用乘积量化器编码残差，得到 8 字节码； 4. 将对附加到单元的倒排列表，其中是数据库向量的索引，是第 3 步产生的 m 字节 PQ 码。 残差集中在零附近，因此 PQ 码本将其比特用于建模粗量化器尚未捕获的变异。相同的字节预算，更好的重建。 **搜索**一个查询： 1. 找到个离最近的粗质心； 2. 计算每个选定单元的查询残差； 3. 预计算**距离查找表 (LUT)**。索引选择 PQ 子块（在我们的设置中），索引选择该子块码本中的一个质心（）。每个条目是查询子块残差与该质心之间的平方距离： 4. 扫描倒排列表：对于每个存储的条目（其 PQ 码为每个子块一个索引），从 LUT 读取预计算的距离并求和：每次候选进行次表读取、次加法。没有浮点乘法，没有平方根。 5. 维护一个固定容量的最大堆，大小为，保存迄今为止看到的最佳候选，按 ADC 距离排序。当堆不满时，推入每个扫描到的候选。一旦堆满，将每个新 ADC 距离与根（当前最佳中最差的那个）比较：如果更小，弹出根并推入候选，否则丢弃它。这个堆在个探测的倒排列表之间共享。扫描完成后，按升序提取堆中的内容，得到查询的近似最近邻。 下面的演练使用每一步的实际算术运行完整流水线。 阶段 1 · 训练（离线） ### 1. 训练粗量化器 在训练集上运行 k-means 以学习个粗质心。它们将空间划分成 Voronoi 单元，这些就是 IVFADC 的*邻域*。 论文设置：对于 SIFT（Jégou 等人§IV.A (https://hal.inria.fr/inria-00514462/document)）。我们这里只渲染 8 个单元，以便每个单元保持可读。 ### 2. 从原始向量到残差 对于每个训练向量，减去其最近的粗质心： 减去质心将每个单元的点重新居中到其自身坐标系的原点。所有单元合并在一起共享一个*中心化*的*残差*分布，这就是乘积量化器要学习的内容。 ### 3. 为什么残差量化更好 残差携带的能量比原始向量少得多：平均而言。相同的质心预算，更小的覆盖区域，更精细的量化步长。 **类比：视频压缩。** 先发送一个粗糙帧，然后将其比特花费在差异上。你编码运动和细节，而不是从头开始编码整幅图像。 在视觉效果上切换**覆盖**以叠加两个直方图并读取方差比。 ### 4. 在残差上训练 PQ 将每个残差分割成个子向量。在每个块上独立运行 Lloyd，每个块一个码本。 这些码本对粗量化器已经减去的变异进行建模。相同的 8 字节占用，更锐利的重建。训练到此结束；接下来的三章在高速阶段消耗这些码本。 阶段 2 · 编码数据库 ### 5. 分配粗桶 一个新的数据库向量到达。将其与每个粗质心比较，选择最近的。调用其索引。该索引标识将存储的倒排列表。 ### 6. 使用 PQ 编码残差 计算，分割成八个子向量，每个吸附到中最远的质心。发出 8 个块索引。 每块比特，8 块，每块一字节。八字节的码将其位置定位在倒排列表内部，无需存储原始坐标。 ### 7. 附加到倒排列表 将附加到倒排列表。的原始坐标被*丢弃*。每个向量的内存： 点击左侧的**Voronoi 单元**以检查其列表中存储的条目。 阶段 3 · 搜索 ### 8. 查询到达，探测 w 个单元 你的真正最近邻可能落在单元边界附近。为了捕获它，检查距离最近的粗质心，而不仅仅是最近的那个（*多重分配*）。