knn search algorithm comparison

KNN搜索算法比较

1951年推出的K-Nearest邻居（K-NN）算法已被广泛用于分类和回归任务。核心概念涉及将K最相似的实例（邻居）识别到数据集中的给定查询点，并使用这些邻居进行预测或分类。近年来，矢量数据库和向量索引的重要性已经增长，特别是在处理大量文本和其他数据数据集中，信息检索以支持大型语言模型（LLMS）。该应用程序的一个突出例子是检索仪器（RAG）。

该项目比较了各种数据集大小和尺寸的不同K-NN搜索算法的性能。比较的算法是：

1. KD-Tree
2. 球树
3. 蛮力（完整的KNN）
4. HNSW（层次可导航的小世界）

算法解释

1. KD-Tree（K维树）：

   • 用于在K维空间中组织点的空间分分数据结构。
   • 通过递归将空间沿不同的维度划分为二进制树。
   • 低维空间的有效效率（通常<20维）。
   • 搜索的平均时间复杂度：O（log n），其中n是点数。
   • 由于“维度的诅咒”，在高维空间中的有效性较差。示例：在2D空间中，KD-Tree可能会垂直划分平面，然后在每个级别上交替：

      y
      |
   4  |    C
      |  A   D
   2  |    B
      |___________
      0    2    4  x
   

   点：A（1,3），B（3,1），C（4,3），D（3,3）树结构：root（x = 2） - >左（y = 2） - >右（x = 3）

2. 球树：

   • 将分区的二进制树数据结构指向嵌套的超球。
   • 每个节点代表一个包含点子集的球（超孔）。
   • 对于高维空间，比KD-Tree更有效。
   • 搜索的平均时间复杂性：O（log n），但恒定因子的恒定因子比KD-Tree更高。
   • 当尺寸> 20时，通常比KD-Tree的表现更好。例如：在2D空间中，球树可能会产生嵌套圆：

      y
      |
   4  |    (C)
      |  (A)  (D)
   2  |    (B)
      |___________
      0    2    4  x
   

   外圆包含所有点，内圆圈划分子集。

3. 完整的KNN（蛮力）：

   • 计算从查询点到数据集中所有其他点的距离。
   • 易于实现，但对于大型数据集而言，计算昂贵。
   • 时间复杂度：O（n * d），其中n是点的数量，d是尺寸的数量。
   • 随着数据集大小或维数的增加而变得效率低下。
   • 保证找到最接近的邻居。示例：对于查询点q（2,2）和k = 2：

      y
      |
   4  |    C
      |  A   D
   2  |----Q--B
      |___________
      0    2    4  x
   

   计算距离：QA = 1.41，QB = 1，QC = 2.24，QD = 1.41结果：最近的2个邻居是B和A（或D）

4. HNSW（层次可通航的小世界）：

   • 一个大约最近的邻居搜索算法。
   • 构建一个多层图结构，以进行有效的导航。
   • 提供搜索速度和准确性之间的权衡。
   • 在高维空间和大数据集中表现良好。
   • 搜索的平均时间复杂性：O（log n），但常数比基于树的方法更好。
   • 通过牺牲一些准确性，可以更快地进行搜索。示例：HNSW层的简化2D表示：

    Layer 2:   A --- C
              |
   Layer 1:   A --- B --- C
              |    |    |
   Layer 0:   A --- B --- C --- D --- E
   

   搜索从顶层的随机点开始，然后下降，探索每个级别的邻居直到到达底部。

这些算法之间的选择取决于数据集大小，维度，所需的准确性和查询速度。 KD-Tree和Ball树提供了精确的结果，并且对于低到中等的尺寸有效。完整的KNN很简单，但对于大型数据集而言变得很慢。 HNSW在速度和准确性之间提供了良好的平衡，尤其是对于高维数据或大型数据集提供了良好的平衡。

安装

1. 克隆这个存储库：

    git clone https://github.com/yourusername/knn-search-comparison.git
   cd knn-search-comparison
   

2. 创建虚拟环境（可选但建议）：

    python -m venv venv
   source venv/bin/activate  # On Windows, use `venvScriptsactivate`
   

3. 安装所需的依赖项：

    pip install -r requirements.txt
   

   这将安装requirements.txt文件中列出的所有必要软件包，包括numpy，scipy，scikit-learn，hnswlib，tabulate和tqdm。

用法

使用默认参数运行比较测试：

 python app.py

您还可以使用命令行参数自定义测试参数：

 python app.py --vectors 1000 10000 100000 --dimensions 4 16 256 --num-tests 5 --k 5

可用参数：

• --vectors ：测试矢量计数列表（默认值：1000，2000，5000，10000，20000，50000，100000，200000）
• --dimensions ：要测试的尺寸列表（默认值：4 16 256 1024）
• --num-tests ：每种组合要运行的测试数（默认值：10）
• --k ：搜索最近的邻居数（默认：10）

该脚本将在执行过程中显示进度条，使您估算剩余时间。

可以随时按CTRL+C来中断该脚本。即使在耗时的操作中，例如构建HNSW索引，它也将尝试优雅地退出。

输出

该脚本将在控制台中显示进度并结果。完成后，您会看到：

1. 矢量数和维度的每种组合的结果摘要，包括：
   • 为KD-Tree，Ball Tree和HNSW索引建立时间
   • 每种算法的平均搜索时间
2. 所有结果表
3. 包含详细结果的CSV文件的位置

单个组合的示例输出：

 Results for 10000 vectors with 256 dimensions:
KD-Tree build time:       0.123456 seconds
Ball Tree build time:     0.234567 seconds
HNSW build time:          0.345678 seconds
KD-Tree search time:      0.001234 seconds
Ball Tree search time:    0.002345 seconds
Brute Force search time:  0.012345 seconds
HNSW search time:         0.000123 seconds

最终结果表和CSV文件将包括每种算法的构建时间和搜索时间，从而可以全面比较不同矢量计数和维度的性能。

定制

您可以在app.py中修改以下变量以调整测试参数：

• NUM_VECTORS_LIST ：测试矢量计数列表
• NUM_DIMENSIONS_LIST ：测试尺寸列表
• NUM_TESTS ：每种组合要运行的测试数
• K ：最近的邻居数量

贡献

欢迎捐款！请随时提交拉动请求。

执照

该项目是开源的，并根据麻省理工学院许可获得。

KNN搜索结果图

以下是显示KNN搜索结果的图表：

包含详细结果的CSV文件可在此处找到。